train_predictor.py

#!/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()