"""Training orchestration for distribution matching (paper Section 3.3). Start at _run_distribution_matching_step: collect the distribution batch, form the feature-space objective, replay the differentiation batch, and update only the generator. The long initialization/resume code preserves the released asset and checkpoint contracts; it does not define a separate score model. """ import copy import gc import hashlib import logging import math from collections.abc import Mapping from utils.dataset import ShardingLMDBDataset, cycle from utils.dataset import ( OpenVidPromptVideoDataset, ODERegressionLMDBDataset, TextDataset, ) from utils.distributed import ( EMA_FSDP, fsdp_wrap, fsdp_state_dict, fsdp_optim_state_dict, fsdp_load_optim_state_dict, launch_distributed_job, ) from utils.misc import ( init_wandb_from_config, set_seed, merge_dict_list, ) from trainer.fresh_video_gradcache import FreshVideoGradCache from trainer.feature_gradient_replay import ( FeatureGradientReplay, ) import torch.distributed as dist from omegaconf import OmegaConf from model import DistributionMatchingModel import numpy as np import random import imageio.v3 as iio import torch from utils.optional_wandb import wandb import time import os import shutil def _canonicalize_experiment_identity_paths(value): """Make shared asset identities independent of a compute-node mount alias.""" local_root = "/data/shared/video-generation-assets" logical_root = "/assets/video-generation-assets" if isinstance(value, str): if value == local_root or value.startswith(local_root + "/"): return logical_root + value[len(local_root):] return value if isinstance(value, Mapping): return { key: _canonicalize_experiment_identity_paths(item) for key, item in value.items() } if isinstance(value, tuple): return tuple( _canonicalize_experiment_identity_paths(item) for item in value) if isinstance(value, list): return [ _canonicalize_experiment_identity_paths(item) for item in value] return value def _dense_video_metric_contract(feature_type): """Return the backend-owned prefix and universally required replay keys.""" prefixes = { "vjepa2_dense_latent": "vjepa2_dense", "dinov3_dense_latent": "dinov3_dense", "dual_dinov3_vjepa2_dense_latent": "dinov3_vjepa2_dense", "triple_dinov3_vjepa2_videomae_v2_dense_latent": "triple_dense", "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent": ( "triple_dense"), } try: prefix = prefixes[str(feature_type)] except KeyError as exc: raise ValueError( "fresh-video GradCache requires a supported dense feature type; " f"got {feature_type!r}" ) from exc return prefix, frozenset({ "combined_latent_grad_norm", f"{prefix}_feature_grad_norm", f"{prefix}_video_grad_norm", f"{prefix}_replay_max_abs", f"{prefix}_latent_replay_max_abs", }) def _validated_resume_data_counter_override(value, *, resume_checkpoint): """Validate an explicit lightweight-resume prompt-stream position. The counter is deliberately never derived from ``step`` or the current world size. Changing the rank topology changes how many per-rank batches correspond to one optimizer update, so callers must provide the desired absolute per-rank batch counter explicitly. """ if value is None: return None if not resume_checkpoint: raise ValueError( "resume_data_counter_override is only valid with " "resume_checkpoint") if type(value) is not int: raise TypeError( "resume_data_counter_override must be a non-negative integer") if value < 0: raise ValueError( "resume_data_counter_override must be a non-negative integer") return value def _validated_fresh_video_generator_grad_sample_scale( value, *, partial_gradient_enabled, ): """Validate the explicit suffix-replay gradient compensation.""" try: scale = float(value) except (TypeError, ValueError) as exc: raise ValueError( "fresh_video_generator_grad_sample_scale must be finite and " "strictly positive" ) from exc if not math.isfinite(scale) or scale <= 0.0: raise ValueError( "fresh_video_generator_grad_sample_scale must be finite and " "strictly positive") if scale != 1.0 and not bool(partial_gradient_enabled): raise ValueError( "fresh_video_generator_grad_sample_scale may differ from 1 only " "when a strict generator-gradient subsample is enabled") return scale def _generator_lr_for_completed_step(config, completed_step): """Return the LR used by the optimizer update producing ``completed_step``. A cosine window keeps ``config.lr`` through its inclusive start step, then decays to the configured minimum at the inclusive end step. Naming the public/checkpoint step explicitly avoids a resume off-by-one: checkpoint 75 next performs the update that publishes checkpoint/log step 76. """ if type(completed_step) is not int or completed_step < 0: raise ValueError("completed_step must be a non-negative integer") schedule = str(getattr(config, "generator_lr_schedule", "")).strip() if schedule != "cosine_window": raise ValueError( "generator_lr_schedule must be 'cosine_window'; got " f"{schedule!r}") base_lr = float(config.lr) min_lr = float(config.generator_lr_min) start_step = int(config.generator_lr_cosine_start_step) end_step = int(config.generator_lr_cosine_end_step) if ( not math.isfinite(base_lr) or not math.isfinite(min_lr) or base_lr <= 0.0 or min_lr <= 0.0 or min_lr > base_lr ): raise ValueError( "cosine-window generator learning rates must satisfy " "0 < generator_lr_min <= lr") if start_step < 0 or end_step <= start_step: raise ValueError( "generator LR cosine window requires " "0 <= start_step < end_step") if completed_step <= start_step: return base_lr if completed_step >= end_step: return min_lr progress = (completed_step - start_step) / (end_step - start_step) return min_lr + 0.5 * (base_lr - min_lr) * ( 1.0 + math.cos(math.pi * progress)) def _scale_replay_video_cotangents(cotangents, scale): """Scale replay gradients without changing the normalized MMD value. GradCache has already applied the configured MMD loss weight and the FSDP world-size compensation at this point. The explicit partial-gradient compensation is therefore applied exactly once, immediately before prepare/replay consumes these cotangents. """ values = tuple(cotangents) if float(scale) == 1.0: return values def scaled(cotangent): if isinstance(cotangent, Mapping): return { key: (value * float(scale)).contiguous() for key, value in cotangent.items() } return (cotangent * float(scale)).contiguous() return tuple(scaled(cotangent) for cotangent in values) def _atomic_torch_save(value, path): """Publish a checkpoint file only after ``torch.save`` fully succeeds.""" temporary_path = f"{path}.tmp" torch.save(value, temporary_path) os.replace(temporary_path, path) def _fresh_video_generator_grad_microbatches( *, local_samples, grad_local_samples, batch_size, ): """Return how many trailing microbatches build generator graphs. ``grad_local_samples=0`` is the compatibility path and replays every microbatch. A positive value keeps the complete forward feature batch for the MMD scalar but replays only the trailing whole microbatches. No inverse-probability rescaling is applied: earlier features are constants, exactly matching a partial derivative of the full MMD objective. """ values = { "local_samples": local_samples, "grad_local_samples": grad_local_samples, "batch_size": batch_size, } if any(type(value) is not int for value in values.values()): raise TypeError( "fresh-video generator-gradient geometry must use integers") if local_samples <= 0 or batch_size <= 0: raise ValueError("local_samples and batch_size must be positive") if local_samples % batch_size != 0: raise ValueError("local_samples must be divisible by batch_size") if grad_local_samples < 0 or grad_local_samples > local_samples: raise ValueError( "grad_local_samples must be zero or lie in [1, local_samples]") selected_samples = ( local_samples if grad_local_samples == 0 else grad_local_samples) if selected_samples % batch_size != 0: raise ValueError( "grad_local_samples must select complete physical microbatches") return selected_samples // batch_size def _transformer_block_types(wrapper, label): model = getattr(wrapper, "model", None) blocks = getattr(model, "blocks", None) if blocks is None or len(blocks) == 0: raise ValueError( f"{label} transformer FSDP strategy requires model.blocks") return {type(block) for block in blocks} def _strict_sha256(value, label): """Return a normalized SHA-256 string or reject ambiguous identity.""" normalized = str(value or "").strip().lower() if ( len(normalized) != 64 or any(character not in "0123456789abcdef" for character in normalized) ): raise ValueError(f"{label} must be a lowercase 64-character SHA-256") return normalized def _file_sha256(path): digest = hashlib.sha256() with open(path, "rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() class DistributionMatchingTrainer: def __init__(self, config): self.config = config self.step = 0 self.data_counter = 0 # Step 1: Initialize the distributed training environment (rank, seed, dtype, logging etc.) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True launch_distributed_job() global_rank = dist.get_rank() self.world_size = dist.get_world_size() self.uses_fresh_video_gradcache = ( str(getattr(config, "distribution_loss", "")) == "wan_kv_mmd" and str(getattr(config, "wan_mmd_estimator", "ema_rff")) == "nystrom_exact_repulsion" ) self.fresh_video_local_samples = int( getattr(config, "fresh_video_local_samples", config.batch_size)) if self.uses_fresh_video_gradcache: if self.fresh_video_local_samples <= 0: raise ValueError( "fresh_video_local_samples must be positive") if self.fresh_video_local_samples % int(config.batch_size) != 0: raise ValueError( "fresh_video_local_samples must be divisible by " f"batch_size; got {self.fresh_video_local_samples} and " f"{int(config.batch_size)}") self.fresh_video_microbatches = ( self.fresh_video_local_samples // int(config.batch_size)) else: self.fresh_video_microbatches = 1 expected_total_batch = ( int(config.batch_size) * self.world_size * self.fresh_video_microbatches ) configured_total_batch = int( getattr(config, "total_batch_size", expected_total_batch)) if configured_total_batch != expected_total_batch: raise ValueError( "total_batch_size must equal per-rank batch_size * world_size " "* fresh microbatches; " f"got {configured_total_batch} != {int(config.batch_size)} * " f"{self.world_size} * {self.fresh_video_microbatches} " f"({expected_total_batch})" ) self.fresh_video_total_samples = configured_total_batch configured_grad_local_samples = int(getattr( config, "fresh_video_generator_grad_local_samples", self.fresh_video_local_samples, )) self.fresh_video_generator_grad_microbatches = ( _fresh_video_generator_grad_microbatches( local_samples=self.fresh_video_local_samples, grad_local_samples=configured_grad_local_samples, batch_size=int(config.batch_size), ) if self.uses_fresh_video_gradcache else 1 ) self.fresh_video_generator_grad_local_samples = ( configured_grad_local_samples if self.uses_fresh_video_gradcache else int(config.batch_size) ) self.fresh_video_generator_grad_subsample_enabled = bool( self.uses_fresh_video_gradcache and self.fresh_video_generator_grad_local_samples < self.fresh_video_local_samples ) self.fresh_video_full_student_partial_gradient = bool(getattr( config, "fresh_video_full_student_partial_gradient", False)) if ( self.fresh_video_full_student_partial_gradient and not self.fresh_video_generator_grad_subsample_enabled ): raise ValueError( "fresh_video_full_student_partial_gradient requires a strict " "gradient suffix smaller than fresh_video_local_samples") self.fresh_video_generator_grad_global_samples = ( self.fresh_video_generator_grad_local_samples * self.world_size ) self.fresh_video_generator_grad_sample_scale = ( _validated_fresh_video_generator_grad_sample_scale( getattr( config, "fresh_video_generator_grad_sample_scale", 1.0, ), partial_gradient_enabled=( self.fresh_video_generator_grad_subsample_enabled), ) ) configured_teacher_samples = getattr( config, "expected_current_teacher_samples", None) if ( configured_teacher_samples is not None and int(configured_teacher_samples) != self.fresh_video_generator_grad_global_samples ): raise ValueError( "expected_current_teacher_samples must equal the global " "gradient suffix, got " f"{int(configured_teacher_samples)} != " f"{self.fresh_video_generator_grad_global_samples}") self.dtype = torch.bfloat16 if config.mixed_precision else torch.float32 self.device = torch.cuda.current_device() self.is_main_process = global_rank == 0 self.causal = config.causal self.disable_wandb = config.disable_wandb # use a random seed for the training if config.seed == 0: random_seed = torch.randint(0, 10000000, (1,), device=self.device) dist.broadcast(random_seed, src=0) config.seed = random_seed.item() set_seed(config.seed + global_rank) if self.is_main_process and not self.disable_wandb: init_wandb_from_config(config) self.output_path = config.logdir # Step 2: Initialize the model and optimizer self.model = DistributionMatchingModel(config, device=self.device) self.uses_critic = bool(getattr(self.model, "uses_critic", True)) if ( bool(getattr( self.model, "uses_fresh_video_gradcache", False)) != self.uses_fresh_video_gradcache ): raise RuntimeError( "trainer/model fresh-video GradCache configuration mismatch") fresh_cache_class = ( FeatureGradientReplay if bool(getattr(self.model, "dual_video_mmd_enabled", False)) else FreshVideoGradCache ) self.fresh_video_gradcache = ( fresh_cache_class( self.model.latent_mmd, target_local_samples=self.fresh_video_local_samples, gradient_local_samples=( self.fresh_video_generator_grad_local_samples), require_sample_id_alignment=( str(getattr(config, "current_teacher_source", "")) .strip().lower() in { "matched_manifest_video", "matched_dual_feature_bank", "matched_triple_feature_bank", "matched_dinov3_vjepa2_feature_bank", }), full_student_partial_gradient=( self.fresh_video_full_student_partial_gradient), loss_scale=float(getattr( config, "wan_kv_loss_weight", 1.0)), ) if self.uses_fresh_video_gradcache else None ) # Save pretrained model state_dicts to CPU self.fake_score_state_dict_cpu = ( self.model.fake_score.state_dict() if getattr(self.model, "fake_score", None) is not None else None ) generator_wrap_strategy = config.generator_fsdp_wrap_strategy self.model.generator = fsdp_wrap( self.model.generator, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=generator_wrap_strategy, transformer_module=( _transformer_block_types(self.model.generator, "generator") if generator_wrap_strategy == "transformer" else None ), ) if getattr(self.model, "real_score", None) is not None: real_wrap_strategy = config.real_score_fsdp_wrap_strategy self.model.real_score = fsdp_wrap( self.model.real_score, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=real_wrap_strategy, transformer_module=( _transformer_block_types(self.model.real_score, "real_score") if real_wrap_strategy == "transformer" else None ), ) if getattr(self.model, "fake_score", None) is not None: self.model.fake_score = fsdp_wrap( self.model.fake_score, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=config.fake_score_fsdp_wrap_strategy ) self.model.text_encoder = fsdp_wrap( self.model.text_encoder, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=config.text_encoder_fsdp_wrap_strategy, cpu_offload=getattr(config, "text_encoder_cpu_offload", False) ) if ( getattr(self.model, "uses_vae_for_loss", False) or ( not config.no_visualize and getattr(self.model, "uses_vae_for_visualization", True) ) or bool(getattr(config, "load_raw_video", False)) or int(getattr(config, "eval_interval", 0)) > 0 ): self.model.vae = self.model.vae.to( device=self.device, dtype=torch.bfloat16 if config.mixed_precision else torch.float32) self.generator_optimizer = torch.optim.AdamW( [param for param in self.model.generator.parameters() if param.requires_grad], lr=config.lr, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay ) if self.uses_critic: self.critic_optimizer = torch.optim.AdamW( [param for param in self.model.fake_score.parameters() if param.requires_grad], lr=config.lr_critic if hasattr(config, "lr_critic") else config.lr, betas=(config.beta1_critic, config.beta2_critic), weight_decay=config.weight_decay ) else: self.critic_optimizer = None # Step 3: Initialize the dataloader current_teacher_enabled = bool(getattr( self.model, "current_teacher_attraction_enabled", False)) configured_teacher_source = getattr( config, "current_teacher_source", None) requested_teacher_source = str( configured_teacher_source if configured_teacher_source is not None else "auto" ).strip().lower() self.current_teacher_source = ( "lmdb_latent" if current_teacher_enabled else "none" ) if requested_teacher_source == "auto" else requested_teacher_source if self.current_teacher_source not in { "none", "lmdb_latent", "openvid_video", "matched_manifest_video", "matched_dual_feature_bank", "matched_triple_feature_bank", "matched_dinov3_vjepa2_feature_bank", }: raise ValueError( "current_teacher_source must be auto, none, lmdb_latent, " "openvid_video, matched_manifest_video, or " "matched_dual_feature_bank, matched_triple_feature_bank, or " "matched_dinov3_vjepa2_feature_bank; got " f"{self.current_teacher_source!r}") if current_teacher_enabled == (self.current_teacher_source == "none"): raise ValueError( "current_teacher_source must name exactly one source iff " "exact current-teacher attraction is enabled") self.current_teacher_dataset = None self.matched_counter_sampling_enabled = False self.experiment_identity = None if self.config.distribution_loss == "wan_kv_mmd": if self.current_teacher_source in { "matched_dual_feature_bank", "matched_triple_feature_bank", "matched_dinov3_vjepa2_feature_bank", }: source = self.current_teacher_source triple_bank = source == "matched_triple_feature_bank" dino_vjepa_bank = ( source == "matched_dinov3_vjepa2_feature_bank") if triple_bank: from utils.triple_video_feature_bank import ( TripleVideoFeatureBank as FeatureBank, ) manifest_field = "triple_teacher_bank_manifest" content_field = "triple_teacher_bank_content_hash" expected_branch_names = { "dinov3", "vjepa2", "videomae_v2"} expected_feature_type = ( "triple_dinov3_vjepa2_videomae_v2_dense_latent") elif dino_vjepa_bank: from utils.dinov3_vjepa2_feature_bank import ( DINOv3VJEPA2FeatureBank as FeatureBank, ) manifest_field = "dinov3_vjepa2_teacher_bank_manifest" content_field = ( "dinov3_vjepa2_teacher_bank_content_hash") expected_branch_names = {"dinov3", "vjepa2"} expected_feature_type = ( "dual_dinov3_vjepa2_dense_latent") else: from utils.dual_video_feature_bank import ( DualVideoFeatureBank as FeatureBank, ) manifest_field = "dual_teacher_bank_manifest" content_field = "dual_teacher_bank_content_hash" expected_branch_names = {"vjepa2", "videomae_v2"} expected_feature_type = ( "dual_vjepa2_videomae_v2_dense_latent") expected_feature_types = ( { "triple_dinov3_vjepa2_videomae_v2_dense_latent", "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent", } if triple_bank else {expected_feature_type} ) if not bool(getattr(self.model, "dual_video_mmd_enabled", False)): raise ValueError( "matched feature bank requires multi-representation mode") if str(getattr( self.model, "feature_type", "")) not in expected_feature_types: raise ValueError( "teacher bank branch set does not match model mode") if not self.uses_fresh_video_gradcache: raise ValueError("matched feature bank requires exact GradCache") if int(config.batch_size) != 1: raise ValueError( "matched feature-bank training requires batch_size=1") manifest_sha_field = f"{manifest_field}_sha256" kwargs = { "expected_manifest_sha256": _strict_sha256( getattr(config, manifest_sha_field), manifest_sha_field), "expected_content_hash": _strict_sha256( getattr(config, content_field), content_field), "expected_source_manifest_sha256": _strict_sha256( config.matched_manifest_sha256, "matched_manifest_sha256"), "expected_sample_count": int(getattr( config, "matched_manifest_samples", 5000)), "expected_vjepa_model_sha256": _strict_sha256( config.vjepa2_model_sha256, "vjepa2_model_sha256"), "validate_all_shards": True, } if triple_bank or dino_vjepa_bank: kwargs.update({ "expected_source_csv_sha256": _strict_sha256( config.matched_source_csv_sha256, "matched_source_csv_sha256"), "expected_sample_identity_sha256": _strict_sha256( config.matched_sample_identity_sha256, "matched_sample_identity_sha256"), }) if not dino_vjepa_bank: kwargs.update({ "expected_videomae_revision": str( config.videomae_v2_model_revision), "expected_videomae_download_manifest_content_hash": ( _strict_sha256( config.videomae_v2_download_manifest_content_hash, "videomae_v2_download_manifest_content_hash")), }) if triple_bank or dino_vjepa_bank: kwargs.update({ "expected_dinov3_checkpoint_sha256": _strict_sha256( config.dinov3_checkpoint_sha256, "dinov3_checkpoint_sha256"), "expected_dinov3_source_commit": str( config.dinov3_source_commit), "expected_shard_size": 32, }) dataset = FeatureBank( str(getattr(config, manifest_field)), **kwargs) model_bank_hash = str(getattr( self.model, "triple_teacher_bank_content_hash" if triple_bank else "dual_teacher_bank_content_hash")) if dataset.content_hash != model_bank_hash: raise ValueError( "training bank identity differs from all references") self.current_teacher_dataset = dataset self.matched_counter_sampling_enabled = True branches = copy.deepcopy(self.model.dual_branch_registry) if set(branches) != expected_branch_names: raise ValueError( "model branch registry does not match teacher bank") self.experiment_identity = { "current_teacher_source": self.current_teacher_source, f"{manifest_field}_sha256": dataset.manifest_sha256, content_field: dataset.content_hash, "source_manifest_sha256": _strict_sha256( config.matched_manifest_sha256, "matched_manifest_sha256"), "sample_count": len(dataset), "branches": branches, "forward_global_samples": self.fresh_video_total_samples, "gradient_global_samples": ( self.fresh_video_generator_grad_global_samples), "teacher_global_samples": int( config.expected_current_teacher_samples), "full_student_partial_gradient": ( self.fresh_video_full_student_partial_gradient), "kss_denominator": self.fresh_video_total_samples ** 2, "fixed_attraction_denominator": ( self.fresh_video_total_samples), "teacher_cross_denominator": ( self.fresh_video_total_samples * int(config.expected_current_teacher_samples)), "teacher_self_denominator": ( int(config.expected_current_teacher_samples) ** 2), "replay_cotangent_compensation": ( self.fresh_video_generator_grad_sample_scale), "branch_normalization": ( "raw_div_stop_gradient_raw_plus_1e-4"), "branch_aggregation": ( "sum_15_to_1_to_2_then_divide_10" if triple_bank else "sum_10_to_1_without_divide_2"), } elif self.current_teacher_source == "matched_manifest_video": from utils.matched_video_dataset import ( MatchedManifestVideoDataset, ) from utils.matched_video_feature_cache import ( MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256, MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256, ) if not self.uses_fresh_video_gradcache: raise ValueError( "current_teacher_source=matched_manifest_video requires " "the fresh-video exact Nyström path") if int(config.batch_size) != 1: raise ValueError( "lazy matched-video teacher decoding requires " "batch_size=1") manifest_path = os.path.realpath(os.path.expanduser(str( config.matched_manifest_path))) expected_manifest_sha256 = _strict_sha256( config.matched_manifest_sha256, "matched_manifest_sha256", ) actual_manifest_sha256 = _file_sha256(manifest_path) if actual_manifest_sha256 != expected_manifest_sha256: raise ValueError( "frozen matched-video manifest SHA-256 mismatch: " f"{actual_manifest_sha256} != " f"{expected_manifest_sha256}") source_csv_sha256 = _strict_sha256( config.matched_source_csv_sha256, "matched_source_csv_sha256", ) expected_sample_identity_sha256 = _strict_sha256( config.matched_sample_identity_sha256, "matched_sample_identity_sha256", ) model_sha256 = _strict_sha256( config.vjepa2_model_sha256, "vjepa2_model_sha256") feature_contract_sha256 = _strict_sha256( config.matched_teacher_feature_contract_sha256, "matched_teacher_feature_contract_sha256", ) if ( feature_contract_sha256 != MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256 ): raise ValueError( "configured matched-video feature contract does not " "match runtime") decoder_contract_sha256 = _strict_sha256( config.matched_teacher_decoder_contract_sha256, "matched_teacher_decoder_contract_sha256", ) if ( decoder_contract_sha256 != MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256 ): raise ValueError( "configured matched-video decoder contract does not " "match runtime") sampling_algorithm = str( config.matched_sampling_algorithm) if sampling_algorithm != "sha256_seed_step_position_mod_n_v1": raise ValueError( "unsupported matched counter sampling algorithm: " f"{sampling_algorithm!r}") if not bool(config.matched_sampling_with_replacement): raise ValueError( "matched sampler must explicitly enable replacement") expected_samples = int(getattr( config, "matched_manifest_samples", 5000)) dataset = MatchedManifestVideoDataset( manifest_path=manifest_path, video_root=str(config.matched_video_root), cache_root=str(config.matched_teacher_cache_root), expected_model_sha256=model_sha256, expected_feature_contract_sha256=( feature_contract_sha256), expected_source_csv_sha256=source_csv_sha256, expected_count=expected_samples, ) actual_sample_identity_sha256 = str( dataset.manifest_metadata.get( "sample_identity_sha256", "")) if ( actual_sample_identity_sha256 != expected_sample_identity_sha256 ): raise ValueError( "matched-video sample identity SHA-256 mismatch: " f"{actual_sample_identity_sha256} != " f"{expected_sample_identity_sha256}") self.current_teacher_dataset = dataset self.matched_counter_sampling_enabled = True self.experiment_identity = { "current_teacher_source": "matched_manifest_video", "matched_manifest_sha256": expected_manifest_sha256, "matched_source_csv_sha256": source_csv_sha256, "matched_sample_identity_sha256": ( expected_sample_identity_sha256), "matched_manifest_samples": expected_samples, "vjepa2_reference_model_id": str( config.vjepa2_reference_model_id), "vjepa2_model_sha256": model_sha256, "matched_teacher_feature_contract_sha256": ( feature_contract_sha256), "matched_teacher_decoder_contract_sha256": ( decoder_contract_sha256), "student_frame_contract": "last64_frames_17_80", "live_teacher_frame_contract": "last64_frames_17_80", "fixed_reference_frame_contract": "central64_existing_pack", "counter_sampling_algorithm": ( "sha256_seed_step_position_mod_n_v1"), "counter_sampling_seed": int(config.seed), "forward_global_samples": ( self.fresh_video_total_samples), "gradient_global_samples": ( self.fresh_video_generator_grad_global_samples), "explicit_gradient_compensation": ( self.fresh_video_generator_grad_sample_scale), "partial_mmd_gradient_compensation": ( self.fresh_video_generator_grad_sample_scale), "teacher_global_samples": int( config.expected_current_teacher_samples), "partial_mmd_objective": ( "full_student_biased_mmd_suffix_partial_backward_v1" if self.fresh_video_full_student_partial_gradient else "rectangular_gradient_student_v1" ), "full_student_partial_gradient": ( self.fresh_video_full_student_partial_gradient), "partial_mmd_kss_denominator": ( self.fresh_video_total_samples ** 2 if self.fresh_video_full_student_partial_gradient else self.fresh_video_total_samples * self.fresh_video_generator_grad_global_samples ), "partial_mmd_fixed_attraction_denominator": ( self.fresh_video_total_samples if self.fresh_video_full_student_partial_gradient else self.fresh_video_generator_grad_global_samples ), "partial_mmd_live_cross_denominator": ( self.fresh_video_total_samples * int(config.expected_current_teacher_samples) if self.fresh_video_full_student_partial_gradient else self.fresh_video_generator_grad_global_samples * int(config.expected_current_teacher_samples) ), "partial_mmd_teacher_self_denominator": ( int(config.expected_current_teacher_samples) ** 2), "partial_replay_upstream_commit": ( "ef43f2d616ff31ba34904ae9743b0de4dda95afd"), "fixed_reference_file_sha256": _strict_sha256( config.vjepa2_nystrom_reference_sha256, "vjepa2_nystrom_reference_sha256"), "fixed_reference_hash": str(getattr( self.model, "vjepa2_nystrom_reference_hash", "")), "live_teacher_baseline_commit": str( config.live_teacher_baseline_commit), "partial_gradient_reference_commit": str( config.partial_gradient_reference_commit), "feature_protocol_identity": str( config.feature_protocol_identity), } elif self.current_teacher_source == "openvid_video": from utils.openvid_feature_cache import ( OPENVID_PYAV_DECODER_CONTRACT_SHA256, OPENVID_VJEPA2_L15_FEATURE_CONTRACT_SHA256, ) if not self.uses_fresh_video_gradcache: raise ValueError( "current_teacher_source=openvid_video requires the " "fresh-video exact Nyström path") if int(config.batch_size) != 1: raise ValueError( "OpenVid lazy teacher decoding currently requires " "batch_size=1") manifest_path = os.path.realpath(os.path.expanduser(str( config.openvid_manifest_path))) expected_manifest_sha256 = _strict_sha256( config.openvid_manifest_sha256, "openvid_manifest_sha256", ) actual_manifest_sha256 = _file_sha256(manifest_path) if actual_manifest_sha256 != expected_manifest_sha256: raise ValueError( "frozen OpenVid manifest SHA-256 mismatch: " f"{actual_manifest_sha256} != " f"{expected_manifest_sha256}") model_sha256 = _strict_sha256( config.vjepa2_model_sha256, "vjepa2_model_sha256") feature_contract_sha256 = _strict_sha256( config.openvid_teacher_feature_contract_sha256, "openvid_teacher_feature_contract_sha256", ) if ( feature_contract_sha256 != OPENVID_VJEPA2_L15_FEATURE_CONTRACT_SHA256 ): raise ValueError( "configured OpenVid feature contract does not match " "runtime") decoder_contract_sha256 = _strict_sha256( config.openvid_teacher_decoder_contract_sha256, "openvid_teacher_decoder_contract_sha256", ) if ( decoder_contract_sha256 != OPENVID_PYAV_DECODER_CONTRACT_SHA256 ): raise ValueError( "configured OpenVid decoder contract does not match " "runtime") dataset = OpenVidPromptVideoDataset( manifest_path=manifest_path, video_root=str(config.openvid_video_root), cache_root=str(config.openvid_teacher_cache_root), expected_model_sha256=model_sha256, expected_feature_contract_sha256=( feature_contract_sha256), ) expected_samples = int(getattr( config, "openvid_manifest_samples", 12800)) if len(dataset) != expected_samples: raise ValueError( "frozen OpenVid dataset size mismatch: " f"{len(dataset)} != {expected_samples}") self.current_teacher_dataset = dataset self.experiment_identity = { "current_teacher_source": "openvid_video", "openvid_manifest_sha256": expected_manifest_sha256, "vjepa2_reference_model_id": str( config.vjepa2_reference_model_id), "vjepa2_model_sha256": model_sha256, "openvid_teacher_feature_contract_sha256": ( feature_contract_sha256), "openvid_teacher_decoder_contract_sha256": ( decoder_contract_sha256), } else: load_matched_terminal_teacher = bool( self.uses_fresh_video_gradcache and self.current_teacher_source == "lmdb_latent") dataset = ODERegressionLMDBDataset( config.teacher_latent_data_path, max_pair=int(getattr( config, "teacher_latent_count", 1e8)), load_latents=( not self.uses_fresh_video_gradcache or load_matched_terminal_teacher ), load_terminal_latent_only=( load_matched_terminal_teacher), ) elif self.config.i2v: dataset = ShardingLMDBDataset(config.data_path, max_pair=int(1e8)) else: dataset = TextDataset(config.data_path) self.dataset = dataset if self.matched_counter_sampling_enabled: self.dataloader_len = 1 self.dataloader = None else: sampler = torch.utils.data.distributed.DistributedSampler( dataset, shuffle=True, drop_last=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=config.batch_size, sampler=sampler, num_workers=int(getattr(config, "num_workers", 4)), drop_last=True, collate_fn=( dataset.collate_fn if self.current_teacher_dataset is dataset else None), ) self.dataloader_len = max(1, len(dataloader)) self.dataloader = cycle(dataloader) if dist.get_rank() == 0: print("DATASET SIZE %d" % len(dataset)) ############################################################################################################## # 6. Record trainable parameters. Generator EMA is initialized only # after loading the generator checkpoint, below. rename_param = ( lambda name: name.replace("_fsdp_wrapped_module.", "") .replace("_checkpoint_wrapped_module.", "") .replace("_orig_mod.", "") ) self.name_to_trainable_params = {} for n, p in self.model.generator.named_parameters(): if not p.requires_grad: continue renamed_n = rename_param(n) self.name_to_trainable_params[renamed_n] = p ema_weight = float(getattr(config, "ema_weight", 0.0) or 0.0) self.generator_ema = None resume_ema_state = None ############################################################################################################## # 7. Load model/distribution state before constructing an EMA shadow. resume_checkpoint = getattr(config, "resume_checkpoint", None) self.resume_data_counter_override = ( _validated_resume_data_counter_override( getattr(config, "resume_data_counter_override", None), resume_checkpoint=resume_checkpoint, ) ) if resume_checkpoint: # Restore generator weights, optional generator EMA, distribution # history and step. Optimizer/RNG/dataloader state is restored later # when save_full_state is enabled and optim.pt is present. print(f"Resuming training state from {resume_checkpoint}") ckpt = torch.load(resume_checkpoint, map_location="cpu", weights_only=False) if self.experiment_identity is not None: saved_identity = ( ckpt.get("experiment_identity") if isinstance(ckpt, dict) else None ) saved_identity = _canonicalize_experiment_identity_paths( saved_identity) runtime_identity = _canonicalize_experiment_identity_paths( self.experiment_identity) if saved_identity != runtime_identity: raise ValueError( "resume checkpoint experiment identity does not match " "the frozen teacher/V-JEPA2 experiment identity: " f"{saved_identity!r} != " f"{runtime_identity!r}") gen_sd = ckpt["generator"] if (isinstance(ckpt, dict) and "generator" in ckpt) else ckpt self.model.generator.load_state_dict(gen_sd, strict=False) if isinstance(ckpt, dict): resume_ema_state = ckpt.get("generator_ema") if ( isinstance(ckpt, dict) and "fd" in ckpt and hasattr(self.model, "load_fd_state_dict") ): self.model.load_fd_state_dict(ckpt["fd"]) print(" restored distribution state") elif getattr( self.model, "requires_distribution_state_on_resume", False ): raise ValueError( "resume checkpoint is missing the required " "distribution history under key 'fd'") if isinstance(ckpt, dict) and "step" in ckpt: self.step = int(ckpt["step"]) else: import re m = re.search(r"checkpoint_model_(\d+)", str(resume_checkpoint)) self.step = int(m.group(1)) if m else 0 if isinstance(ckpt, dict) and "data_counter" in ckpt: self.data_counter = int(ckpt["data_counter"]) print(f" resumed at step {self.step}") elif getattr(config, "generator_ckpt", False): print(f"Loading pretrained generator from {config.generator_ckpt}") state_dict = torch.load(config.generator_ckpt, map_location="cpu") fd_state_dict = state_dict.get("fd") if isinstance(state_dict, dict) else None if isinstance(state_dict, dict): if "generator" in state_dict: state_dict = state_dict["generator"] elif "model" in state_dict: state_dict = state_dict["model"] self.model.generator.load_state_dict( state_dict, strict=False ) if fd_state_dict is not None and hasattr(self.model, "load_fd_state_dict"): is_compatible_initialization_state = ( self.model.can_load_initialization_fd_state(fd_state_dict) if hasattr( self.model, "can_load_initialization_fd_state") else not getattr( self.model, "requires_distribution_state_on_resume", False, ) ) if is_compatible_initialization_state: # Known dense-MMD states are always passed through the # model's strict metadata/cross-feature compatibility # checks. Unknown distribution-state payloads in an initialization-only # checkpoint remain intentionally ignorable. self.model.load_fd_state_dict(fd_state_dict) elif self.is_main_process: print( " initialization checkpoint has unrelated or " "incompatible FD state; using generator weights only") ############################################################################################################## ema_start_step = int(getattr(config, "ema_start_step", 0)) if ema_weight > 0.0 and self.step >= ema_start_step: print(f"Setting up EMA with weight {ema_weight}") # This now snapshots the loaded generator, never the stale # pre-checkpoint parameters. self.generator_ema = EMA_FSDP( self.model.generator, decay=ema_weight) if resume_ema_state is not None: self.generator_ema.load_state_dict(resume_ema_state) print(" restored EMA shadow") elif resume_checkpoint and self.step > ema_start_step: raise ValueError( "resume checkpoint is past ema_start_step but has no " "generator_ema state") self.max_grad_norm_generator = getattr(config, "max_grad_norm_generator", 10.0) self.max_grad_norm_critic = getattr(config, "max_grad_norm_critic", 10.0) self.previous_time = None # ---- precise-resume + periodic eval bookkeeping ---- # ``save_full_state`` additionally writes the optimizer state, per-rank RNG # states and the dataloader counter so a resume continues bit-for-bit (up # to multi-worker dataloader ordering). Off by default so other trainers # (DMD etc.) keep their lightweight checkpoints. self.save_full_state = bool(getattr(config, "save_full_state", False)) self.save_full_state_final_only = bool( getattr(config, "save_full_state_final_only", False)) self.checkpoint_keep_last = int( getattr(config, "checkpoint_keep_last", 0)) if self.checkpoint_keep_last < 0: raise ValueError("checkpoint_keep_last must be non-negative") self.eval_interval = int(getattr(config, "eval_interval", 0)) self.eval_seed = int(getattr(config, "eval_seed", 12345)) self.eval_num_prompts = int(getattr(config, "eval_num_prompts", 4)) self.eval_prompts = self._load_eval_prompts() if self.eval_interval > 0 else [] self._last_saved_step = None # If we resumed and the optimizer/RNG/data state was restored, this is set. self._restore_full_training_state(resume_checkpoint) self._active_generator_lr = self._set_generator_learning_rate( self.step + 1) if self.is_main_process and str(getattr( self.config, "generator_lr_schedule", "")).strip(): print( "[generator-lr] restored_optimizer_override " f"next_completed_step={self.step + 1} " f"lr={self._active_generator_lr:.12g}", flush=True, ) def _all_ranks_true(self, local_value): """Create an all-rank decision before a rank can branch or raise.""" flag = torch.tensor( int(bool(local_value)), device=self.device, dtype=torch.int32) if dist.is_initialized() and dist.get_world_size() > 1: dist.all_reduce(flag, op=dist.ReduceOp.MIN) return bool(flag.item()) def _set_generator_learning_rate(self, completed_step): """Apply an opt-in schedule after optimizer restore and before update.""" schedule = str(getattr( self.config, "generator_lr_schedule", "")).strip() if not schedule: # Preserve historical resume behavior for every configuration that # has not explicitly opted into a generator LR schedule. return float(self.generator_optimizer.param_groups[0]["lr"]) lr = _generator_lr_for_completed_step( self.config, completed_step) for param_group in self.generator_optimizer.param_groups: param_group["lr"] = lr return lr def save(self): if ( self.fresh_video_gradcache is not None and ( self.fresh_video_gradcache.local_samples or self.fresh_video_gradcache.finalized ) ): raise RuntimeError( "refusing to checkpoint a partial fresh-video optimizer " "update") if ( hasattr(self.model, "distribution_state_pending") and self.model.distribution_state_pending() ): raise RuntimeError( "refusing to checkpoint pending distribution state") # max_steps can coincide with log_iters; avoid gathering and rewriting # the same full FSDP checkpoint twice. if self._last_saved_step == self.step: return max_steps = int(getattr(self.config, "max_steps", 0)) write_full_state = ( self.save_full_state and ( not self.save_full_state_final_only or (max_steps > 0 and self.step >= max_steps) ) ) required_free_gib = float( getattr( self.config, ( "checkpoint_full_min_free_gib" if write_full_state else "checkpoint_model_min_free_gib" ), 0.0, ) ) disk_ok = True disk_error = None if self.is_main_process and required_free_gib > 0: try: free_gib = ( shutil.disk_usage(self.output_path).free / (1024 ** 3)) if free_gib < required_free_gib: disk_ok = False disk_error = ( f"checkpoint at step {self.step} needs at least " f"{required_free_gib:.1f} GiB free, only " f"{free_gib:.1f} GiB is available" ) except Exception as exc: disk_ok = False disk_error = f"checkpoint disk-space check failed: {exc}" disk_flag = torch.tensor( int(disk_ok), device=self.device, dtype=torch.int32) if dist.is_initialized(): dist.broadcast(disk_flag, src=0) if not bool(disk_flag.item()): raise RuntimeError( disk_error or "rank 0 checkpoint disk-space check failed") print("Start gathering distributed model states...") generator_state_dict = fsdp_state_dict( self.model.generator) critic_state_dict = ( fsdp_state_dict(self.model.fake_score) if self.uses_critic else None ) if self.generator_ema is not None: state_dict = { "generator": generator_state_dict, "generator_ema": self.generator_ema.state_dict(), } else: state_dict = { "generator": generator_state_dict, } if self.uses_critic: state_dict["critic"] = critic_state_dict if hasattr(self.model, "fd_state_dict"): state_dict["fd"] = self.model.fd_state_dict() # record the step so a resume continues numbering / EMA timing. state_dict["step"] = self.step state_dict["data_counter"] = self.data_counter if self.experiment_identity is not None: state_dict["experiment_identity"] = ( _canonicalize_experiment_identity_paths( self.experiment_identity)) # ---- full training state for precise resume (optimizer + RNG + data) ---- gen_optim_full = None critic_optim_full = None if write_full_state: # collective: every rank participates; full state gathered on rank 0. gen_optim_full = fsdp_optim_state_dict(self.model.generator, self.generator_optimizer) if self.uses_critic: critic_optim_full = fsdp_optim_state_dict(self.model.fake_score, self.critic_optimizer) ckpt_dir = os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}") rank0_write_error = None if self.is_main_process: try: os.makedirs(ckpt_dir, exist_ok=True) model_path = os.path.join(ckpt_dir, "model.pt") _atomic_torch_save(state_dict, model_path) print("Model saved to", model_path) if write_full_state: optim_path = os.path.join(ckpt_dir, "optim.pt") _atomic_torch_save( { "generator_optimizer": gen_optim_full, "critic_optimizer": critic_optim_full, "step": self.step, "data_counter": self.data_counter, "world_size": self.world_size, }, optim_path, ) print( "Optimizer/training state saved to", optim_path, ) except Exception as exc: rank0_write_error = exc write_ok = torch.tensor( int(rank0_write_error is None), device=self.device, dtype=torch.int32, ) if dist.is_initialized(): dist.broadcast(write_ok, src=0) if not bool(write_ok.item()): if rank0_write_error is not None: raise RuntimeError( f"checkpoint write failed at step {self.step}" ) from rank0_write_error raise RuntimeError( f"rank 0 checkpoint write failed at step {self.step}") # Per-rank RNG state (each rank writes its own) for reproducible streams. rng_write_ok = True rng_write_error = None if write_full_state: try: os.makedirs(ckpt_dir, exist_ok=True) rng_state = { "torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state(self.device), "numpy": np.random.get_state(), "python": random.getstate(), "data_counter": self.data_counter, } _atomic_torch_save( rng_state, os.path.join( ckpt_dir, f"rng_rank{dist.get_rank():02d}.pt"), ) except Exception as exc: rng_write_ok = False rng_write_error = exc rng_ok = torch.tensor( int(rng_write_ok), device=self.device, dtype=torch.int32) if dist.is_initialized(): dist.all_reduce(rng_ok, op=dist.ReduceOp.MIN) if not bool(rng_ok.item()): if rng_write_error is not None: raise RuntimeError( f"rank {dist.get_rank()} RNG checkpoint write failed at " f"step {self.step}" ) from rng_write_error raise RuntimeError( f"an RNG checkpoint write failed at step {self.step}") if dist.is_initialized(): dist.barrier() retention_error = None if self.is_main_process and self.checkpoint_keep_last > 0: try: checkpoints = [] for name in os.listdir(self.output_path): prefix = "checkpoint_model_" if not name.startswith(prefix): continue try: checkpoint_step = int(name[len(prefix):]) except ValueError: continue path = os.path.join(self.output_path, name) if not os.path.isdir(path): continue if not os.path.isfile(os.path.join(path, "model.pt")): shutil.rmtree(path) print("Removed incomplete checkpoint", path) continue checkpoints.append((checkpoint_step, path)) checkpoints.sort() for _, path in checkpoints[:-self.checkpoint_keep_last]: shutil.rmtree(path) print("Removed old checkpoint", path) except Exception as exc: retention_error = exc retention_ok = torch.tensor( int(retention_error is None), device=self.device, dtype=torch.int32, ) if dist.is_initialized(): dist.broadcast(retention_ok, src=0) if not bool(retention_ok.item()): if retention_error is not None: raise RuntimeError( "checkpoint retention failed after publishing a valid " f"checkpoint at step {self.step}" ) from retention_error raise RuntimeError( "rank 0 checkpoint retention failed after publishing a " f"valid checkpoint at step {self.step}") if dist.is_initialized(): dist.barrier() self._last_saved_step = self.step def _restore_full_training_state(self, resume_checkpoint): """Restore optimizer + per-rank RNG + dataloader position for a precise resume. Called at the end of ``__init__`` (after the optimizer and generator weights exist). Silently no-ops if the resume checkpoint has no ``optim.pt`` (e.g. an old lightweight checkpoint), so older runs still resume their weights/EMA/queue/step as before. """ if not resume_checkpoint: return if not self.save_full_state: # Lightweight checkpoints still record data_counter. They do not # promise bit-exact optimizer/RNG continuation, but they must not # restart the prompt stream at its first batch. self._apply_lightweight_resume_data_counter_override() self._fast_forward_dataloader(self.data_counter) if self.is_main_process: print( " lightweight resume: restored generator/distribution " "state and data position; optimizer and RNG start fresh", flush=True, ) return ckpt_dir = os.path.dirname(resume_checkpoint) optim_path = os.path.join(ckpt_dir, "optim.pt") if os.path.exists(optim_path): if getattr(self, "resume_data_counter_override", None) is not None: raise ValueError( "resume_data_counter_override is only valid for a " "lightweight resume; precise optimizer state exists at " f"{optim_path}") # rank0 holds the full optimizer state; FSDP broadcasts/shards it. full = None optim_load_error = None if self.is_main_process: try: full = torch.load( optim_path, map_location="cpu", weights_only=False, ) except Exception as exc: optim_load_error = exc optim_load_ok = torch.tensor( int(optim_load_error is None), device=self.device, dtype=torch.int32, ) if dist.is_initialized(): dist.broadcast(optim_load_ok, src=0) if not bool(optim_load_ok.item()): if optim_load_error is not None: raise RuntimeError( f"failed to load optimizer checkpoint {optim_path}" ) from optim_load_error raise RuntimeError( f"rank 0 failed to load optimizer checkpoint {optim_path}") saved_world_size = torch.tensor( [ int(full.get("world_size", self.world_size)) if full is not None else 0 ], device=self.device, dtype=torch.long, ) dist.broadcast(saved_world_size, src=0) if ( int(saved_world_size.item()) != self.world_size and not bool(getattr( self.config, "allow_resume_world_size_change", False)) ): raise ValueError( "precise resume world-size mismatch: checkpoint used " f"{int(saved_world_size.item())}, current run uses " f"{self.world_size}; set allow_resume_world_size_change " "only if non-bit-exact data/RNG continuation is acceptable") gen_osd = full["generator_optimizer"] if full is not None else None fsdp_load_optim_state_dict(self.model.generator, self.generator_optimizer, gen_osd) if self.uses_critic: critic_osd = full["critic_optimizer"] if full is not None else None if critic_osd is not None: fsdp_load_optim_state_dict(self.model.fake_score, self.critic_optimizer, critic_osd) # broadcast the data counter from rank0 dc = torch.tensor( [int(full["data_counter"]) if full is not None else 0], device=self.device, dtype=torch.long, ) if dist.is_initialized(): dist.broadcast(dc, src=0) self.data_counter = int(dc.item()) if self.is_main_process: print(f" restored optimizer state; data_counter={self.data_counter}") else: self._apply_lightweight_resume_data_counter_override() if self.is_main_process: print( " lightweight resume: restored generator/MMD history " "and data position; optimizer and RNG start fresh", flush=True, ) # Read RNG now, but restore it only after DataLoader fast-forwarding: # creating/advancing workers consumes the main process RNG. rng_path = os.path.join(ckpt_dir, f"rng_rank{dist.get_rank():02d}.pt") rng = None if os.path.exists(rng_path): rng = torch.load(rng_path, map_location="cpu", weights_only=False) # fast-forward the prompt stream to its position self._fast_forward_dataloader(self.data_counter) if rng is not None: torch.set_rng_state(rng["torch"]) torch.cuda.set_rng_state(rng["cuda"], self.device) np.random.set_state(rng["numpy"]) random.setstate(rng["python"]) if self.is_main_process: print(" restored per-rank RNG states") def _apply_lightweight_resume_data_counter_override(self): override = getattr(self, "resume_data_counter_override", None) if override is None: return checkpoint_counter = self.data_counter self.data_counter = override if self.is_main_process: print( " lightweight resume data cursor: explicitly overriding " f"checkpoint data_counter {checkpoint_counter} -> {override} " "from resume_data_counter_override", flush=True, ) def _matched_global_indices_for_step(self): """Counter-based with-replacement draw shared identically by all ranks.""" if not getattr(self, "matched_counter_sampling_enabled", False): raise RuntimeError("matched counter sampler is not enabled") population = len(self.dataset) if population <= 0: raise RuntimeError("matched counter sampler has an empty population") seed = int(self.config.seed) step = int(self.step) indices = [] for position in range(self.fresh_video_total_samples): payload = ( f"sf-matched-counter-v1\0{seed}\0{step}\0{position}" .encode("ascii") ) value = int.from_bytes( hashlib.sha256(payload).digest()[:8], "little") indices.append(value % population) return tuple(indices) def _matched_batch_for_microbatch(self, microbatch_index): if not getattr(self, "matched_counter_sampling_enabled", False): raise RuntimeError("matched counter sampler is not enabled") batch_size = int(self.config.batch_size) local_offset = ( dist.get_rank() * self.fresh_video_local_samples + int(microbatch_index) * batch_size ) indices = self._matched_global_indices_for_step() selected = indices[local_offset:local_offset + batch_size] if len(selected) != batch_size: raise RuntimeError("matched counter sampler produced a short slice") items = [self.dataset[index] for index in selected] self.data_counter += 1 return self.dataset.collate_fn(items) def _fast_forward_dataloader(self, n_consumed): if getattr(self, "matched_counter_sampling_enabled", False): return if n_consumed <= 0: return if hasattr(self.dataloader, "seek"): skip = self.dataloader.seek(n_consumed) else: skip = n_consumed % self.dataloader_len for _ in range(skip): next(self.dataloader) if self.is_main_process and skip: print( f" restored dataloader epoch/offset " f"({n_consumed // self.dataloader_len}, {skip})") def _next_batch(self): if getattr(self, "matched_counter_sampling_enabled", False): raise RuntimeError( "matched-manifest training must use the counter sampler") self.data_counter += 1 return next(self.dataloader) def _load_eval_prompts(self): path = getattr(self.config, "eval_prompt_path", "prompts/test_8.txt") try: with open(path, encoding="utf-8") as f: lines = [l.strip() for l in f if l.strip()] except FileNotFoundError: if self.is_main_process: print(f"[eval] prompt file {path} not found; eval disabled") return [] return lines[: self.eval_num_prompts] @torch.no_grad() def _fd_eval_inference(self): """Generate fixed-prompt videos with the live generator and log to wandb. Runs the model's own few-step (here single-step, no-CFG) generation under a fixed seed so successive evals are directly comparable. All ranks participate (FSDP forward is collective); only rank 0 writes/logs. """ if not self.eval_prompts: return # Frozen encoders stay in eval mode; generator gradients are enabled only on replay. self.model.eval() # fix the RNG for reproducible eval, then restore so training is unperturbed cpu_rng = torch.get_rng_state() cuda_rng = torch.cuda.get_rng_state(self.device) torch.manual_seed(self.eval_seed) prompts = self.eval_prompts # Evaluation is deliberately serialized at one video per rank. The # training forward is also microbatch 1; evaluating all fixed prompts # together would multiply the peak generator/VAE memory before the # first optimizer update. rank0_pixels = [] try: for prompt in prompts: conditional_dict = self.model.text_encoder( text_prompts=[prompt]) unconditional_dict = self.model.text_encoder( text_prompts=[self.config.negative_prompt]) shape = [1] + list( self.config.image_or_video_shape[1:]) latent = self.model._run_few_step_generator( image_or_video_shape=shape, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, ) pixels = self.model.vae.decode_to_pixel( latent.to(self.dtype)) pixels = ((pixels + 1.0) * 0.5).clamp(0.0, 1.0) if self.is_main_process: rank0_pixels.append(pixels[0].float().cpu()) del latent, pixels, conditional_dict, unconditional_dict finally: torch.set_rng_state(cpu_rng) torch.cuda.set_rng_state(cuda_rng, self.device) if self.is_main_process: eval_dir = os.path.join(self.output_path, "eval", f"step_{self.step:06d}") os.makedirs(eval_dir, exist_ok=True) log_videos = {} for i in range(len(prompts)): vid = ( rank0_pixels[i].permute(0, 2, 3, 1).numpy() * 255.0 ).astype("uint8") # [T,H,W,C] out_path = os.path.join(eval_dir, f"sample_{i}.mp4") wrote = False try: iio.imwrite( out_path, vid, fps=16, codec="libx264", plugin="FFMPEG", ) wrote = True except Exception as exc: # pragma: no cover - codec/runtime issues print(f"[eval] failed to write {out_path}: {exc}") # Log the encoded mp4 by PATH so wandb does not require moviepy. if not self.disable_wandb and wrote: try: log_videos[f"eval/sample_{i}"] = wandb.Video(out_path, format="mp4") except Exception as exc: # pragma: no cover print(f"[eval] wandb.Video failed for {out_path}: {exc}") if not self.disable_wandb and log_videos: wandb.log(log_videos, step=self.step) print(f"[eval] step {self.step}: wrote {len(prompts)} videos to {eval_dir}", flush=True) if dist.is_initialized(): dist.barrier() @staticmethod def _batch_string_list(value, *, key, batch_size): if isinstance(value, str): values = [value] elif isinstance(value, (list, tuple)): values = list(value) else: raise TypeError(f"{key} must collate to a list of strings") if ( len(values) != batch_size or any(not isinstance(item, str) or not item for item in values) ): raise ValueError( f"{key} must contain one non-empty string per prompt") return values def _attach_matched_teacher_payload(self, batch): """Load cache/decode teacher data only for a selected replay suffix.""" if self.current_teacher_source != "matched_manifest_video": return if self.current_teacher_dataset is None: raise RuntimeError("matched teacher dataset is unavailable") batch_size = len(batch.get("prompts", ())) sample_ids = self._batch_string_list( batch.get("sample_ids"), key="sample_ids", batch_size=batch_size) cache_keys = self._batch_string_list( batch.get("teacher_cache_keys"), key="teacher_cache_keys", batch_size=batch_size) local_error = None payloads = None try: payloads = [ self.current_teacher_dataset.load_teacher_payload(sample_id) for sample_id in sample_ids ] for sample_id, cache_key, payload in zip( sample_ids, cache_keys, payloads ): if ( payload.get("sample_id") != sample_id or payload.get("teacher_cache_key") != cache_key ): raise ValueError( "lazy teacher payload identity does not match the " "selected student") hit_values = [ bool(payload["current_teacher_cache_hit"]) for payload in payloads ] if not (all(hit_values) or not any(hit_values)): raise ValueError( "a metadata microbatch cannot mix cached features and " "decoded videos") batch["current_teacher_cache_hit"] = torch.tensor( hit_values, dtype=torch.bool) batch["current_teacher_decode_seconds"] = sum( float(payload["current_teacher_decode_seconds"]) for payload in payloads ) if all(hit_values): batch["current_teacher_features"] = torch.stack([ payload["current_teacher_features"] for payload in payloads ]) else: batch["current_teacher_video01"] = torch.stack([ payload["current_teacher_video01"] for payload in payloads ]) except BaseException as exc: local_error = exc if not self._all_ranks_true(local_error is None): message = ( f": {type(local_error).__name__}: {local_error}" if local_error is not None else " on another rank") raise RuntimeError( "matched live-teacher lazy I/O failed" + message ) from local_error def _prepare_current_teacher_inputs(self, batch, batch_size): """Validate and move exactly one configured teacher payload.""" if not bool(getattr( self.model, "current_teacher_attraction_enabled", False )): return {}, None, None local_error = None forward_inputs = None cache_records = None cache_log = None try: if self.current_teacher_source in { "matched_dual_feature_bank", "matched_triple_feature_bank", "matched_dinov3_vjepa2_feature_bank", }: if self.current_teacher_dataset is None: raise RuntimeError("matched teacher bank is unavailable") sample_ids = self._batch_string_list( batch.get("sample_ids"), key="sample_ids", batch_size=batch_size) features = self.current_teacher_dataset.load_rows(sample_ids) registry = getattr(self.model, "dual_branch_registry", None) if isinstance(registry, Mapping): expected_shapes = { name: tuple(branch.get( "raw_sample_shape", branch["sample_shape"])) for name, branch in registry.items() } else: expected_shapes = getattr( self.model.latent_mmd, "sample_shapes", {}) if set(features) != set(expected_shapes): raise ValueError("teacher bank/model feature branches differ") for name, tensor in features.items(): expected = (batch_size, *expected_shapes[name]) if tensor.dtype != torch.float16 or tuple(tensor.shape) != expected: raise ValueError( f"teacher bank {name} row shape/dtype mismatch: " f"{tuple(tensor.shape)}/{tensor.dtype} != " f"{expected}/float16") features[name] = tensor.to( device=self.device, dtype=torch.float32) forward_inputs = {"current_teacher_features": features} cache_records = () cache_log = { "hits": batch_size, "misses": 0, "decode_seconds": 0.0, "bank_rows_read": batch_size, } elif self.current_teacher_source == "lmdb_latent": latent = batch.get("ode_latent") if ( not isinstance(latent, torch.Tensor) or latent.ndim != 5 or int(latent.shape[0]) != batch_size or not latent.is_floating_point() ): raise ValueError( "exact LMDB current-teacher MMD requires one terminal " "latent per prompt with shape " "[B,frames,channels,height,width]") latent = latent.to(device=self.device, dtype=self.dtype) if not bool(torch.isfinite(latent).all().item()): raise FloatingPointError( "matched current teacher latent contains non-finite " "values") forward_inputs = {"current_teacher_latent": latent} elif self.current_teacher_source in { "openvid_video", "matched_manifest_video" }: features = batch.get("current_teacher_features") video01 = batch.get("current_teacher_video01") if (features is None) == (video01 is None): raise ValueError( "video teacher batch must provide exactly one of cached " "current_teacher_features or current_teacher_video01") cache_hit = batch.get("current_teacher_cache_hit") if ( not isinstance(cache_hit, torch.Tensor) or cache_hit.dtype != torch.bool or tuple(cache_hit.shape) != (batch_size,) ): raise ValueError( "current_teacher_cache_hit must be BoolTensor[B]") sample_ids = self._batch_string_list( batch.get("sample_ids"), key="sample_ids", batch_size=batch_size) cache_keys = self._batch_string_list( batch.get("teacher_cache_keys"), key="teacher_cache_keys", batch_size=batch_size) decode_value = batch.get("current_teacher_decode_seconds", 0.0) if isinstance(decode_value, torch.Tensor): decode_seconds = float( decode_value.detach().float().sum().item()) else: decode_seconds = float(decode_value) if not np.isfinite(decode_seconds) or decode_seconds < 0.0: raise ValueError( "current_teacher_decode_seconds must be finite and " "non-negative") if features is not None: expected_shape = ( batch_size, int(self.model.dense_num_tokens), int(self.model.dense_feature_dim), ) if ( not isinstance(features, torch.Tensor) or not features.is_floating_point() or tuple(features.shape) != expected_shape or not bool(cache_hit.all().item()) ): raise ValueError( "cached video teacher features must be floating " f"{expected_shape} with every cache-hit flag set") features = features.to( device=self.device, dtype=torch.float32) if not bool(torch.isfinite(features).all().item()): raise FloatingPointError( "cached video teacher features are non-finite") forward_inputs = { "current_teacher_features": features} cache_hits = batch_size cache_records = () else: expected_prefix = ( batch_size, int(self.model.dense_num_frames), 3, ) if ( not isinstance(video01, torch.Tensor) or not video01.is_floating_point() or video01.ndim != 5 or tuple(video01.shape[:3]) != expected_prefix or bool(cache_hit.any().item()) ): raise ValueError( "uncached video teacher video must be floating " "[B,64,3,H,W] with every cache-hit flag clear") if ( not bool(torch.isfinite(video01).all().item()) or float(video01.min().item()) < 0.0 or float(video01.max().item()) > 1.0 ): raise ValueError( "current_teacher_video01 must contain finite " "values in [0,1]") video01 = video01.to( device=self.device, dtype=torch.float32) forward_inputs = { "current_teacher_video01": video01} cache_hits = 0 cache_records = tuple(zip(cache_keys, sample_ids)) cache_log = { "hits": cache_hits, "misses": batch_size - cache_hits, "decode_seconds": decode_seconds, } else: raise RuntimeError( "enabled exact current teacher has no configured source") except BaseException as exc: local_error = exc if not self._all_ranks_true(local_error is None): message = ( f": {type(local_error).__name__}: {local_error}" if local_error is not None else " on another rank") raise RuntimeError( "current teacher batch validation failed" + message ) from local_error return forward_inputs, cache_records, cache_log def _write_current_teacher_cache(self, records, features): if not records: return if self.current_teacher_dataset is None: raise RuntimeError("live-teacher feature cache writer is unavailable") cpu_features = features.detach().float().cpu() if int(cpu_features.shape[0]) != len(records): raise ValueError( "teacher feature count does not match cache-miss records") for (cache_key, sample_id), feature in zip(records, cpu_features): self.current_teacher_dataset.write_feature_cache( cache_key, feature, sample_id=sample_id) def _encode_fresh_video_text_conditioning( self, text_prompts, batch_size, ): """Encode prompts and keep the joint-feature mask out-of-band.""" if len(text_prompts) != int(batch_size): raise ValueError( "text prompt count must match the fresh microbatch size") joint_enabled = bool(getattr( self.model, "triple_t5_joint_enabled", False)) conditional_dict = self.model.text_encoder( text_prompts=text_prompts, **( {"return_attention_mask": True} if joint_enabled else {} ), ) t5_attention_mask = conditional_dict.pop( "prompt_attention_mask", None) if joint_enabled: if not isinstance(t5_attention_mask, torch.Tensor): raise RuntimeError( "triple joint training requires the text encoder to " "return prompt_attention_mask") prompt_embeds = conditional_dict.get("prompt_embeds") if ( not isinstance(prompt_embeds, torch.Tensor) or t5_attention_mask.shape != prompt_embeds.shape[:2] ): raise ValueError( "prompt_attention_mask must match prompt_embeds [B,S]") if t5_attention_mask.device != prompt_embeds.device: raise ValueError( "prompt_attention_mask and prompt_embeds must share " "a device") elif t5_attention_mask is not None: raise RuntimeError( "pure-video text encoding unexpectedly returned an " "attention mask") cached_unconditional = getattr( self, "unconditional_dict", None) cached_batch = ( next(iter(cached_unconditional.values())).shape[0] if cached_unconditional else None ) if cached_unconditional is None or cached_batch != batch_size: unconditional_dict = self.model.text_encoder( text_prompts=[ self.config.negative_prompt] * batch_size) unconditional_dict = { key: value.detach() for key, value in unconditional_dict.items() } self.unconditional_dict = unconditional_dict else: unconditional_dict = cached_unconditional return conditional_dict, unconditional_dict, t5_attention_mask def _collect_rollout_microbatch( self, batch, *, gradient_selected=True, vjepa2_window_sample_keys=None, ): """Collect generated features and replay inputs without a generator graph. Reference minibatch rows are attached only to selected differentiation samples; all generated features still enter the distribution loss.""" text_prompts = batch["prompts"] batch_size = len(text_prompts) sample_ids = ( self._batch_string_list( batch.get("sample_ids"), key="sample_ids", batch_size=batch_size) if batch.get("sample_ids") is not None else None ) if ( gradient_selected and self.current_teacher_source == "matched_manifest_video" ): self._attach_matched_teacher_payload(batch) if gradient_selected: teacher_forward_inputs, cache_records, teacher_cache_log = ( self._prepare_current_teacher_inputs(batch, batch_size)) else: teacher_forward_inputs, cache_records, teacher_cache_log = ( {}, None, None) # Transfer ownership out of the DataLoader batch before model entry. # Each tensor is popped from teacher_forward_inputs as the call binds # it, so the model can release its final reference before decoding the # student and avoid a teacher+student pixel-memory peak. for key in ( "current_teacher_features", "current_teacher_video01", "ode_latent", ): batch.pop(key, None) image_or_video_shape = list(self.config.image_or_video_shape) image_or_video_shape[0] = batch_size first_pass_started = time.perf_counter() with torch.no_grad(): text_conditioning_error = None conditional_dict = None unconditional_dict = None t5_attention_mask = None try: ( conditional_dict, unconditional_dict, t5_attention_mask, ) = self._encode_fresh_video_text_conditioning( text_prompts, batch_size, ) except BaseException as exc: text_conditioning_error = exc if not self._all_ranks_true(text_conditioning_error is None): message = ( f": {type(text_conditioning_error).__name__}: " f"{text_conditioning_error}" if text_conditioning_error is not None else " on another rank" ) raise RuntimeError( "fresh text conditioning failed" + message ) from text_conditioning_error features, replay_descriptor = ( self.model.collect_rollout_features( image_or_video_shape=image_or_video_shape, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, t5_attention_mask=t5_attention_mask, initial_latent=None, vjepa2_window_sample_keys=( vjepa2_window_sample_keys), require_current_teacher=bool(gradient_selected), current_teacher_features=teacher_forward_inputs.pop( "current_teacher_features", None), current_teacher_video01=teacher_forward_inputs.pop( "current_teacher_video01", None), current_teacher_latent=teacher_forward_inputs.pop( "current_teacher_latent", None), ) ) first_pass_seconds = time.perf_counter() - first_pass_started current_teacher_features = replay_descriptor.pop( "current_teacher_features", None) teacher_extract_seconds = replay_descriptor.pop( "current_teacher_extract_seconds", None) if ( gradient_selected and bool(getattr( self.model, "current_teacher_attraction_enabled", False)) and current_teacher_features is None ): raise RuntimeError( "fresh first pass omitted matched current teacher features") cache_write_error = None try: if cache_records: self._write_current_teacher_cache( cache_records, current_teacher_features) except BaseException as exc: cache_write_error = exc if not self._all_ranks_true(cache_write_error is None): message = ( f": {type(cache_write_error).__name__}: {cache_write_error}" if cache_write_error is not None else " on another rank") raise RuntimeError( "live-teacher feature cache write failed" + message ) from cache_write_error # The call-site pops every mutually exclusive input while binding it; # retaining anything here would defeat the teacher payload handoff. if teacher_forward_inputs: raise RuntimeError( "unconsumed current teacher forward input remained") if teacher_cache_log is not None: teacher_cache_log["first_pass_seconds"] = first_pass_seconds if teacher_extract_seconds is not None: teacher_cache_log["extract_seconds"] = float( teacher_extract_seconds) entry = { "features": features, "replay_descriptor": replay_descriptor, "image_or_video_shape": image_or_video_shape, "conditional_dict": conditional_dict, "unconditional_dict": unconditional_dict, "gradient_selected": bool(gradient_selected), "sample_ids": sample_ids, } if current_teacher_features is not None: entry["current_teacher_features"] = current_teacher_features if teacher_cache_log is not None: entry["current_teacher_cache_log"] = teacher_cache_log return entry def _fresh_video_window_log(self, entries, reference): if getattr( self.model, "vjepa2_dense_frame_sampling", "center_contiguous", ) != "random_contiguous_per_video": return {} num_starts = ( int(self.model.vjepa2_dense_source_num_frames) - int(self.model.vjepa2_dense_num_frames) + 1 ) local_error = None local_starts = None try: start_parts = [] for entry in entries: indices = entry["replay_descriptor"][ "vjepa2_frame_indices"] if ( not isinstance(indices, torch.Tensor) or indices.ndim != 2 or int(indices.shape[1]) != int(self.model.vjepa2_dense_num_frames) ): raise ValueError( "invalid V-JEPA2 frame indices in GradCache entry") start_parts.append(indices[:, 0].to( device=reference.device, dtype=torch.long)) local_starts = torch.cat(start_parts) if int(local_starts.numel()) != self.fresh_video_local_samples: raise ValueError( "V-JEPA2 window count does not match local samples") if ( int(local_starts.min().item()) < 0 or int(local_starts.max().item()) >= num_starts ): raise ValueError("V-JEPA2 window start is out of range") except BaseException as exc: local_error = exc if not self._all_ranks_true(local_error is None): message = ( f": {type(local_error).__name__}: {local_error}" if local_error is not None else " on another rank" ) raise RuntimeError( "V-JEPA2 window logging validation failed" + message ) from local_error histogram = torch.bincount( local_starts, minlength=num_starts)[:num_starts] if dist.is_initialized() and self.world_size > 1: dist.all_reduce(histogram, op=dist.ReduceOp.SUM) if int(histogram.sum().item()) != self.fresh_video_total_samples: raise RuntimeError( "global V-JEPA2 window histogram has an invalid sample count") used = histogram.nonzero(as_tuple=False).flatten() if used.numel() == 0: raise RuntimeError("global V-JEPA2 window histogram is empty") return { "vjepa2_dense_window_start_min": used[0], "vjepa2_dense_window_start_max": used[-1], "vjepa2_dense_window_unique_starts": used.new_tensor( int(used.numel())), "vjepa2_dense_source_num_frames": used.new_tensor( int(self.model.vjepa2_dense_source_num_frames)), } def _validate_partial_mmd_geometry(self, result): return None def _run_distribution_matching_step(self): """Perform one distribution-matching update (paper Sec. 3.3). Collect B generated representations, differentiate the feature objective for b selected rollouts, then replay encoders, VAE and generator. The released selection and loss conventions are documented in paper-to-code.md.""" # Frozen encoders stay in eval mode; generator gradients are enabled only on replay. self.model.eval() dense_log_prefix = getattr( self.model, "dense_log_prefix", "vjepa2_dense") def dense_key(suffix): return f"{dense_log_prefix}_{suffix}" if self.step % 20 == 0: torch.cuda.empty_cache() # This cache lasts one optimizer update; it is not a history of older rollouts. cache = self.fresh_video_gradcache if cache is None: raise RuntimeError("fresh video GradCache is not initialized") if cache.local_samples or cache.finalized: raise RuntimeError( "fresh video GradCache leaked state across optimizer updates") if ( hasattr(self.model, "distribution_state_pending") and self.model.distribution_state_pending() ): raise RuntimeError( "dynamic distribution state leaked across optimizer updates") # This consensus must happen before the first model forward because # enabled ranks enter a pixel-cache collective inside that forward. use_prepared_replay = bool( self.uses_fresh_video_gradcache and self.model.fresh_video_prepared_replay_enabled ) all_use_prepared = self._all_ranks_true(use_prepared_replay) none_use_prepared = self._all_ranks_true( not use_prepared_replay) if not (all_use_prepared or none_use_prepared): raise RuntimeError( "fresh-video prepared replay setting differs across ranks") use_prepared_replay = all_use_prepared entries = [] try: for microbatch_index in range(self.fresh_video_microbatches): window_sample_keys = None if getattr( self.model, "vjepa2_dense_frame_sampling", "center_contiguous", ) == "random_contiguous_per_video": local_start = ( microbatch_index * int(self.config.batch_size)) global_start = ( dist.get_rank() * self.fresh_video_local_samples + local_start ) update_start = ( int(self.step) * self.fresh_video_total_samples) window_sample_keys = tuple( update_start + global_start + offset for offset in range(int(self.config.batch_size)) ) gradient_selected = bool( microbatch_index >= ( self.fresh_video_microbatches - self.fresh_video_generator_grad_microbatches) ) batch = ( self._matched_batch_for_microbatch(microbatch_index) if getattr( self, "matched_counter_sampling_enabled", False) else self._next_batch() ) entry = self._collect_rollout_microbatch( batch, gradient_selected=gradient_selected, vjepa2_window_sample_keys=window_sample_keys, ) add_error = None try: student_ids = entry.get("sample_ids") if hasattr(cache, "partial_gradient_enabled"): cache.add( entry["features"], entry.get("current_teacher_features"), gradient_selected=gradient_selected, sample_ids=student_ids, current_teacher_sample_ids=( student_ids if entry.get( "current_teacher_features") is not None else None ), ) else: # Preserve the original one-argument cache protocol # used by legacy trainer fixtures and adapters. cache.add(entry["features"]) entry.pop("current_teacher_features", None) except BaseException as exc: add_error = exc if not self._all_ranks_true(add_error is None): message = ( f": {type(add_error).__name__}: {add_error}" if add_error is not None else " on another rank" ) raise RuntimeError( "fresh-video GradCache add failed" + message ) from add_error if not gradient_selected: descriptor = entry.get("replay_descriptor", {}) descriptor.pop("cached_student_video01", None) entry.pop("conditional_dict", None) entry.pop("unconditional_dict", None) entry.pop("image_or_video_shape", None) entries.append(entry) if ( use_prepared_replay and self.fresh_video_generator_grad_subsample_enabled and microbatch_index < ( self.fresh_video_microbatches - self.fresh_video_generator_grad_microbatches) ): # These feature values participate in the full MMD scalar, # but no generator graph will be built for them. Drop the # update-local decoded pixel cache immediately instead of # retaining fourteen unnecessary videos on host memory. entry["replay_descriptor"].pop( "cached_student_video01", None) # Differentiate the global feature-space loss before constructing generator graphs. result = cache.finalize() if ( self.fresh_video_generator_grad_subsample_enabled and hasattr(cache, "partial_gradient_enabled") ): self._validate_partial_mmd_geometry(result) window_log = self._fresh_video_window_log(entries, result.loss) video_cotangents = result.fsdp_cotangents expected_replay_start = ( len(entries) - self.fresh_video_generator_grad_microbatches ) expected_replay_indices = tuple( range(expected_replay_start, len(entries))) replay_indices = tuple(getattr( result, "replay_batch_indices", expected_replay_indices, )) if replay_indices != expected_replay_indices: raise RuntimeError( "GradCache replay indices are not the configured suffix: " f"{replay_indices} != {expected_replay_indices}") replay_entries = tuple(entries[index] for index in replay_indices) if ( not replay_entries or len(replay_entries) != self.fresh_video_generator_grad_microbatches or ( hasattr(cache, "partial_gradient_enabled") and any( not entry.get("gradient_selected", False) for entry in replay_entries ) ) ): raise RuntimeError( "generator-gradient replay microbatch selection is invalid") if len(video_cotangents) == len(replay_indices): replay_video_cotangents = tuple(video_cotangents) elif len(video_cotangents) == len(entries): replay_video_cotangents = tuple( video_cotangents[index] for index in replay_indices) else: raise RuntimeError( "fresh GradCache cotangent/replay count mismatch") # Apply B/b selection compensation once; FSDP averaging has a # separate factor. replay_video_cotangents = _scale_replay_video_cotangents( replay_video_cotangents, self.fresh_video_generator_grad_sample_scale, ) # Replay frozen encoder/VAE VJPs before generator activations. prepared_replays = [None] * len(replay_entries) if use_prepared_replay: prepared_replays = [ self.model.prepare_latent_gradients( replay_descriptor=entry["replay_descriptor"], expected_features=entry["features"], feature_cotangent=cotangent, replay_atol=float(getattr( self.config, "fresh_video_replay_atol", 1e-5)), replay_rtol=float(getattr( self.config, "fresh_video_replay_rtol", 1e-4)), ) for entry, cotangent in zip( replay_entries, replay_video_cotangents) ] replay_logs = [] for entry, cotangent, prepared_replay in zip( replay_entries, replay_video_cotangents, prepared_replays, ): surrogate, replay_log = self.model.replay_generator_gradients( replay_descriptor=entry["replay_descriptor"], expected_features=entry["features"], feature_cotangent=cotangent, image_or_video_shape=entry["image_or_video_shape"], conditional_dict=entry["conditional_dict"], unconditional_dict=entry["unconditional_dict"], initial_latent=None, prepared_replay=prepared_replay, replay_atol=float(getattr( self.config, "fresh_video_replay_atol", 1e-5)), replay_rtol=float(getattr( self.config, "fresh_video_replay_rtol", 1e-4)), ) surrogate.backward() replay_logs.append(replay_log) generator_grad_norm = self.model.generator.clip_grad_norm_( self.max_grad_norm_generator) log_dict = { key: torch.stack([ item[key].detach().float() for item in replay_logs ]).mean() for key in replay_logs[0] } meta = result.metadata generator_loss = result.optimized_loss multilayer_metadata = meta.get("layers") displayed_mmd = ( meta["raw_mmd_sum"] if isinstance(multilayer_metadata, dict) else result.loss ) log_dict.update({ "generator_loss": generator_loss, "generator_grad_norm": generator_grad_norm, "dmdtrain_gradient_norm": log_dict[ "combined_latent_grad_norm"], dense_key("mmd"): displayed_mmd, dense_key("optimized_mmd"): result.optimized_loss, dense_key("exact_repulsion"): meta["repulsion"], dense_key("nystrom_attraction"): meta["attraction"], dense_key("nystrom_attraction_weight"): meta.get( "attraction_weight", result.loss.new_ones(())), dense_key("reference_constant"): meta[ "reference_constant"], dense_key("local_video_samples"): meta[ "gradcache_local_samples"], dense_key("global_video_samples"): meta[ "gradcache_global_samples"], dense_key("gradient_local_video_samples"): meta.get( "gradcache_gradient_local_samples", result.loss.new_tensor( self.fresh_video_generator_grad_local_samples, dtype=torch.long, ), ), dense_key("gradient_global_video_samples"): meta.get( "gradcache_gradient_global_samples", result.loss.new_tensor( self.fresh_video_generator_grad_global_samples, dtype=torch.long, ), ), dense_key("current_teacher_global_samples"): meta.get( "gradcache_current_teacher_global_samples", result.loss.new_zeros((), dtype=torch.long), ), dense_key("kss_denominator"): meta.get( "gradcache_kss_denominator", result.loss.new_tensor( self.fresh_video_total_samples ** 2, dtype=torch.long, ), ), dense_key("fresh_microbatches"): result.loss.new_tensor( len(entries), dtype=torch.long), dense_key("replay_microbatches"): result.loss.new_tensor( len(replay_entries), dtype=torch.long), dense_key("raw_tokens"): result.loss.new_tensor( getattr( self.model, "dense_raw_tokens", self.model.vjepa2_dense_raw_tokens, ), dtype=torch.long), dense_key("sampled_tokens"): result.loss.new_tensor( getattr( self.model, "dense_num_tokens", self.model.vjepa2_dense_num_tokens, ), dtype=torch.long), dense_key("feature_dim"): result.loss.new_tensor( getattr( self.model, "dense_feature_dim", self.model.vjepa2_dense_feature_dim, ), dtype=torch.long), dense_key("num_frames"): result.loss.new_tensor( getattr( self.model, "dense_num_frames", self.model.vjepa2_dense_num_frames, ), dtype=torch.long), dense_key("layer"): result.loss.new_tensor( getattr( self.model, "dense_layer", self.model.vjepa2_dense_layer, ), dtype=torch.long), dense_key("prefix_blocks"): result.loss.new_tensor( getattr( self.model, "dense_layer", self.model.vjepa2_dense_layer, ) + 1, dtype=torch.long), dense_key("nystrom_landmarks"): result.loss.new_tensor( self.model.latent_mmd.landmarks.shape[0], dtype=torch.long), dense_key("reference_samples"): result.loss.new_tensor( self.model.vjepa2_nystrom_reference_count, dtype=torch.long), "fresh_video_optimizer_step_skipped": ( result.loss.new_zeros((), dtype=torch.long)), "fresh_video_generator_grad_sample_scale": ( result.loss.new_tensor( self.fresh_video_generator_grad_sample_scale)), "fresh_video_generator_effective_replay_scale": ( result.loss.new_tensor( self.fresh_video_generator_grad_sample_scale * float(getattr(self, "world_size", 1)))), "fresh_video_generator_grad_subsample_enabled": ( result.loss.new_tensor( int(self.fresh_video_generator_grad_subsample_enabled), dtype=torch.long, ) ), "fresh_video_generator_grad_local_samples": ( result.loss.new_tensor( self.fresh_video_generator_grad_local_samples, dtype=torch.long, ) ), "fresh_video_generator_grad_global_samples": ( result.loss.new_tensor( self.fresh_video_generator_grad_global_samples, dtype=torch.long, ) ), "fresh_video_generator_grad_microbatches": ( result.loss.new_tensor( self.fresh_video_generator_grad_microbatches, dtype=torch.long, ) ), }) if isinstance(meta.get("branches"), Mapping): for branch_name, branch in meta["branches"].items(): prefix = f"{branch_name}_dense" log_dict.update({ f"{prefix}_raw_mmd": branch["raw_mmd"], f"{prefix}_normalized_mmd": branch["normalized_mmd"], f"{prefix}_normalization_denominator": ( branch["normalization_denominator"]), f"{prefix}_exact_repulsion": branch["repulsion"], f"{prefix}_fixed_teacher_attraction": ( branch["fixed_teacher_attraction"]), f"{prefix}_current_teacher_attraction": ( branch["current_teacher_attraction"]), f"{prefix}_reference_constant": ( branch["reference_constant"]), f"{prefix}_bandwidth": branch["sigma"], f"{prefix}_kss_denominator": ( branch["repulsion_denominator"]), f"{prefix}_fixed_attraction_denominator": ( branch["fixed_teacher_attraction_denominator"]), f"{prefix}_teacher_cross_denominator": ( branch["current_teacher_attraction_denominator"]), f"{prefix}_teacher_self_denominator": ( branch["current_teacher_self_kernel_denominator"]), }) log_dict["dual_dense_normalized_mmd_sum"] = meta[ "normalized_mmd_sum"] teacher_cache_logs = [ entry["current_teacher_cache_log"] for entry in entries if "current_teacher_cache_log" in entry ] if teacher_cache_logs: cache_stats = result.loss.new_tensor([ sum(item["hits"] for item in teacher_cache_logs), sum(item["misses"] for item in teacher_cache_logs), sum(item["decode_seconds"] for item in teacher_cache_logs), sum(item["first_pass_seconds"] for item in teacher_cache_logs), len(teacher_cache_logs), sum(item.get("extract_seconds", 0.0) for item in teacher_cache_logs), sum("extract_seconds" in item for item in teacher_cache_logs), ], dtype=torch.float64) if dist.is_initialized() and self.world_size > 1: dist.all_reduce(cache_stats, op=dist.ReduceOp.SUM) total_items = cache_stats[0] + cache_stats[1] teacher_metric_prefix = ( "triple_bank" if self.current_teacher_source == "matched_triple_feature_bank" else ("dinov3_vjepa2_bank" if self.current_teacher_source == "matched_dinov3_vjepa2_feature_bank" else ("dual_bank" if self.current_teacher_source == "matched_dual_feature_bank" else ( "matched" if self.current_teacher_source == "matched_manifest_video" else "openvid" ))) ) log_dict.update({ f"{teacher_metric_prefix}_teacher_cache_hits": cache_stats[0], f"{teacher_metric_prefix}_teacher_cache_misses": cache_stats[1], f"{teacher_metric_prefix}_teacher_cache_hit_rate": ( cache_stats[0] / total_items.clamp_min(1.0)), f"{teacher_metric_prefix}_teacher_decode_seconds": cache_stats[2], f"{teacher_metric_prefix}_teacher_first_pass_seconds_per_microbatch": ( cache_stats[3] / cache_stats[4].clamp_min(1.0)), }) if float(cache_stats[6].item()) > 0: log_dict[ f"{teacher_metric_prefix}_teacher_extract_seconds_per_microbatch" ] = cache_stats[5] / cache_stats[6] if "current_attraction" in meta: log_dict.update({ dense_key("nystrom_current_attraction"): ( meta["current_attraction"]), dense_key("nystrom_attraction_ema_updates"): ( meta["attraction_ema_updates"]), dense_key("nystrom_attraction_current_coefficient"): ( meta["attraction_current_coefficient"]), }) if "current_teacher_attraction" in meta: log_dict.update({ dense_key("fixed_teacher_attraction"): ( meta["fixed_teacher_attraction"]), dense_key("current_teacher_attraction"): ( meta["current_teacher_attraction"]), dense_key("mixed_current_attraction"): ( meta["mixed_current_attraction"]), dense_key("fixed_teacher_normalized_weight"): ( meta[ "fixed_teacher_attraction_normalized_weight"]), dense_key("current_teacher_normalized_weight"): ( meta[ "current_teacher_attraction_normalized_weight"]), dense_key("current_teacher_self_kernel"): ( meta["current_teacher_self_kernel"]), dense_key("current_teacher_exact_mmd"): ( meta["current_teacher_exact_mmd"]), dense_key("current_teacher_samples"): ( meta["current_teacher_samples"]), }) if "attraction_ema_blend_current_coefficient" in meta: log_dict.update({ dense_key("attraction_ema_blend_coefficient"): ( meta[ "attraction_ema_blend_current_coefficient"]), dense_key("attraction_gradient_coefficient"): ( meta[ "attraction_gradient_current_coefficient"]), dense_key("effective_ema_loss"): ( meta["effective_ema_loss"]), }) if "repulsion_fifo_enabled" in meta: log_dict.update({ dense_key("repulsion_fifo_capacity"): ( meta["repulsion_fifo_capacity"]), dense_key("repulsion_fifo_fill_before"): ( meta["repulsion_fifo_fill_before"]), dense_key("repulsion_fifo_fill_after_commit"): ( meta["repulsion_fifo_fill_after_commit"]), dense_key("repulsion_fifo_commits_before"): ( meta["repulsion_fifo_commits"]), dense_key("repulsion_current_offdiag_kernel_sum"): ( meta["repulsion_current_offdiag_kernel_sum"]), dense_key("repulsion_history_kernel_sum"): ( meta["repulsion_history_kernel_sum"]), dense_key("repulsion_denominator"): ( meta["repulsion_denominator"]), }) if isinstance(multilayer_metadata, dict): log_dict.update({ "vjepa2_dense_multilayer_normalized_sum": result.loss, "vjepa2_dense_multilayer_weighted_sum": ( result.optimized_loss), "vjepa2_dense_num_layers": meta["num_layers"], "vjepa2_dense_sampled_tokens_per_layer": ( result.loss.new_tensor( self.model.vjepa2_dense_num_tokens, dtype=torch.long, ) ), "vjepa2_dense_gradcache_transport_tokens": ( result.loss.new_tensor( self.model.vjepa2_dense_num_tokens * len(self.model.vjepa2_dense_layers), dtype=torch.long, ) ), "vjepa2_dense_self_normalized": ( meta["self_normalized"]), "vjepa2_dense_normalization_epsilon": ( meta["normalization_epsilon"]), "vjepa2_dense_layer_min": result.loss.new_tensor( min(self.model.vjepa2_dense_layers), dtype=torch.long, ), "vjepa2_dense_layer_max": result.loss.new_tensor( max(self.model.vjepa2_dense_layers), dtype=torch.long, ), }) for layer in self.model.vjepa2_dense_layers: layer_meta = multilayer_metadata[str(layer)] prefix = f"vjepa2_dense_l{layer}" log_dict.update({ f"{prefix}_mmd": layer_meta["raw_mmd"], f"{prefix}_normalized_mmd": ( layer_meta["normalized_mmd"]), f"{prefix}_weighted_mmd": ( layer_meta["weighted_mmd"]), f"{prefix}_repulsion": ( layer_meta["repulsion"]), f"{prefix}_nystrom_attraction": ( layer_meta["attraction"]), f"{prefix}_reference_constant": ( layer_meta["reference_constant"]), f"{prefix}_bandwidth": layer_meta["sigma"], }) if "current_attraction" in layer_meta: log_dict.update({ f"{prefix}_nystrom_current_attraction": ( layer_meta["current_attraction"]), f"{prefix}_attraction_ema_updates": ( layer_meta["attraction_ema_updates"]), f"{prefix}_attraction_current_coefficient": ( layer_meta[ "attraction_current_coefficient"]), }) elif not isinstance(meta.get("branches"), Mapping): log_dict[dense_key("bandwidth")] = meta["sigma"] log_dict.update(window_log) return log_dict except FloatingPointError as exc: cache.discard() if hasattr(self.model, "discard_distribution_state"): self.model.discard_distribution_state() if self.is_main_process: print( "[warn] discarding complete fresh-video update after " f"non-finite value: {exc}", flush=True, ) nonfinite = torch.tensor( float("nan"), device=self.device, dtype=torch.float32) return { "generator_loss": nonfinite, "generator_grad_norm": nonfinite, "dmdtrain_gradient_norm": nonfinite, dense_key("update_discarded"): nonfinite.new_ones(()), } except BaseException: cache.discard() if hasattr(self.model, "discard_distribution_state"): self.model.discard_distribution_state() raise finally: # Prepared-replay pixels are update-local. Never let a decoded # batch leak into the next update after success or failure. for entry in entries: descriptor = entry.get("replay_descriptor", {}) if isinstance(descriptor, dict): descriptor.pop("cached_student_video01", None) def fwdbwd_one_step(self, batch, train_generator): # Frozen encoders stay in eval mode; generator gradients are enabled only on replay. self.model.eval() # prevent any randomness (e.g. dropout) if self.step % 20 == 0: torch.cuda.empty_cache() # Step 1: Get the next batch of text prompts text_prompts = batch["prompts"] if self.config.distribution_loss == "wan_kv_mmd": teacher_trajectory = batch["ode_latent"] local_shape_ok = teacher_trajectory.ndim == 6 if not self._all_ranks_true(local_shape_ok): raise ValueError( "wan_kv_mmd teacher LMDB yielded an invalid shape on at " "least one rank; expected " "[B, steps, frames, channels, height, width]") clean_latent = teacher_trajectory[:, -1].to( device=self.device, dtype=self.dtype) local_finite = bool(torch.isfinite(clean_latent).all().item()) if not self._all_ranks_true(local_finite): raise FloatingPointError( "teacher latent contains non-finite values on at least one rank") image_latent = None elif self.config.i2v: clean_latent = None image_latent = batch["ode_latent"][:, -1][:, 0:1, ].to( device=self.device, dtype=self.dtype) else: clean_latent = None image_latent = None batch_size = len(text_prompts) image_or_video_shape = list(self.config.image_or_video_shape) image_or_video_shape[0] = batch_size # Step 2: Extract the conditional infos with torch.no_grad(): conditional_dict = self.model.text_encoder( text_prompts=text_prompts) if not getattr(self, "unconditional_dict", None): unconditional_dict = self.model.text_encoder( text_prompts=[self.config.negative_prompt] * batch_size) unconditional_dict = {k: v.detach() for k, v in unconditional_dict.items()} self.unconditional_dict = unconditional_dict # cache the unconditional_dict else: unconditional_dict = self.unconditional_dict # Step 3: Store gradients for the generator (if training the generator) if train_generator: generator_loss, generator_log_dict = self.model.generator_loss( image_or_video_shape=image_or_video_shape, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, clean_latent=clean_latent, initial_latent=image_latent if self.config.i2v else None ) generator_loss.backward() generator_grad_norm = self.model.generator.clip_grad_norm_( self.max_grad_norm_generator) generator_log_dict.update({"generator_loss": generator_loss, "generator_grad_norm": generator_grad_norm}) return generator_log_dict else: generator_log_dict = {} # Step 4: Store gradients for the critic (if training the critic) critic_loss, critic_log_dict = self.model.critic_loss( image_or_video_shape=image_or_video_shape, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, clean_latent=clean_latent, initial_latent=image_latent if self.config.i2v else None ) critic_loss.backward() critic_grad_norm = self.model.fake_score.clip_grad_norm_( self.max_grad_norm_critic) critic_log_dict.update({"critic_loss": critic_loss, "critic_grad_norm": critic_grad_norm}) return critic_log_dict def generate_video(self, pipeline, prompts, image=None): batch_size = len(prompts) if image is not None: image = image.squeeze(0).unsqueeze(0).unsqueeze(2).to(device="cuda", dtype=torch.bfloat16) # Encode the input image as the first latent initial_latent = pipeline.vae.encode_to_latent(image).to(device="cuda", dtype=torch.bfloat16) initial_latent = initial_latent.repeat(batch_size, 1, 1, 1, 1) sampled_noise = torch.randn( [batch_size, self.model.num_training_frames - 1, 16, 60, 104], device="cuda", dtype=self.dtype ) else: initial_latent = None sampled_noise = torch.randn( [batch_size, self.model.num_training_frames, 16, 60, 104], device="cuda", dtype=self.dtype ) video, _ = pipeline.inference( noise=sampled_noise, text_prompts=prompts, return_latents=True, initial_latent=initial_latent ) current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0 return current_video def train(self): start_step = self.step max_steps = int(getattr(self.config, "max_steps", 0)) if max_steps > 0 and self.step >= max_steps: if self.is_main_process: print( f"[done] checkpoint step {self.step} already reached " f"max_steps={max_steps}; nothing to train.", flush=True, ) return # baseline eval before any optimization (captures the init / resumed model) if self.eval_interval > 0: self._fd_eval_inference() while True: # ``self.step`` is the last committed/public step. Apply the LR # for the update about to produce the next completed step. self._active_generator_lr = self._set_generator_learning_rate( self.step + 1) TRAIN_GENERATOR = ( (not self.uses_critic) or self.step % self.config.dfake_gen_update_ratio == 0 ) update_committed = True # Train the generator if TRAIN_GENERATOR: self.generator_optimizer.zero_grad(set_to_none=True) if self.uses_fresh_video_gradcache: generator_log_dict = ( self._run_distribution_matching_step()) else: extras_list = [] batch = self._next_batch() extra = self.fwdbwd_one_step(batch, True) extras_list.append(extra) generator_log_dict = merge_dict_list(extras_list) # Skip the step if the generator gradient is non-finite (avoids # corrupting weights from a stray bf16 overflow); the gradients are # zeroed at the top of the next iteration. _ggn = generator_log_dict.get("generator_grad_norm", None) _local_ggn_ok = ( (_ggn is None) or bool(torch.isfinite( torch.as_tensor(_ggn).float()).all().item()) ) _ggn_ok = self._all_ranks_true(_local_ggn_ok) if _ggn_ok: prepare_error = None try: if hasattr( self.model, "prepare_distribution_state" ): self.model.prepare_distribution_state() except Exception as exc: prepare_error = exc prepare_ok = self._all_ranks_true( prepare_error is None) if not prepare_ok: if self.fresh_video_gradcache is not None: self.fresh_video_gradcache.discard() if hasattr(self.model, "discard_distribution_state"): self.model.discard_distribution_state() if prepare_error is not None: raise RuntimeError( "distribution transaction preparation failed " "before optimizer.step") from prepare_error raise RuntimeError( "distribution transaction preparation failed on " "another rank before optimizer.step") try: # Commit after every selected rollout contributed its gradient. self.generator_optimizer.step() if hasattr(self.model, "commit_distribution_state"): self.model.commit_distribution_state() if self.fresh_video_gradcache is not None: self.fresh_video_gradcache.commit() except BaseException: if self.fresh_video_gradcache is not None: self.fresh_video_gradcache.discard() if hasattr(self.model, "discard_distribution_state"): self.model.discard_distribution_state() raise update_committed = True if self.generator_ema is not None: self.generator_ema.update(self.model.generator) else: update_committed = False if self.fresh_video_gradcache is not None: self.fresh_video_gradcache.discard() if hasattr(self.model, "discard_distribution_state"): self.model.discard_distribution_state() if self.is_main_process: print(f"[warn] step {self.step}: non-finite generator_grad_norm={_ggn}; skipping optimizer step", flush=True) # Train the critic critic_log_dict = {} if self.uses_critic: self.critic_optimizer.zero_grad(set_to_none=True) extras_list = [] batch = self._next_batch() extra = self.fwdbwd_one_step(batch, False) extras_list.append(extra) critic_log_dict = merge_dict_list(extras_list) self.critic_optimizer.step() # Increment the step since we finished gradient update self.step += 1 # Create EMA params (if not already created) ema_start_step = int(getattr(self.config, "ema_start_step", 0)) ema_weight = float(getattr(self.config, "ema_weight", 0.0) or 0.0) if (self.step >= ema_start_step) and \ (self.generator_ema is None) and (ema_weight > 0): self.generator_ema = EMA_FSDP( self.model.generator, decay=ema_weight) # Save the model if (not self.config.no_save) and (self.step - start_step) > 0 and self.step % self.config.log_iters == 0: torch.cuda.empty_cache() self.save() torch.cuda.empty_cache() # Periodic eval inference (all ranks participate; rank 0 logs/saves). # Remember whether this step was evaluated so max_steps does not # write the same fixed-seed samples twice. did_periodic_eval = ( self.eval_interval > 0 and (self.step - start_step) > 0 and self.step % self.eval_interval == 0 ) if did_periodic_eval: torch.cuda.empty_cache() self._fd_eval_inference() torch.cuda.empty_cache() # Logging if self.is_main_process: wandb_loss_dict = {} if TRAIN_GENERATOR: wandb_loss_dict.update( { "generator_loss": generator_log_dict["generator_loss"].mean().item(), "generator_grad_norm": generator_log_dict["generator_grad_norm"].mean().item(), "dmdtrain_gradient_norm": generator_log_dict["dmdtrain_gradient_norm"].mean().item(), "generator_lr": float(self._active_generator_lr), } ) for key, value in generator_log_dict.items(): if key in wandb_loss_dict or key in { "generator_loss", "generator_grad_norm", "dmdtrain_gradient_norm", }: continue if isinstance(value, torch.Tensor) and value.numel() == 1: # Metadata includes integral counters (for example # Wan-KV history_updates). ``Tensor.mean`` rejects # integer dtypes; a one-element tensor needs no # reduction. wandb_loss_dict[key] = value.detach().item() if self.uses_critic: wandb_loss_dict.update( { "critic_loss": critic_log_dict["critic_loss"].mean().item(), "critic_grad_norm": critic_log_dict["critic_grad_norm"].mean().item() } ) if not self.disable_wandb: wandb.log(wandb_loss_dict, step=self.step) if (self.step % getattr(self.config, "stdout_log_iters", 1)) == 0: summary = " ".join( f"{k}={v:.4g}" for k, v in wandb_loss_dict.items() if isinstance(v, (int, float)) ) print(f"[step {self.step}] {summary}", flush=True) if self.step % self.config.gc_interval == 0: if dist.get_rank() == 0: logging.info("DistGarbageCollector: Running GC.") gc.collect() torch.cuda.empty_cache() if self.is_main_process: current_time = time.time() if self.previous_time is None: self.previous_time = current_time else: if not self.disable_wandb: wandb.log({"per iteration time": current_time - self.previous_time}, step=self.step) self.previous_time = current_time # Stop only after logging the final step. This is used by smoke tests # and sequential launchers, so suppressing its metrics would hide # precisely the EMA transition they need to verify. max_steps = int(getattr(self.config, "max_steps", 0)) if max_steps > 0 and self.step >= max_steps: if not self.config.no_save: torch.cuda.empty_cache() self.save() torch.cuda.empty_cache() if self.eval_interval > 0 and not did_periodic_eval: self._fd_eval_inference() if self.is_main_process: print( f"[done] reached max_steps={max_steps}; exiting.", flush=True, ) break # Historical entry points remain aliases for external callers. _prepare_fresh_video_microbatch = _collect_rollout_microbatch _fresh_video_generator_update = _run_distribution_matching_step # Old import/class spelling remains valid for existing callers. Trainer = DistributionMatchingTrainer