| """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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
| ) |
| |
| 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 |
|
|
| |
| 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)) |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| 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: |
| |
| |
| |
| 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: |
| |
| |
| |
| |
| 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}") |
| |
| |
| 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 |
|
|
| |
| |
| |
| |
| |
| 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 |
| |
| 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: |
| |
| |
| 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") |
| |
| |
| 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() |
| |
| 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)) |
|
|
| |
| gen_optim_full = None |
| critic_optim_full = None |
| if write_full_state: |
| |
| 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}") |
|
|
| |
| 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: |
| |
| |
| |
| 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}") |
| |
| 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) |
| |
| 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, |
| ) |
|
|
| |
| |
| 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) |
|
|
| |
| 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 |
| |
| self.model.eval() |
| |
| 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 |
| |
| |
| |
| |
| 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") |
| 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: |
| print(f"[eval] failed to write {out_path}: {exc}") |
| |
| if not self.disable_wandb and wrote: |
| try: |
| log_videos[f"eval/sample_{i}"] = wandb.Video(out_path, format="mp4") |
| except Exception as exc: |
| 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) |
|
|
| |
| |
| |
| |
| 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 |
| |
| |
| 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.""" |
| |
| 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() |
| |
| 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") |
| |
| |
| 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: |
| |
| |
| 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) |
| ): |
| |
| |
| |
| |
| entry["replay_descriptor"].pop( |
| "cached_student_video01", None) |
| |
| 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") |
| |
| |
| replay_video_cotangents = _scale_replay_video_cotangents( |
| replay_video_cotangents, |
| self.fresh_video_generator_grad_sample_scale, |
| ) |
|
|
| |
| 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: |
| |
| |
| 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): |
| |
| self.model.eval() |
|
|
| if self.step % 20 == 0: |
| torch.cuda.empty_cache() |
|
|
| |
| 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 |
|
|
| |
| 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 |
| else: |
| unconditional_dict = self.unconditional_dict |
|
|
| |
| 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 = {} |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| if self.eval_interval > 0: |
| self._fd_eval_inference() |
|
|
| while True: |
| |
| |
| 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 |
|
|
| |
| 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) |
| |
| |
| |
| _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: |
| |
| 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) |
|
|
| |
| 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() |
|
|
| |
| self.step += 1 |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
| |
| |
| |
| 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() |
|
|
| |
| 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: |
| |
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| 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 |
|
|
| |
| _prepare_fresh_video_microbatch = _collect_rollout_microbatch |
| _fresh_video_generator_update = _run_distribution_matching_step |
|
|
| |
| Trainer = DistributionMatchingTrainer |
|
|
