"""Eight-GPU DDP trainer for the offline Self-Forcing Predictor-v4.""" from __future__ import annotations import gc import json import math import os import random import time from pathlib import Path from typing import Any import torch import torch.distributed as dist import torch.nn.functional as F from omegaconf import OmegaConf from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader from predictor_training import ( DistributedContextBucketBatchSampler, PredictorV4PairDataset, move_batch_to_device, predictor_v4_collate, ) from predictor_training.cache import build_cross_attention_cache from predictor_training.checkpoint import ( atomic_torch_save, capture_rng_state, restore_rng_state, save_predictor_weights, trainable_state_dict, ) def _cosine_with_warmup(step: int, warmup: int, total: 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)) def _make_scheduler( optimizer: torch.optim.Optimizer, *, scheduler_warmup_steps: int, block_warmup_steps: int, max_steps: int, ) -> LambdaLR: return LambdaLR( optimizer, [ lambda step: _cosine_with_warmup( step, scheduler_warmup_steps, max_steps ), lambda step: ( 0.0 if step < block_warmup_steps else _cosine_with_warmup( step - block_warmup_steps, scheduler_warmup_steps, max(1, max_steps - block_warmup_steps), ) ), ], ) def _distributed_mean(values: torch.Tensor) -> torch.Tensor: dist.all_reduce(values, op=dist.ReduceOp.SUM) return values / dist.get_world_size() def _atomic_json(path: Path, payload: dict[str, Any]) -> None: temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") temporary.write_text( json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True, default=str) + "\n", encoding="utf-8", ) os.replace(temporary, path) 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 _load_teacher_and_predictor( cfg, ) -> torch.nn.Module: """Load DMD EMA on CPU and retain only Predictor/frozen projections.""" from model.predictor_v4 import SelfForcingPredictorV4 from utils.wan_wrapper import WanDiffusionWrapper wrapper = WanDiffusionWrapper( model_name=str(cfg.teacher_model_name), timestep_shift=float(cfg.timestep_shift), is_causal=True, local_attn_size=int(cfg.local_attn_size), sink_size=int(cfg.sink_size), ) checkpoint = torch.load( str(cfg.teacher_checkpoint), map_location="cpu", weights_only=False, ) checkpoint_key = str(cfg.teacher_checkpoint_key) if checkpoint_key not in checkpoint: raise KeyError( f"{cfg.teacher_checkpoint} has no {checkpoint_key!r}; " f"available keys are {sorted(checkpoint)}" ) result = wrapper.load_state_dict( checkpoint[checkpoint_key], strict=bool(cfg.strict_teacher_load), ) if not bool(cfg.strict_teacher_load): if result.unexpected_keys: raise RuntimeError( f"Unexpected teacher weights: {result.unexpected_keys[:20]}" ) teacher_model = wrapper.model source_blocks = tuple(int(value) for value in cfg.source_block_ids) predictor = SelfForcingPredictorV4.from_teacher( teacher_model, source_block_ids=source_blocks, ) del checkpoint, wrapper, teacher_model gc.collect() return predictor def _configure_precision( module: torch.nn.Module, *, device: torch.device, activation_dtype: torch.dtype, fp32_trainable_params: bool, ) -> None: module.to(device=device, dtype=activation_dtype) if fp32_trainable_params: with torch.no_grad(): for parameter in module.parameters(): if parameter.requires_grad: parameter.data = parameter.data.float() def _optimizer_groups( model: torch.nn.Module, ) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter], dict[str, str]]: fusion_parameters = [] block_parameters = [] parameter_group_names: dict[str, str] = {} for name, parameter in model.named_parameters(): if not parameter.requires_grad: continue lowered = name.lower() is_block = ( "predictor_blocks" in lowered or "source_blocks" in lowered or "double_blocks" in lowered or lowered.startswith("blocks.") ) if is_block: block_parameters.append(parameter) parameter_group_names[name] = "blocks" else: fusion_parameters.append(parameter) parameter_group_names[name] = "fusion_residual" if not fusion_parameters or not block_parameters: raise RuntimeError( "Could not form both Predictor optimizer groups. " f"fusion={len(fusion_parameters)}, blocks={len(block_parameters)}. " "Block modules must include 'predictor_blocks', 'source_blocks', " "'double_blocks', or begin with 'blocks'." ) return fusion_parameters, block_parameters, parameter_group_names def _set_gradient_checkpointing(model: torch.nn.Module, enabled: bool) -> None: if hasattr(model, "enable_gradient_checkpointing"): try: model.enable_gradient_checkpointing(enabled) except TypeError: if enabled: model.enable_gradient_checkpointing() elif hasattr(model, "disable_gradient_checkpointing"): model.disable_gradient_checkpointing() class Trainer: """Trainer dispatch target used by ``train.py``.""" def __init__(self, config) -> None: if not torch.cuda.is_available(): raise RuntimeError("Predictor-v4 training requires CUDA") required_env = {"RANK", "WORLD_SIZE", "LOCAL_RANK"} if not required_env.issubset(os.environ): raise RuntimeError( "Launch Predictor-v4 training with torchrun; missing " f"{sorted(required_env.difference(os.environ))}" ) self.root_config = config self.cfg = config.predictor_v4 self.rank = int(os.environ["RANK"]) self.world_size = int(os.environ["WORLD_SIZE"]) self.local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(self.local_rank) self.device = torch.device("cuda", self.local_rank) dist.init_process_group(backend="nccl") self.is_main_process = self.rank == 0 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision("high") base_seed = int(self.root_config.seed) seed = base_seed + self.rank random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) self.activation_dtype = torch.bfloat16 self.source_blocks = tuple(int(value) for value in self.cfg.source_block_ids) self.output_dir = Path(str(self.cfg.output_dir)).resolve() if self.is_main_process: self.output_dir.mkdir(parents=True, exist_ok=True) dist.barrier() self.log_path = self.output_dir / "train_log.jsonl" self.dataset = PredictorV4PairDataset( self.cfg.manifest, source_block_ids=self.source_blocks, max_records=( None if self.cfg.max_records is None else int(self.cfg.max_records) ), ) self.sampler = DistributedContextBucketBatchSampler( self.dataset, batch_size=int(self.cfg.batch_size), rank=self.rank, world_size=self.world_size, seed=base_seed, drop_last=True, ) self.loader = DataLoader( self.dataset, batch_sampler=self.sampler, num_workers=int(self.cfg.num_workers), pin_memory=bool(self.cfg.pin_memory), persistent_workers=( int(self.cfg.num_workers) > 0 and bool(self.cfg.persistent_workers) ), collate_fn=predictor_v4_collate, ) if len(self.loader) == 0: raise ValueError( "Predictor-v4 loader has zero batches; lower batch_size or " "provide more records for every context bucket" ) self.model = _load_teacher_and_predictor(self.cfg) _configure_precision( self.model, device=self.device, activation_dtype=self.activation_dtype, fp32_trainable_params=bool(self.cfg.fp32_trainable_params), ) _set_gradient_checkpointing( self.model, bool(self.cfg.gradient_checkpointing) ) self.model.train() fusion_parameters, block_parameters, group_names = _optimizer_groups( self.model ) self.optimizer = AdamW( [ { "params": fusion_parameters, "lr": float(self.cfg.fusion_lr), "name": "fusion_residual", }, { "params": block_parameters, "lr": float(self.cfg.blocks_lr), "name": "blocks", }, ], betas=(float(self.cfg.beta1), float(self.cfg.beta2)), weight_decay=float(self.cfg.weight_decay), ) self.scheduler = _make_scheduler( self.optimizer, scheduler_warmup_steps=int(self.cfg.scheduler_warmup_steps), block_warmup_steps=int(self.cfg.block_warmup_steps), max_steps=int(self.cfg.max_steps), ) self.global_step = 0 self.micro_step = 0 self.epoch = 0 self.batch_in_epoch = 0 self._resume(group_names) self.ddp = DDP( self.model, device_ids=[self.local_rank], output_device=self.local_rank, broadcast_buffers=False, gradient_as_bucket_view=True, find_unused_parameters=bool(self.cfg.find_unused_parameters), ) self.run_config = self._run_config(group_names) self.wandb_run = self._initialize_wandb() self.swanlab_run = self._initialize_swanlab() if self.is_main_process: _atomic_json(self.output_dir / "train_config.json", self.run_config) print(json.dumps(self.run_config, indent=2, default=str), flush=True) def _resume(self, group_names: dict[str, str]) -> None: resume = self.cfg.resume if resume is None or str(resume).lower() in {"", "null", "none"}: return resume_path = Path(str(resume)).resolve() state = torch.load(resume_path, map_location="cpu", weights_only=False) if tuple(state["source_block_ids"]) != self.source_blocks: raise ValueError("Resume source_block_ids do not match current config") if state.get("parameter_groups") != group_names: raise ValueError("Resume optimizer parameter grouping has changed") result = self.model.load_state_dict(state["model"], strict=False) trainable = { name for name, parameter in self.model.named_parameters() if parameter.requires_grad } missing_trainable = sorted(trainable.intersection(result.missing_keys)) if result.unexpected_keys or missing_trainable: raise RuntimeError( "Resume model mismatch: " f"unexpected={result.unexpected_keys}, " f"missing_trainable={missing_trainable}" ) self.optimizer.load_state_dict(state["optimizer"]) self.scheduler.load_state_dict(state["scheduler"]) self.global_step = int(state["global_step"]) self.micro_step = int(state["micro_step"]) self.epoch = int(state["epoch"]) self.batch_in_epoch = int(state["batch_in_epoch"]) rng_path = resume_path.parent / f"rng_rank_{self.rank:02d}.pt" if not rng_path.is_file(): raise FileNotFoundError( f"Full resume requires per-rank RNG checkpoint {rng_path}" ) restore_rng_state( torch.load(rng_path, map_location="cpu", weights_only=False) ) def _run_config(self, parameter_groups: dict[str, str]) -> dict[str, Any]: trainable = sum( parameter.numel() for parameter in self.model.parameters() if parameter.requires_grad ) return { **OmegaConf.to_container(self.cfg, resolve=True), "manifest": str(Path(str(self.cfg.manifest)).resolve()), "teacher_checkpoint": str( Path(str(self.cfg.teacher_checkpoint)).resolve() ), "output_dir": str(self.output_dir), "source_block_ids": list(self.source_blocks), "prompt_count": int(self.cfg.prompt_count), "supervision_pairs": [list(pair) for pair in self.dataset.PAIRS], "manifest_records": len(self.dataset.records), "dataset_pairs": len(self.dataset), "steps_per_epoch_per_rank": len(self.loader), "world_size": self.world_size, "global_batch_size": ( int(self.cfg.batch_size) * self.world_size * int(self.cfg.gradient_accumulation_steps) ), "trainable_parameters": trainable, "parameter_groups": parameter_groups, "activation_dtype": "bfloat16", "trainable_parameter_dtype": ( "float32" if bool(self.cfg.fp32_trainable_params) else "bfloat16" ), } def _initialize_wandb(self): if ( not self.is_main_process or bool(self.root_config.disable_wandb) or not bool(self.cfg.use_wandb) ): return None import wandb return wandb.init( project=str(self.cfg.wandb_project), entity=( None if self.cfg.wandb_entity is None else str(self.cfg.wandb_entity) ), name=str(self.cfg.wandb_name), dir=str(self.output_dir), config=self.run_config, ) def _initialize_swanlab(self): if ( not self.is_main_process or not bool(getattr(self.cfg, "use_swanlab", False)) ): return None import swanlab mode = str(getattr(self.cfg, "swanlab_mode", "cloud")) if mode == "cloud": api_key = os.environ.get("SWANLAB_API_KEY") if api_key: swanlab.login(api_key=api_key, save=False) else: # Reuse the host's existing ~/.swanlab/.netrc credential. swanlab.login() workspace = getattr(self.cfg, "swanlab_workspace", None) run = swanlab.init( project=str(self.cfg.swanlab_project), workspace=None if workspace is None else str(workspace), experiment_name=str(self.cfg.swanlab_experiment), description=str(self.cfg.swanlab_description), tags=list(self.cfg.swanlab_tags), config=self.run_config, logdir=str(self.output_dir / "swanlab"), mode=mode, ) run_id = getattr(run, "id", None) self.run_config["swanlab_run_id"] = run_id print(json.dumps({"swanlab_run_id": run_id}), flush=True) return run def _save(self) -> None: dist.barrier() atomic_torch_save( capture_rng_state(), self.output_dir / f"rng_rank_{self.rank:02d}.pt", ) if self.is_main_process: weights_path = ( self.output_dir / f"predictor_step_{self.global_step:05d}.safetensors" ) save_predictor_weights( self.model, weights_path, metadata={ "source_block_ids": list(self.source_blocks), "global_step": self.global_step, "teacher_checkpoint": str(self.cfg.teacher_checkpoint), "teacher_checkpoint_key": str( self.cfg.teacher_checkpoint_key ), "predictor_config": self.model.config_dict, "schema_version": str(self.cfg.schema_version), }, ) atomic_torch_save( { "model": trainable_state_dict(self.model), "optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict(), "global_step": self.global_step, "micro_step": self.micro_step, "epoch": self.epoch, "batch_in_epoch": self.batch_in_epoch, "source_block_ids": self.source_blocks, "parameter_groups": self.run_config["parameter_groups"], "config": self.run_config, "weights_path": str(weights_path), }, self.output_dir / "training_latest.pt", ) _prune_snapshots(self.output_dir, int(self.cfg.keep_snapshots)) dist.barrier() def _forward_loss( self, batch: dict[str, Any], ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: context_frames_values = set(int(value) for value in batch["context_frames"]) if len(context_frames_values) != 1: raise ValueError("A Predictor batch contains mixed Context lengths") context_frames = context_frames_values.pop() current_tokens = int(batch["anchor_hidden"].shape[1]) tokens_per_frame = int(self.model.predictor_config.tokens_per_frame) current_start = context_frames * tokens_per_frame with torch.no_grad(): kv_cache = self.model.build_history_kv_cache( batch["clean_prefeature"], current_start=current_start, current_tokens=current_tokens, start_frames=0, ) crossattn_cache = build_cross_attention_cache( batch["text_kv"], self.source_blocks ) output = self.ddp( target_latent=batch["target_latent"], target_timestep=batch["target_timestep"], anchor_hidden=batch["anchor_hidden"], previous_chunk_hidden=batch["previous_chunk_hidden"], kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, ) if not isinstance(output, dict) or not { "pred_hidden", "pred_flow", }.issubset(output): raise TypeError( "SelfForcingPredictorV4.forward must return a dict containing " "pred_hidden and pred_flow" ) hidden_loss = F.mse_loss( output["pred_hidden"].float(), batch["target_hidden"].float() ) flow_loss = F.mse_loss( output["pred_flow"].float(), batch["target_flow"].float() ) loss = ( float(self.cfg.hidden_loss_weight) * hidden_loss + float(self.cfg.flow_loss_weight) * flow_loss ) return loss, hidden_loss, flow_loss def train(self) -> None: accumulation = int(self.cfg.gradient_accumulation_steps) max_steps = int(self.cfg.max_steps) self.optimizer.zero_grad(set_to_none=True) running = torch.zeros(4, device=self.device, dtype=torch.float64) running_count = 0 try: while self.global_step < max_steps: self.sampler.set_epoch(self.epoch) for batch_index, cpu_batch in enumerate(self.loader): if batch_index < self.batch_in_epoch: continue self.batch_in_epoch = batch_index + 1 started = time.perf_counter() batch = move_batch_to_device( cpu_batch, device=self.device, dtype=self.activation_dtype, ) sync_gradients = (self.micro_step + 1) % accumulation == 0 sync_context = ( torch.enable_grad() if sync_gradients else self.ddp.no_sync() ) with sync_context: with torch.autocast( device_type="cuda", dtype=self.activation_dtype ): loss, hidden_loss, flow_loss = self._forward_loss(batch) scaled_loss = loss / accumulation scaled_loss.backward() self.micro_step += 1 running += torch.tensor( [ float(loss.detach()), float(hidden_loss.detach()), float(flow_loss.detach()), time.perf_counter() - started, ], device=self.device, dtype=torch.float64, ) running_count += 1 del ( cpu_batch, batch, loss, hidden_loss, flow_loss, scaled_loss, ) if not sync_gradients: continue grad_norm = torch.nn.utils.clip_grad_norm_( self.model.parameters(), float(self.cfg.grad_clip) ) self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad(set_to_none=True) self.global_step += 1 if ( self.global_step == 1 or self.global_step % int(self.cfg.log_every) == 0 ): torch.cuda.synchronize(self.device) averaged = _distributed_mean(running.clone()) averaged /= running_count max_grad = torch.tensor( float(grad_norm), device=self.device, dtype=torch.float64, ) dist.all_reduce(max_grad, op=dist.ReduceOp.MAX) if self.is_main_process: record = { "global_step": self.global_step, "epoch": self.epoch, "loss": float(averaged[0]), "hidden_mse": float(averaged[1]), "flow_mse": float(averaged[2]), "avg_micro_time_s": float(averaged[3]), "grad_norm_max": float(max_grad), "lr_fusion": self.optimizer.param_groups[0]["lr"], "lr_blocks": self.optimizer.param_groups[1]["lr"], "peak_memory_gib": ( torch.cuda.max_memory_allocated(self.device) / 2**30 ), } with self.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) if self.wandb_run is not None: self.wandb_run.log(record, step=self.global_step) if self.swanlab_run is not None: import swanlab swanlab.log(record, step=self.global_step) running.zero_() running_count = 0 should_save = ( not bool(self.root_config.no_save) and ( self.global_step % int(self.cfg.save_every) == 0 or self.global_step == max_steps ) ) if should_save: self._save() if self.global_step >= max_steps: break if self.global_step < max_steps: self.epoch += 1 self.batch_in_epoch = 0 except BaseException as error: if self.swanlab_run is not None: try: import swanlab swanlab.finish(error=str(error)) except Exception: pass raise else: if self.swanlab_run is not None: import swanlab swanlab.finish() finally: if self.wandb_run is not None: self.wandb_run.finish() dist.destroy_process_group()