#!/usr/bin/env python3 """Single-GPU training for the dense image/text HY-WorldPlay Predictor.""" from __future__ import annotations import argparse import json import math import os import random import time from pathlib import Path 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 HYWorldPlayPredictor from predictor_training.checkpoint import ( load_predictor_weights, save_predictor_weights, save_training_state, trainable_state_dict, ) from predictor_training.dataset import ( PredictorPairDataset, move_pair_to_device, predictor_collate, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--manifest", default="datasets/predictor_v1/train_eval_manifest.jsonl") parser.add_argument( "--teacher_checkpoint", default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors", ) parser.add_argument( "--output_dir", default=( "checkpoints/predictor_v1_blocks1-52_bf16_1gpu_" "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100" ), ) parser.add_argument("--max_steps", type=int, default=2000) 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) 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_blocks_trainable(model: HYWorldPlayPredictor, enabled: bool) -> None: model.double_blocks.requires_grad_(enabled) 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 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 = PredictorPairDataset(args.manifest, max_records=args.max_records) generator = torch.Generator().manual_seed(args.seed) loader = DataLoader( dataset, batch_size=1, shuffle=True, num_workers=args.num_workers, pin_memory=True, collate_fn=predictor_collate, generator=generator, ) model = HYWorldPlayPredictor() model.load_teacher_initialization(args.teacher_checkpoint) model.to(device=device, dtype=dtype) model.enable_gradient_checkpointing(args.gradient_checkpointing) model.train() optimizer = AdamW( [ { "params": list(model.img_fusion.parameters()) + list(model.txt_fusion.parameters()) + list(model.residual_out.parameters()), "lr": args.fusion_lr, }, {"params": model.double_blocks.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_blocks_trainable(model, args.warmup_blocks_steps <= 0) if args.resume: state = torch.load(args.resume, map_location="cpu", weights_only=False) 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_blocks_trainable(model, global_step >= args.warmup_blocks_steps) print( json.dumps( { "dataset_pairs": len(dataset), "trainable_parameters": model.trainable_parameter_count(), "breakdown": model.trainable_parameter_breakdown(), "start_step": global_step, }, indent=2, ), flush=True, ) optimizer.zero_grad(set_to_none=True) running: list[dict[str, float | str]] = [] epoch = 0 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_blocks_trainable(model, True) inputs, target_hidden, target_velocity = move_pair_to_device( 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": 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, } ) del output, loss, scaled_loss, hidden_loss, velocity_loss 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, } with log_path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True) + "\n") 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) if __name__ == "__main__": main()