| |
| """Single-GPU training for AR-context HY-WorldPlay Predictor-v2.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import random |
| import time |
| from pathlib import Path |
|
|
| import swanlab |
| import torch |
| import torch.nn.functional as F |
| from torch.optim import AdamW |
| from torch.optim.lr_scheduler import LambdaLR |
| from torch.utils.data import DataLoader |
|
|
| from hyvideo.commons.infer_state import initialize_infer_state |
| from models import HYWorldPlayPredictorV2, HYWorldPlayPredictorV3, PredictorV2Config |
| from predictor_training.checkpoint import ( |
| save_predictor_weights, |
| save_training_state, |
| trainable_state_dict, |
| ) |
| from predictor_training.dataset_v2 import ( |
| PredictorV2PairDataset, |
| move_v2_pair_to_device, |
| predictor_v2_collate, |
| ) |
| from predictor_training.dataset_v3 import ( |
| PredictorV3PairDataset, |
| move_v3_pair_to_device, |
| predictor_v3_collate, |
| ) |
|
|
|
|
| def parse_source_blocks(value: str) -> tuple[int, int]: |
| result = tuple(int(part.strip()) for part in value.split(",") if part.strip()) |
| if result not in {(0, 53), (1, 52)}: |
| raise argparse.ArgumentTypeError("source_blocks must be 0,53 or 1,52") |
| return result |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--manifest", default="datasets/predictor_v2/train_eval_manifest.jsonl") |
| parser.add_argument("--predictor_version", choices=("v2", "v3"), default="v2") |
| parser.add_argument("--source_blocks", type=parse_source_blocks, required=True) |
| parser.add_argument( |
| "--teacher_checkpoint", |
| default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors", |
| ) |
| parser.add_argument("--output_dir", required=True) |
| parser.add_argument("--max_steps", type=int, default=2000) |
| parser.add_argument("--batch_size", type=int, default=1) |
| parser.add_argument("--gradient_accumulation_steps", type=int, default=4) |
| parser.add_argument("--warmup_blocks_steps", type=int, default=100) |
| parser.add_argument("--scheduler_warmup_steps", type=int, default=100) |
| parser.add_argument("--fusion_lr", type=float, default=1e-4) |
| parser.add_argument("--blocks_lr", type=float, default=1e-5) |
| parser.add_argument("--weight_decay", type=float, default=0.01) |
| parser.add_argument("--hidden_weight", type=float, default=0.1) |
| parser.add_argument("--velocity_weight", type=float, default=1.0) |
| parser.add_argument("--grad_clip", type=float, default=1.0) |
| parser.add_argument("--log_every", type=int, default=10) |
| parser.add_argument("--save_every", type=int, default=500) |
| parser.add_argument("--keep_snapshots", type=int, default=2) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--num_workers", type=int, default=0) |
| parser.add_argument("--max_records", type=int, default=None) |
| parser.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True) |
| parser.add_argument("--save_optimizer_state", action=argparse.BooleanOptionalAction, default=True) |
| parser.add_argument("--resume", default=None) |
| parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v2") |
| parser.add_argument("--swanlab_workspace", default=None) |
| parser.add_argument("--swanlab_experiment", default=None) |
| parser.add_argument("--swanlab_mode", choices=("cloud", "offline", "local", "disabled"), default="cloud") |
| return parser.parse_args() |
|
|
|
|
| def initialize_runtime_state() -> None: |
| class Args: |
| sage_blocks_range = "0-53" |
| use_sageattn = False |
| enable_torch_compile = False |
| use_fp8_gemm = False |
| quant_type = "fp8-per-block" |
| include_patterns = "double_blocks" |
| use_vae_parallel = False |
|
|
| initialize_infer_state(Args()) |
|
|
|
|
| def make_scheduler(optimizer: torch.optim.Optimizer, warmup: int, total: int) -> LambdaLR: |
| def factor(step: int) -> float: |
| if step < warmup: |
| return max(1e-8, float(step + 1) / max(1, warmup)) |
| progress = min(1.0, float(step - warmup) / max(1, total - warmup)) |
| return 0.5 * (1.0 + math.cos(math.pi * progress)) |
|
|
| return LambdaLR(optimizer, factor) |
|
|
|
|
| def set_ar_blocks_trainable(model: HYWorldPlayPredictorV2, enabled: bool) -> None: |
| for block in model.double_blocks: |
| for name, parameter in block.named_parameters(): |
| parameter.requires_grad_(enabled and not name.startswith("txt_")) |
|
|
|
|
| def prune_snapshots(output_dir: Path, keep: int) -> None: |
| snapshots = sorted(output_dir.glob("predictor_step_*.safetensors")) |
| for path in snapshots[:-keep] if keep > 0 else snapshots: |
| path.unlink() |
|
|
|
|
| def init_swanlab(args: argparse.Namespace, config: dict): |
| if args.swanlab_mode == "cloud": |
| api_key = os.environ.get("SWANLAB_API_KEY") |
| swanlab.login(api_key=api_key, save=False) if api_key else swanlab.login() |
| fusion_lr = f"{args.fusion_lr:.0e}".replace("e-0", "e-").replace("e+0", "e+") |
| blocks_lr = f"{args.blocks_lr:.0e}".replace("e-0", "e-").replace("e+0", "e+") |
| experiment = args.swanlab_experiment or ( |
| f"predictor-{args.predictor_version}-blocks" |
| + "-".join(str(value) for value in args.source_blocks) |
| + f"-flr{fusion_lr}-blr{blocks_lr}-schedsteps{args.max_steps}" |
| + f"-swarmup{args.scheduler_warmup_steps}" |
| + f"-bwarmup{args.warmup_blocks_steps}" |
| ) |
| return swanlab.init( |
| project=args.swanlab_project, |
| workspace=args.swanlab_workspace, |
| experiment_name=experiment, |
| description=( |
| f"HY-WorldPlay {args.predictor_version} AR-context Predictor, " |
| "fresh exact 10-case targets" |
| ), |
| tags=[ |
| f"predictor-{args.predictor_version}", |
| "ar-context-kv", |
| f"blocks-{args.source_blocks[0]}-{args.source_blocks[1]}", |
| ], |
| config=config, |
| logdir=str(Path(args.output_dir).resolve() / "swanlab"), |
| mode=args.swanlab_mode, |
| ) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA is required") |
| random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| torch.cuda.manual_seed_all(args.seed) |
| initialize_runtime_state() |
| device = torch.device("cuda:0") |
| dtype = torch.bfloat16 |
| output_dir = Path(args.output_dir).resolve() |
| output_dir.mkdir(parents=True, exist_ok=True) |
| log_path = output_dir / "train_log.jsonl" |
|
|
| dataset_class = ( |
| PredictorV3PairDataset if args.predictor_version == "v3" else PredictorV2PairDataset |
| ) |
| collate_fn = ( |
| predictor_v3_collate if args.predictor_version == "v3" else predictor_v2_collate |
| ) |
| move_pair = ( |
| move_v3_pair_to_device |
| if args.predictor_version == "v3" |
| else move_v2_pair_to_device |
| ) |
| dataset = dataset_class( |
| args.manifest, |
| source_block_ids=args.source_blocks, |
| max_records=args.max_records, |
| ) |
| generator = torch.Generator().manual_seed(args.seed) |
| loader = DataLoader( |
| dataset, |
| batch_size=args.batch_size, |
| shuffle=True, |
| num_workers=args.num_workers, |
| pin_memory=False, |
| collate_fn=collate_fn, |
| generator=generator, |
| ) |
|
|
| model_class = ( |
| HYWorldPlayPredictorV3 if args.predictor_version == "v3" else HYWorldPlayPredictorV2 |
| ) |
| model = model_class(PredictorV2Config(source_block_ids=args.source_blocks)) |
| model.load_teacher_initialization(args.teacher_checkpoint) |
| model.to(device=device, dtype=dtype) |
| model.enable_gradient_checkpointing(args.gradient_checkpointing) |
| model.train() |
|
|
| fusion_parameters = list(model.img_fusion.parameters()) + list(model.residual_out.parameters()) |
| block_parameters = [ |
| parameter for parameter in model.double_blocks.parameters() if parameter.requires_grad |
| ] |
| intended_trainable_parameters = sum( |
| parameter.numel() for parameter in fusion_parameters + block_parameters |
| ) |
| intended_parameter_breakdown = model.trainable_parameter_breakdown() |
| optimizer = AdamW( |
| [ |
| {"params": fusion_parameters, "lr": args.fusion_lr}, |
| {"params": block_parameters, "lr": args.blocks_lr}, |
| ], |
| betas=(0.9, 0.95), |
| weight_decay=args.weight_decay, |
| ) |
| scheduler = make_scheduler(optimizer, args.scheduler_warmup_steps, args.max_steps) |
| global_step = 0 |
| micro_step = 0 |
| set_ar_blocks_trainable(model, args.warmup_blocks_steps <= 0) |
|
|
| if args.resume: |
| state = torch.load(args.resume, map_location="cpu", weights_only=False) |
| if tuple(state["predictor_config"]["source_block_ids"]) != args.source_blocks: |
| raise ValueError("Resume checkpoint source blocks do not match") |
| model.load_state_dict(state["model"], strict=False) |
| optimizer.load_state_dict(state["optimizer"]) |
| scheduler.load_state_dict(state["scheduler"]) |
| global_step = int(state["global_step"]) |
| micro_step = int(state.get("micro_step", global_step * args.gradient_accumulation_steps)) |
| if "torch_rng_state" in state: |
| torch.set_rng_state(state["torch_rng_state"]) |
| if "cuda_rng_state" in state: |
| torch.cuda.set_rng_state_all(state["cuda_rng_state"]) |
| set_ar_blocks_trainable(model, global_step >= args.warmup_blocks_steps) |
|
|
| run_config = { |
| **vars(args), |
| "source_blocks": list(args.source_blocks), |
| "manifest": str(Path(args.manifest).resolve()), |
| "teacher_checkpoint": str(Path(args.teacher_checkpoint).resolve()), |
| "dataset_pairs": len(dataset), |
| "effective_batch_size": args.batch_size * args.gradient_accumulation_steps, |
| "trainable_parameters": intended_trainable_parameters, |
| "parameter_breakdown": intended_parameter_breakdown, |
| "dtype": "bfloat16", |
| "predictor_version": args.predictor_version, |
| "gpu": torch.cuda.get_device_name(device), |
| } |
| run = init_swanlab(args, run_config) |
| print( |
| json.dumps( |
| { |
| "dataset_pairs": len(dataset), |
| "source_blocks": args.source_blocks, |
| "trainable_parameters": intended_trainable_parameters, |
| "breakdown": intended_parameter_breakdown, |
| "start_step": global_step, |
| "swanlab_run_id": getattr(run, "id", None), |
| }, |
| indent=2, |
| ), |
| flush=True, |
| ) |
|
|
| optimizer.zero_grad(set_to_none=True) |
| running: list[dict[str, float | str]] = [] |
| epoch = 0 |
| try: |
| while global_step < args.max_steps: |
| epoch += 1 |
| for item in loader: |
| if global_step >= args.max_steps: |
| break |
| if global_step == args.warmup_blocks_steps: |
| set_ar_blocks_trainable(model, True) |
|
|
| inputs, target_hidden, target_velocity = move_pair( |
| item, device=device, dtype=dtype |
| ) |
| started = time.perf_counter() |
| with torch.autocast(device_type="cuda", dtype=dtype): |
| output = model(**inputs) |
| hidden_loss = F.mse_loss( |
| output["pred_hidden"].float(), target_hidden.float() |
| ) |
| velocity_loss = F.mse_loss( |
| output["pred_velocity"].float(), target_velocity.float() |
| ) |
| loss = args.hidden_weight * hidden_loss + args.velocity_weight * velocity_loss |
| scaled_loss = loss / args.gradient_accumulation_steps |
| scaled_loss.backward() |
| micro_step += 1 |
| running.append( |
| { |
| "pair_id": ",".join(item["pair_id"]), |
| "loss": float(loss.detach()), |
| "hidden_mse": float(hidden_loss.detach()), |
| "velocity_mse": float(velocity_loss.detach()), |
| "micro_time_s": time.perf_counter() - started, |
| "context_frames": sum(item["context_frames"]) |
| / len(item["context_frames"]), |
| } |
| ) |
| del ( |
| output, |
| loss, |
| scaled_loss, |
| hidden_loss, |
| velocity_loss, |
| inputs, |
| target_hidden, |
| target_velocity, |
| item, |
| ) |
|
|
| if micro_step % args.gradient_accumulation_steps: |
| continue |
| grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| optimizer.step() |
| scheduler.step() |
| optimizer.zero_grad(set_to_none=True) |
| global_step += 1 |
|
|
| if global_step % args.log_every == 0 or global_step == 1: |
| torch.cuda.synchronize() |
| count = len(running) |
| record = { |
| "global_step": global_step, |
| "epoch": epoch, |
| "loss": sum(float(x["loss"]) for x in running) / count, |
| "hidden_mse": sum(float(x["hidden_mse"]) for x in running) / count, |
| "velocity_mse": sum(float(x["velocity_mse"]) for x in running) / count, |
| "grad_norm": float(grad_norm), |
| "lr_fusion": optimizer.param_groups[0]["lr"], |
| "lr_blocks": optimizer.param_groups[1]["lr"], |
| "peak_memory_gib": torch.cuda.max_memory_allocated() / 2**30, |
| "avg_micro_time_s": sum(float(x["micro_time_s"]) for x in running) / count, |
| "avg_context_frames": sum(float(x["context_frames"]) for x in running) / count, |
| } |
| with log_path.open("a", encoding="utf-8") as handle: |
| handle.write(json.dumps(record, sort_keys=True) + "\n") |
| swanlab.log(record, step=global_step) |
| print(json.dumps(record, sort_keys=True), flush=True) |
| running.clear() |
|
|
| should_save = global_step % args.save_every == 0 or global_step == args.max_steps |
| if should_save: |
| weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" |
| save_predictor_weights(model, weights_path) |
| if args.save_optimizer_state: |
| latest = { |
| "model": trainable_state_dict(model), |
| "optimizer": optimizer.state_dict(), |
| "scheduler": scheduler.state_dict(), |
| "global_step": global_step, |
| "micro_step": micro_step, |
| "torch_rng_state": torch.get_rng_state(), |
| "cuda_rng_state": torch.cuda.get_rng_state_all(), |
| "args": vars(args), |
| "predictor_config": model.config_dict, |
| "weights_path": str(weights_path), |
| } |
| save_training_state(output_dir / "training_latest.pt", latest) |
| prune_snapshots(output_dir, args.keep_snapshots) |
| except BaseException as error: |
| swanlab.finish(error=error) |
| raise |
| else: |
| swanlab.finish() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
