"""Detached chunk+timestep FPPP rollout training for Self-Forcing Predictor-v4.""" from __future__ import annotations import gc import json import math import os import random import time from contextlib import nullcontext 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 torch.utils.data.distributed import DistributedSampler from model.predictor_v4 import SelfForcingPredictorV4 from pipeline.causal_inference import CausalInferencePipeline from predictor_training.checkpoint import ( atomic_torch_save, capture_rng_state, restore_rng_state, save_predictor_weights, trainable_state_dict, ) from predictor_training.rollout_cache import ( assert_clean_history_extent, build_predictor_workspace, reset_main_caches, reset_predictor_workspace, ) from predictor_training.trajectory_dataset import ( PredictorV4TrajectoryDataset, load_offline_ffff_target, trajectory_collate, ) from utils.misc import set_seed CHUNK_FRAMES = 3 NUM_CHUNKS = 7 NUM_STEPS = 4 TRAIN_STEPS = (1, 2, 3) TOKENS_PER_FRAME = 1560 TOKENS_PER_CHUNK = CHUNK_FRAMES * TOKENS_PER_FRAME EXPECTED_TIMESTEPS = torch.tensor( [1000.0, 937.5, 833.3333129882812, 625.0], dtype=torch.float32 ) def _atomic_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") temporary.write_text( json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True, default=str) + "\n", encoding="utf-8", ) os.replace(temporary, path) def _distributed_mean(value: torch.Tensor) -> torch.Tensor: dist.all_reduce(value, op=dist.ReduceOp.SUM) return value / dist.get_world_size() def _cosine_with_linear_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, *, warmup_steps: int, max_steps: int, ) -> LambdaLR: return LambdaLR( optimizer, [ lambda step: _cosine_with_linear_warmup( step, warmup_steps, max_steps ), lambda step: _cosine_with_linear_warmup( step, warmup_steps, max_steps ), ], ) def _optimizer_groups( model: torch.nn.Module, ) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter], dict[str, str]]: fusion: list[torch.nn.Parameter] = [] blocks: list[torch.nn.Parameter] = [] names: dict[str, str] = {} for name, parameter in model.named_parameters(): if not parameter.requires_grad: continue if "predictor_blocks" in name: blocks.append(parameter) names[name] = "blocks" else: fusion.append(parameter) names[name] = "fusion_residual" if not fusion or not blocks: raise RuntimeError( f"Invalid Stage-2 optimizer groups: fusion={len(fusion)}, " f"blocks={len(blocks)}" ) return fusion, blocks, names 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_stage1_model_state( model: SelfForcingPredictorV4, path: Path, *, expected_step: int, ) -> None: state = torch.load(path, map_location="cpu", weights_only=False) actual_step = int(state.get("global_step", -1)) if actual_step != int(expected_step): raise ValueError( f"Stage-1 initialization must be step {expected_step}, got {actual_step}" ) expected_blocks = tuple(model.source_block_ids) saved_blocks = tuple(int(value) for value in state["source_block_ids"]) if saved_blocks != expected_blocks: raise ValueError( f"Stage-1 source blocks {saved_blocks} != requested {expected_blocks}" ) result = model.load_state_dict(state["model"], strict=False) trainable = { name for name, parameter in model.named_parameters() if parameter.requires_grad } missing_trainable = sorted(trainable.intersection(result.missing_keys)) if result.unexpected_keys or missing_trainable: raise RuntimeError( "Stage-1 Predictor state mismatch: " f"unexpected={result.unexpected_keys}, " f"missing_trainable={missing_trainable}" ) del state def _configure_predictor_precision( model: SelfForcingPredictorV4, *, device: torch.device, ) -> None: model.to(device=device, dtype=torch.bfloat16) with torch.no_grad(): for parameter in model.parameters(): if parameter.requires_grad: parameter.data = parameter.data.float() def _load_pipeline_and_predictor( cfg, *, device: torch.device, ) -> tuple[CausalInferencePipeline, SelfForcingPredictorV4]: dmd_config = OmegaConf.merge( OmegaConf.load(str(cfg.default_config)), OmegaConf.load(str(cfg.dmd_config)), ) if int(dmd_config.num_frame_per_block) != CHUNK_FRAMES: raise ValueError("Stage-2 rollout requires three latent frames per chunk") if bool(dmd_config.independent_first_frame): raise ValueError("Stage-2 rollout requires independent_first_frame=false") if bool(getattr(dmd_config, "reuse_first_step_velocity", False)): raise ValueError("Stage-2 Teacher must use Full FFFF denoising") pipeline = CausalInferencePipeline( dmd_config, device=device, vae=torch.nn.Identity(), ) 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} lacks checkpoint key {checkpoint_key!r}" ) pipeline.generator.load_state_dict( checkpoint[checkpoint_key], strict=bool(cfg.strict_teacher_load), ) del checkpoint predictor = SelfForcingPredictorV4.from_teacher( pipeline.generator.model, source_block_ids=tuple(int(value) for value in cfg.source_block_ids), ) pipeline.generator.eval().requires_grad_(False) pipeline.text_encoder.eval().requires_grad_(False) pipeline.generator.to(device=device, dtype=torch.bfloat16) pipeline.text_encoder.to(device=device, dtype=torch.bfloat16) _configure_predictor_precision(predictor, device=device) # Load the Stage-1 FP32 master state only after trainable parameters have # been promoted back to FP32; loading before .to(bfloat16) would silently # quantize the initialization. _load_stage1_model_state( predictor, Path(str(cfg.stage1_training_state)).resolve(), expected_step=int(cfg.stage1_expected_step), ) predictor.train() pipeline._initialize_kv_cache( batch_size=1, dtype=torch.bfloat16, device=device ) pipeline._initialize_crossattn_cache( batch_size=1, dtype=torch.bfloat16, device=device ) actual = pipeline.denoising_step_list.detach().float().cpu() if not torch.equal(actual, EXPECTED_TIMESTEPS): raise ValueError( f"Stage-2 requires exact FP32 timesteps {EXPECTED_TIMESTEPS.tolist()}, " f"got {actual.tolist()}" ) return pipeline, predictor def _fixed_noise_trajectory( *, device: torch.device, ) -> tuple[torch.Tensor, list[list[torch.Tensor]]]: """Reproduce inference seed-0 initial and per-transition noise draws.""" set_seed(0) initial = torch.randn( [1, NUM_CHUNKS * CHUNK_FRAMES, 16, 60, 104], device=device, dtype=torch.bfloat16, ) transition_noise = [ [ torch.randn( [1, CHUNK_FRAMES, 16, 60, 104], device=device, dtype=torch.bfloat16, ) for _ in range(NUM_STEPS - 1) ] for _ in range(NUM_CHUNKS) ] return initial, transition_noise def _timestep( timesteps: torch.Tensor, step_id: int, *, device: torch.device, ) -> torch.Tensor: return torch.full( (1, CHUNK_FRAMES), float(timesteps[int(step_id)].item()), device=device, dtype=torch.float32, ) def _renoise( pipeline: CausalInferencePipeline, clean: torch.Tensor, noise: torch.Tensor, *, next_timestep: torch.Tensor, ) -> torch.Tensor: return pipeline.scheduler.add_noise( clean.detach().flatten(0, 1), noise.flatten(0, 1), next_timestep.flatten(0, 1), ).unflatten(0, clean.shape[:2]).detach() @torch.no_grad() def _full_step( pipeline: CausalInferencePipeline, *, latent: torch.Tensor, conditional_dict: dict[str, torch.Tensor], timestep: torch.Tensor, current_start: int, capture_hidden: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: with torch.autocast(device_type="cuda", dtype=torch.bfloat16): flow, clean, hidden = pipeline._full_step_with_optional_hidden( noisy_input=latent, conditional_dict=conditional_dict, timestep=timestep, current_start=current_start, capture_hidden=capture_hidden, ) return flow.detach(), clean.detach(), None if hidden is None else hidden.detach() @torch.no_grad() def _commit_clean_history( pipeline: CausalInferencePipeline, *, clean: torch.Tensor, conditional_dict: dict[str, torch.Tensor], current_start: int, expected_end: int, ) -> None: timestep = torch.full( (1, clean.shape[1]), float(pipeline.args.context_noise), device=clean.device, dtype=torch.float32, ) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pipeline.generator( noisy_image_or_video=clean, conditional_dict=conditional_dict, timestep=timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=current_start, ) assert_clean_history_extent( pipeline.kv_cache1, expected_tokens=expected_end ) class Trainer: """Trainer dispatch target for both Stage-2 supervision variants.""" def __init__(self, config) -> None: if not torch.cuda.is_available(): raise RuntimeError("Predictor-v4 Stage-2 requires CUDA") required = {"RANK", "WORLD_SIZE", "LOCAL_RANK"} if not required.issubset(os.environ): raise RuntimeError("Launch Predictor-v4 Stage-2 with torchrun") self.root_config = config self.cfg = config.predictor_v4_rollout self.rank = int(os.environ["RANK"]) self.world_size = int(os.environ["WORLD_SIZE"]) self.local_rank = int(os.environ["LOCAL_RANK"]) if self.world_size != 8: raise ValueError(f"Formal Stage-2 requires 8 ranks, got {self.world_size}") torch.cuda.set_device(self.local_rank) self.device = torch.device("cuda", self.local_rank) dist.init_process_group(backend="nccl") self.is_main = self.rank == 0 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision("high") seed = int(self.root_config.seed) random.seed(seed + self.rank) torch.manual_seed(seed + self.rank) torch.cuda.manual_seed_all(seed + self.rank) self.output_dir = Path(str(self.cfg.output_dir)).resolve() if self.is_main: self.output_dir.mkdir(parents=True, exist_ok=True) dist.barrier() self.log_path = self.output_dir / "train_log.jsonl" self.dataset = PredictorV4TrajectoryDataset( self.cfg.cases_path, data_root=self.cfg.data_root, max_cases=( None if self.cfg.max_cases is None else int(self.cfg.max_cases) ), ) self.sampler = DistributedSampler( self.dataset, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=seed, drop_last=True, ) self.loader = DataLoader( self.dataset, batch_size=1, 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=trajectory_collate, ) if not self.loader: raise ValueError("Stage-2 trajectory loader is empty") self.pipeline, self.model = _load_pipeline_and_predictor( self.cfg, device=self.device ) fusion, blocks, group_names = _optimizer_groups(self.model) self.optimizer = AdamW( [ { "params": fusion, "lr": float(self.cfg.fusion_lr), "name": "fusion_residual", }, { "params": blocks, "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, warmup_steps=int(self.cfg.warmup_steps), max_steps=int(self.cfg.max_steps), ) self.global_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=False, ) self.run_config = self._run_config(group_names) self.swanlab_run = self._initialize_swanlab() if self.is_main: _atomic_json(self.output_dir / "train_config.json", self.run_config) print(json.dumps(self.run_config, indent=2, default=str), flush=True) @property def supervision_mode(self) -> str: return str(self.cfg.supervision_mode) def _resume(self, group_names: dict[str, str]) -> None: resume = self.cfg.resume if resume is None or str(resume).lower() in {"", "none", "null"}: return path = Path(str(resume)).resolve() state = torch.load(path, map_location="cpu", weights_only=False) if str(state["supervision_mode"]) != self.supervision_mode: raise ValueError("Stage-2 resume supervision mode changed") if state["parameter_groups"] != group_names: raise ValueError("Stage-2 resume optimizer grouping 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( "Stage-2 resume 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.epoch = int(state["epoch"]) self.batch_in_epoch = int(state["batch_in_epoch"]) rng_path = path.parent / f"rng_rank_{self.rank:02d}.pt" if not rng_path.is_file(): raise FileNotFoundError(rng_path) restore_rng_state( torch.load(rng_path, map_location="cpu", weights_only=False) ) def _run_config(self, group_names: dict[str, str]) -> dict[str, Any]: return { **OmegaConf.to_container(self.cfg, resolve=True), "trainer_semantics": ( "chunk0_ffff_chunks1-6_fppp_detached_timestep_and_chunk" ), "train_predictor_steps": [1, 2, 3], "persistent_cache_semantics": "final_clean_full_pass_only", "exact_online_timesteps": EXPECTED_TIMESTEPS.tolist(), "offline_target_timestep_semantics": ( "saved_stage1_int64" if self.supervision_mode == "offline_ffff" else None ), "world_size": self.world_size, "global_batch_trajectories": self.world_size, "supervised_states_per_optimizer_step": ( self.world_size * (NUM_CHUNKS - 1) * len(TRAIN_STEPS) ), "parameter_groups": group_names, "trainable_parameters": sum( parameter.numel() for parameter in self.model.parameters() if parameter.requires_grad ), "activation_dtype": "bfloat16", "trainable_parameter_dtype": "float32", } def _initialize_swanlab(self): if not self.is_main or not bool(self.cfg.use_swanlab): return None import swanlab mode = str(self.cfg.swanlab_mode) if mode == "cloud": api_key = os.environ.get("SWANLAB_API_KEY") if api_key: swanlab.login(api_key=api_key, save=False) else: swanlab.login() workspace = self.cfg.swanlab_workspace 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, ) self.run_config["swanlab_run_id"] = getattr(run, "id", None) 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: 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.model.source_block_ids), "global_step": self.global_step, "stage": 2, "training_rollout": "fppp", "supervision_mode": self.supervision_mode, "predictor_config": self.model.config_dict, }, ) atomic_torch_save( { "model": trainable_state_dict(self.model), "optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict(), "global_step": self.global_step, "epoch": self.epoch, "batch_in_epoch": self.batch_in_epoch, "supervision_mode": self.supervision_mode, "source_block_ids": tuple(self.model.source_block_ids), "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 _target( self, *, case_id: int, chunk_id: int, step_id: int, latent: torch.Tensor, conditional_dict: dict[str, torch.Tensor], timestep: torch.Tensor, current_start: int, ) -> tuple[torch.Tensor, torch.Tensor]: if self.supervision_mode == "onpolicy": flow, _, hidden = _full_step( self.pipeline, latent=latent, conditional_dict=conditional_dict, timestep=timestep, current_start=current_start, capture_hidden=True, ) if hidden is None: raise RuntimeError("On-policy Full Teacher did not return hidden") return flow, hidden if self.supervision_mode == "offline_ffff": target = load_offline_ffff_target( self.dataset.offline_step_path(case_id, chunk_id), step_id=step_id, device=self.device, ) return target["flow"], target["hidden"] raise ValueError(f"Unknown supervision mode {self.supervision_mode!r}") def _run_trajectory(self, record: dict[str, Any]) -> dict[str, float]: started = time.perf_counter() case_id = int(record["case_id"]) if int(record["seed"]) != 0: raise ValueError(f"case {case_id} does not use seed 0") reset_main_caches( self.pipeline.kv_cache1, self.pipeline.crossattn_cache ) initial_noise, transition_noise = _fixed_noise_trajectory( device=self.device ) with torch.no_grad(), torch.autocast( device_type="cuda", dtype=torch.bfloat16 ): conditional_dict = self.pipeline.text_encoder( text_prompts=[str(record["prompt"])] ) timesteps = self.pipeline.denoising_step_list.to( device=self.device, dtype=torch.float32 ) previous_hidden: dict[int, torch.Tensor] = {} metric_sums = torch.zeros(7, device=self.device, dtype=torch.float64) # total, h1, h2, h3, flow1, flow2, flow3 num_predictor_calls = (NUM_CHUNKS - 1) * len(TRAIN_STEPS) predictor_call = 0 self.optimizer.zero_grad(set_to_none=True) for chunk_id in range(NUM_CHUNKS): current_start = chunk_id * TOKENS_PER_CHUNK current_end = current_start + TOKENS_PER_CHUNK latent = initial_noise[ :, chunk_id * CHUNK_FRAMES : (chunk_id + 1) * CHUNK_FRAMES ].detach() predictor_cache = None if chunk_id > 0: predictor_cache = build_predictor_workspace( self.pipeline.kv_cache1, source_block_ids=tuple(self.model.source_block_ids), history_tokens=current_start, current_tokens=TOKENS_PER_CHUNK, ) if chunk_id == 0: chunk_hidden: dict[int, torch.Tensor] = {} final_clean = None for step_id in range(NUM_STEPS): timestep = _timestep( timesteps, step_id, device=self.device ) _, clean, hidden = _full_step( self.pipeline, latent=latent, conditional_dict=conditional_dict, timestep=timestep, current_start=current_start, capture_hidden=step_id in TRAIN_STEPS, ) if step_id in TRAIN_STEPS: if hidden is None: raise RuntimeError("Chunk 0 Full hidden is missing") chunk_hidden[step_id] = hidden final_clean = clean if step_id < NUM_STEPS - 1: latent = _renoise( self.pipeline, clean, transition_noise[chunk_id][step_id], next_timestep=_timestep( timesteps, step_id + 1, device=self.device ), ) if final_clean is None: raise RuntimeError("Chunk 0 did not produce a final clean latent") _commit_clean_history( self.pipeline, clean=final_clean, conditional_dict=conditional_dict, current_start=current_start, expected_end=current_end, ) previous_hidden = { step: chunk_hidden[step].detach() for step in TRAIN_STEPS } del chunk_hidden, final_clean continue timestep0 = _timestep(timesteps, 0, device=self.device) _, clean0, anchor = _full_step( self.pipeline, latent=latent, conditional_dict=conditional_dict, timestep=timestep0, current_start=current_start, capture_hidden=True, ) if anchor is None: raise RuntimeError(f"Chunk {chunk_id} has no Full step-0 anchor") latent = _renoise( self.pipeline, clean0, transition_noise[chunk_id][0], next_timestep=_timestep(timesteps, 1, device=self.device), ) current_hidden: dict[int, torch.Tensor] = {} for step_id in TRAIN_STEPS: reset_predictor_workspace( predictor_cache, history_tokens=current_start, ) if latent.requires_grad or anchor.requires_grad: raise RuntimeError( f"Chunk {chunk_id} step {step_id} received a " "gradient-connected detached rollout state" ) if previous_hidden[step_id].requires_grad: raise RuntimeError( f"Chunk {chunk_id} step {step_id} previous hidden " "retained a graph across chunks" ) if any( cache["k"].requires_grad or cache["v"].requires_grad for cache in predictor_cache.values() ): raise RuntimeError( f"Chunk {chunk_id} step {step_id} Predictor workspace " "retained an earlier timestep graph" ) timestep = _timestep(timesteps, step_id, device=self.device) target_flow, target_hidden = self._target( case_id=case_id, chunk_id=chunk_id, step_id=step_id, latent=latent, conditional_dict=conditional_dict, timestep=timestep, current_start=current_start, ) predictor_call += 1 sync = predictor_call == num_predictor_calls sync_context = nullcontext() if sync else self.ddp.no_sync() with sync_context: with torch.autocast( device_type="cuda", dtype=torch.bfloat16 ): output = self.ddp( target_latent=latent, target_timestep=timestep, anchor_hidden=anchor, previous_chunk_hidden=previous_hidden[step_id], kv_cache=predictor_cache, crossattn_cache=self.pipeline.crossattn_cache, current_start=current_start, ) hidden_loss = F.mse_loss( output["pred_hidden"].float(), target_hidden.float(), ) flow_loss = F.mse_loss( output["pred_flow"].float(), target_flow.float(), ) step_loss = ( float(self.cfg.flow_loss_weight) * flow_loss + float(self.cfg.hidden_loss_weight) * hidden_loss ) pred_hidden = output["pred_hidden"].detach() pred_flow = output["pred_flow"].detach() (step_loss / num_predictor_calls).backward() metric_sums[0] += step_loss.detach() metric_sums[step_id] += hidden_loss.detach() metric_sums[3 + step_id] += flow_loss.detach() current_hidden[step_id] = pred_hidden anchor = pred_hidden pred_clean = self.pipeline.generator._convert_flow_pred_to_x0( flow_pred=pred_flow.flatten(0, 1), xt=latent.flatten(0, 1), timestep=timestep.flatten(0, 1), ).unflatten(0, pred_flow.shape[:2]).detach() if step_id < NUM_STEPS - 1: latent = _renoise( self.pipeline, pred_clean, transition_noise[chunk_id][step_id], next_timestep=_timestep( timesteps, step_id + 1, device=self.device ), ) del ( output, target_flow, target_hidden, hidden_loss, flow_loss, step_loss, pred_flow, ) # Training rollout is always FPPP: P3 supplies both final latent and # next-chunk step-3 hidden. Only this final clean state is committed. _commit_clean_history( self.pipeline, clean=pred_clean, conditional_dict=conditional_dict, current_start=current_start, expected_end=current_end, ) previous_hidden = { step: current_hidden[step].detach() for step in TRAIN_STEPS } del ( predictor_cache, current_hidden, pred_clean, clean0, anchor, ) if predictor_call != num_predictor_calls: raise RuntimeError( f"Expected {num_predictor_calls} Predictor calls, got " f"{predictor_call}" ) grad_norm = torch.nn.utils.clip_grad_norm_( self.model.parameters(), float(self.cfg.grad_clip) ) if not torch.isfinite(grad_norm): raise FloatingPointError(f"Non-finite Stage-2 grad norm {grad_norm}") self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad(set_to_none=True) metric_sums[0] /= num_predictor_calls metric_sums[1:] /= NUM_CHUNKS - 1 result = { "loss": float(metric_sums[0]), "hidden_mse_step1": float(metric_sums[1]), "hidden_mse_step2": float(metric_sums[2]), "hidden_mse_step3": float(metric_sums[3]), "flow_mse_step1": float(metric_sums[4]), "flow_mse_step2": float(metric_sums[5]), "flow_mse_step3": float(metric_sums[6]), "grad_norm": float(grad_norm), "trajectory_time_s": time.perf_counter() - started, } del ( conditional_dict, initial_noise, transition_noise, previous_hidden, metric_sums, ) return result def train(self) -> None: max_steps = int(self.cfg.max_steps) try: while self.global_step < max_steps: self.sampler.set_epoch(self.epoch) for batch_index, record in enumerate(self.loader): if batch_index < self.batch_in_epoch: continue self.batch_in_epoch = batch_index + 1 metrics = self._run_trajectory(record) self.global_step += 1 values = torch.tensor( [ metrics["loss"], metrics["hidden_mse_step1"], metrics["hidden_mse_step2"], metrics["hidden_mse_step3"], metrics["flow_mse_step1"], metrics["flow_mse_step2"], metrics["flow_mse_step3"], metrics["trajectory_time_s"], ], device=self.device, dtype=torch.float64, ) averaged = _distributed_mean(values) max_grad = torch.tensor( metrics["grad_norm"], device=self.device, dtype=torch.float64, ) dist.all_reduce(max_grad, op=dist.ReduceOp.MAX) if ( self.global_step == 1 or self.global_step % int(self.cfg.log_every) == 0 ): record_out = { "global_step": self.global_step, "epoch": self.epoch, "supervision_mode": self.supervision_mode, "loss": float(averaged[0]), "hidden_mse_step1": float(averaged[1]), "hidden_mse_step2": float(averaged[2]), "hidden_mse_step3": float(averaged[3]), "flow_mse_step1": float(averaged[4]), "flow_mse_step2": float(averaged[5]), "flow_mse_step3": float(averaged[6]), "trajectory_time_s": float(averaged[7]), "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 ), } if self.is_main: with self.log_path.open( "a", encoding="utf-8" ) as handle: handle.write( json.dumps(record_out, sort_keys=True) + "\n" ) print(json.dumps(record_out, sort_keys=True), flush=True) if self.swanlab_run is not None: import swanlab swanlab.log(record_out, step=self.global_step) 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() del metrics, values, averaged, max_grad 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: gc.collect() torch.cuda.empty_cache() dist.destroy_process_group() __all__ = ["Trainer"]