| """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, |
| move_adamw_moment_state_, |
| 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 safetensors.torch import load_file as load_safetensors_file |
| 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 = "/path/to/local-assets" |
| logical_root = "/path/to/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", |
| } |
| 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 _cosine_learning_rate(update_index, *, start, end, total_updates): |
| """Endpoint-inclusive cosine LR for the optimizer update at update_index.""" |
| if type(update_index) is not int or update_index < 0: |
| raise ValueError( |
| "learning-rate update_index must be a non-negative integer") |
| if type(total_updates) is not int or total_updates <= 0: |
| raise ValueError( |
| "lr_schedule_total_steps must be a positive integer") |
| start = float(start) |
| end = float(end) |
| if not math.isfinite(start) or not math.isfinite(end): |
| raise ValueError("learning-rate endpoints must be finite") |
| if start <= 0.0 or end <= 0.0 or end > start: |
| raise ValueError( |
| "learning-rate endpoints must satisfy start >= end > 0") |
| if total_updates == 1: |
| return end |
| position = min(update_index, total_updates - 1) / (total_updates - 1) |
| return end + 0.5 * (start - end) * ( |
| 1.0 + math.cos(math.pi * position)) |
|
|
|
|
| def _learning_rate_schedule_identity(config): |
| """Validate and freeze the optimizer schedule used by this experiment.""" |
| schedule = str( |
| getattr(config, "lr_schedule", "constant")).strip().lower() |
| if schedule == "constant": |
| value = float(config.lr) |
| if not math.isfinite(value) or value <= 0.0: |
| raise ValueError( |
| "constant learning rate must be finite and positive") |
| return {"name": "constant", "value": value} |
| if schedule != "cosine": |
| raise ValueError(f"unsupported lr_schedule {schedule!r}") |
| total_updates = int(getattr(config, "lr_schedule_total_steps", 0)) |
| max_steps = int(getattr(config, "max_steps", 0)) |
| if not 0 < max_steps <= total_updates: |
| raise ValueError( |
| "endpoint-inclusive cosine schedule requires " |
| "0 < max_steps <= lr_schedule_total_steps") |
| start = float(config.lr) |
| end = float(getattr(config, "lr_final", 0.0)) |
| first = _cosine_learning_rate( |
| 0, start=start, end=end, total_updates=total_updates) |
| last = _cosine_learning_rate( |
| total_updates - 1, |
| start=start, |
| end=end, |
| total_updates=total_updates, |
| ) |
| if first != start or last != end: |
| raise RuntimeError("cosine learning-rate endpoint invariant failed") |
| return { |
| "name": "cosine_endpoint_inclusive_v1", |
| "start": start, |
| "end": end, |
| "total_updates": total_updates, |
| } |
|
|
|
|
| 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 _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() |
|
|
|
|
| def _fixed_openvid_experiment_identity( |
| config, |
| model, |
| ledger_metadata, |
| *, |
| forward_global_samples, |
| gradient_global_samples, |
| gradient_compensation, |
| ): |
| """Build the immutable fixed-only identity saved in every checkpoint.""" |
| if ( |
| str(getattr(model, "feature_type", "")) |
| != "dual_dinov3_vjepa2_dense_latent" |
| or bool(getattr(model, "current_teacher_attraction_enabled", True)) |
| or str(getattr(config, "current_teacher_source", "")) != "none" |
| or int(getattr(config, "expected_current_teacher_samples", -1)) != 0 |
| ): |
| raise ValueError( |
| "fixed OpenVid experiment identity requires dual fixed-only mode") |
| branch_weights = { |
| name: float(value) |
| for name, value in model.latent_mmd.branch_weights.items() |
| } |
| if branch_weights != {"dinov3": 10.0, "vjepa2": 1.0}: |
| raise ValueError("fixed OpenVid branch weights must be exactly 10:1") |
| registry = getattr(model, "dual_branch_registry", {}) |
| if set(registry) != {"dinov3", "vjepa2"}: |
| raise ValueError("fixed OpenVid branch registry is incomplete") |
| branches = {} |
| for name in sorted(registry): |
| branch = registry[name] |
| branches[name] = { |
| "nystrom_landmarks": int(branch["nystrom_landmarks"]), |
| "reference_file_sha256": _strict_sha256( |
| branch["reference_file_sha256"], |
| f"{name} reference_file_sha256"), |
| "reference_hash": _strict_sha256( |
| branch["reference_hash"], f"{name} reference_hash"), |
| "reference_path": str(branch["reference_path"]), |
| "sample_shape": list(branch["sample_shape"]), |
| } |
| expected_branch_geometry = { |
| "dinov3": {"shape": [1024, 1024], "landmarks": 256}, |
| "vjepa2": {"shape": [1024, 1024], "landmarks": 1024}, |
| } |
| for name, expected in expected_branch_geometry.items(): |
| if ( |
| branches[name]["sample_shape"] != expected["shape"] |
| or branches[name]["nystrom_landmarks"] != expected["landmarks"] |
| ): |
| raise ValueError(f"fixed OpenVid {name} geometry changed") |
|
|
| bundle_sha = _strict_sha256( |
| getattr(config, "dinov3_vjepa2_reference_bundle_manifest_sha256", ""), |
| "dinov3_vjepa2_reference_bundle_manifest_sha256", |
| ) |
| if bundle_sha != str(getattr( |
| model, "dinov3_vjepa2_reference_bundle_manifest_sha256", "")): |
| raise ValueError("fixed OpenVid model/config bundle SHA mismatch") |
| ledger_sha = _strict_sha256( |
| ledger_metadata.get("ledger_sha256", ""), "OpenVid ledger SHA-256") |
| sequence_sha = _strict_sha256( |
| ledger_metadata.get("video_name_sequence_sha256", ""), |
| "OpenVid video-name sequence SHA-256", |
| ) |
| if ( |
| int(ledger_metadata.get("reference_count", -1)) != 77_615 |
| or sequence_sha != _strict_sha256( |
| getattr(config, "openvid_video_name_sequence_sha256", ""), |
| "configured OpenVid video-name sequence SHA-256", |
| ) |
| ): |
| raise ValueError("fixed OpenVid ledger identity/count mismatch") |
| generator_sha = _strict_sha256( |
| getattr(config, "generator_ckpt_sha256", ""), |
| "generator_ckpt_sha256", |
| ) |
| generator_model = str(getattr( |
| config, "generator_ckpt_model_name", "")) |
| generator_recipe = str(getattr( |
| config, "generator_ckpt_ode_recipe", "")) |
| generator_step = int(getattr(config, "generator_ckpt_step", -1)) |
| generator_prompt_sha = _strict_sha256( |
| getattr(config, "generator_ckpt_prompt_manifest_sha256", ""), |
| "generator_ckpt_prompt_manifest_sha256", |
| ) |
| model_kwargs = dict(getattr(config, "model_kwargs", {})) |
| schedule = { |
| "denoising_step_list": list(config.denoising_step_list), |
| "model_timestep_shift": float(model_kwargs.get("timestep_shift", -1)), |
| "timestep_shift": float(config.timestep_shift), |
| "warp_denoising_step": bool(config.warp_denoising_step), |
| } |
| if ( |
| generator_model != "Wan2.1-T2V-14B" |
| or generator_recipe != "causvid_14b_16k" |
| or generator_step <= 0 |
| or str(model_kwargs.get("model_name", "")) != generator_model |
| or schedule != { |
| "denoising_step_list": [1000, 757, 522, 0], |
| "model_timestep_shift": 8.0, |
| "timestep_shift": 8.0, |
| "warp_denoising_step": False, |
| } |
| ): |
| raise ValueError("fixed OpenVid Wan-14B ODE/schedule identity changed") |
| forward = int(forward_global_samples) |
| gradient = int(gradient_global_samples) |
| if ( |
| (forward, gradient) != (256, 128) |
| or float(gradient_compensation) != 2.0 |
| ): |
| raise ValueError("fixed OpenVid partial-gradient geometry changed") |
| return { |
| "mode": "openvid77615_fixed_only_dual_v1", |
| "current_teacher_source": "none", |
| "reference_bundle": { |
| "path": str(config.dinov3_vjepa2_reference_bundle_manifest), |
| "sha256": bundle_sha, |
| }, |
| "ledger": { |
| "path": str(config.fixed_student_prompt_ledger), |
| "sha256": ledger_sha, |
| "reference_count": 77_615, |
| "video_name_sequence_sha256": sequence_sha, |
| }, |
| "generator_initialization": { |
| "path": str(config.generator_ckpt), |
| "sha256": generator_sha, |
| "model_name": generator_model, |
| "ode_recipe": generator_recipe, |
| "ode_step": generator_step, |
| "prompt_manifest_sha256": generator_prompt_sha, |
| }, |
| "schedule": schedule, |
| "branches": branches, |
| "branch_weights": branch_weights, |
| "partial_gradient": { |
| "forward_global_samples": forward, |
| "gradient_global_samples": gradient, |
| "teacher_global_samples": 0, |
| "full_student_partial_gradient": True, |
| "kss_denominator": 65_536, |
| "fixed_attraction_denominator": 256, |
| "teacher_cross_denominator": 0, |
| "teacher_self_denominator": 0, |
| "replay_cotangent_compensation": float( |
| gradient_compensation), |
| }, |
| "feature_protocol_identity": str(config.feature_protocol_identity), |
| } |
|
|
|
|
| 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() |
| expected_world_size = int( |
| getattr(config, "expected_world_size", 0) or 0) |
| if expected_world_size and self.world_size != expected_world_size: |
| raise ValueError( |
| "distributed world-size mismatch: expected " |
| f"{expected_world_size}, got {self.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)) |
| self.fresh_video_physical_forward_batch_size = int(config.batch_size) |
| self.fresh_video_logical_micro_batch_size = int(getattr( |
| config, |
| "fresh_video_logical_micro_batch_size", |
| self.fresh_video_physical_forward_batch_size, |
| )) |
| self.fresh_video_physical_replay_batch_size = int(getattr( |
| config, |
| "fresh_video_physical_replay_batch_size", |
| self.fresh_video_physical_forward_batch_size, |
| )) |
| if self.uses_fresh_video_gradcache: |
| sizes = { |
| "fresh_video_local_samples": self.fresh_video_local_samples, |
| "fresh_video_logical_micro_batch_size": ( |
| self.fresh_video_logical_micro_batch_size), |
| "batch_size": self.fresh_video_physical_forward_batch_size, |
| "fresh_video_physical_replay_batch_size": ( |
| self.fresh_video_physical_replay_batch_size), |
| } |
| if any(value <= 0 for value in sizes.values()): |
| raise ValueError( |
| f"fresh-video execution sizes must be positive: {sizes}") |
| if ( |
| self.fresh_video_local_samples |
| % self.fresh_video_logical_micro_batch_size |
| or self.fresh_video_logical_micro_batch_size |
| % self.fresh_video_physical_forward_batch_size |
| or self.fresh_video_logical_micro_batch_size |
| % self.fresh_video_physical_replay_batch_size |
| ): |
| raise ValueError( |
| "VACE-aligned logical/physical batches must divide " |
| f"exactly: {sizes}") |
| self.fresh_video_microbatches = ( |
| self.fresh_video_local_samples |
| // self.fresh_video_logical_micro_batch_size) |
| else: |
| self.fresh_video_microbatches = 1 |
| expected_total_batch = self.fresh_video_local_samples * self.world_size |
| 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 fresh_video_local_samples * " |
| f"world_size; got {configured_total_batch} != " |
| f"{self.fresh_video_local_samples} * {self.world_size} " |
| 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=self.fresh_video_logical_micro_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) != 0 |
| 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 |
|
|
| |
| if config.distribution_loss != "wan_kv_mmd": |
| raise ValueError("This release supports distribution matching (wan_kv_mmd).") |
| 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") |
| expected_teacher_samples = ( |
| self.fresh_video_generator_grad_global_samples |
| if bool(getattr( |
| self.model, "current_teacher_attraction_enabled", False)) |
| else 0 |
| ) |
| if ( |
| configured_teacher_samples is not None |
| and int(configured_teacher_samples) != expected_teacher_samples |
| ): |
| raise ValueError( |
| "expected_current_teacher_samples disagrees with the active " |
| "teacher mode: " |
| f"{int(configured_teacher_samples)} != " |
| f"{expected_teacher_samples}") |
| 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 |
| ), |
| backward_prefetch=str(getattr( |
| config, "fsdp_backward_prefetch", "pre")), |
| ) |
|
|
| 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, |
| foreach=bool(getattr(config, "optimizer_foreach", False)), |
| fused=bool(getattr(config, "optimizer_fused", False)), |
| ) |
| self.optimizer_state_cpu_offload_during_backward = bool(getattr( |
| config, "optimizer_state_cpu_offload_during_backward", False)) |
| self.learning_rate_schedule = _learning_rate_schedule_identity(config) |
|
|
| 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.active_mmd_branches import active_mmd_branches |
| 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 = set(active_mmd_branches( |
| getattr(config, "triple_active_branches", None))) |
| 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") |
| 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", "")) != expected_feature_type: |
| 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") |
| 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 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: |
| kwargs["active_branches"] = active_mmd_branches( |
| getattr(config, "triple_active_branches", None)) |
| kwargs["expected_shard_size"] = 32 |
| if "dinov3" in expected_branch_names: |
| 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_10_to_1_to_1_then_divide_3" |
| if "dinov3" in expected_branch_names |
| else "sum_1_to_1_then_divide_3") |
| 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), |
| } |
| elif self.current_teacher_source == "none": |
| fixed_prompt_ledger = getattr( |
| config, "fixed_student_prompt_ledger", None) |
| if fixed_prompt_ledger: |
| from utils.openvid_reference_ledger import ( |
| OpenVidReferencePromptDataset, |
| ) |
| dataset = OpenVidReferencePromptDataset( |
| fixed_prompt_ledger, |
| expected_count=int(getattr( |
| config, "openvid_reference_count", 77615)), |
| ) |
| self.experiment_identity = ( |
| _fixed_openvid_experiment_identity( |
| config, |
| self.model, |
| dataset.metadata, |
| forward_global_samples=( |
| self.fresh_video_total_samples), |
| gradient_global_samples=( |
| self.fresh_video_generator_grad_global_samples), |
| gradient_compensation=( |
| self.fresh_video_generator_grad_sample_scale), |
| ) |
| ) |
| else: |
| dataset = TextDataset(config.data_path) |
| 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)) |
|
|
| |
| |
| |
| if self.experiment_identity is None: |
| self.experiment_identity = {} |
| self.experiment_identity["optimizer_learning_rate_schedule"] = ( |
| copy.deepcopy(self.learning_rate_schedule)) |
| self.experiment_identity["vace_aligned_execution"] = { |
| "fsdp_backward_prefetch": str(getattr( |
| config, "fsdp_backward_prefetch", "pre")), |
| "optimizer_foreach": bool(getattr( |
| config, "optimizer_foreach", False)), |
| "optimizer_fused": bool(getattr( |
| config, "optimizer_fused", False)), |
| "optimizer_state_cpu_offload_during_backward": bool( |
| self.optimizer_state_cpu_offload_during_backward), |
| "logical_micro_batch_size": int(getattr( |
| config, "fresh_video_logical_micro_batch_size", |
| config.batch_size)), |
| "physical_forward_batch_size": int(config.batch_size), |
| "physical_replay_batch_size": int(getattr( |
| config, "fresh_video_physical_replay_batch_size", |
| config.batch_size)), |
| } |
| self.experiment_identity["generator_initialization"] = { |
| "path": str(config.generator_ckpt), |
| "sha256": _strict_sha256( |
| config.generator_ckpt_sha256, "generator_ckpt_sha256"), |
| "replay_activation_cpu_offload": bool( |
| self.model.generator_replay_activation_cpu_offload), |
| "replay_activation_cpu_offload_block_stride": int( |
| self.model. |
| generator_replay_activation_cpu_offload_block_stride), |
| "replay_frozen_encoder_cpu_offload": bool( |
| self.model.generator_replay_frozen_encoder_cpu_offload), |
| } |
|
|
| |
| |
| |
| 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, |
| ) |
| ) |
| fixed_openvid_dual = bool( |
| getattr(self.model, "feature_type", "") |
| == "dual_dinov3_vjepa2_dense_latent" |
| and not getattr( |
| self.model, "current_teacher_attraction_enabled", False) |
| ) |
| if resume_checkpoint: |
| |
| |
| |
| if fixed_openvid_dual: |
| expected_resume_sha256 = _strict_sha256( |
| getattr(config, "resume_checkpoint_sha256", ""), |
| "resume_checkpoint_sha256", |
| ) |
| actual_resume_sha256 = self.model._distributed_file_sha256( |
| str(resume_checkpoint)) |
| if actual_resume_sha256 != expected_resume_sha256: |
| raise ValueError( |
| "fixed-only exact resume checkpoint SHA-256 mismatch") |
| checkpoint_dir = os.path.dirname(str(resume_checkpoint)) |
| required_full_state = [ |
| os.path.join(checkpoint_dir, "optim.pt"), |
| *[ |
| os.path.join( |
| checkpoint_dir, f"rng_rank{rank:02d}.pt") |
| for rank in range(4) |
| ], |
| ] |
| missing = [ |
| path for path in required_full_state |
| if not os.path.isfile(path) |
| ] |
| if missing: |
| raise FileNotFoundError( |
| "fixed-only exact resume requires optimizer and all " |
| f"four RNG states: {missing}") |
| 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=True) |
| 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 Wan KV " |
| "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): |
| if not fixed_openvid_dual: |
| expected_generator_sha256 = _strict_sha256( |
| config.generator_ckpt_sha256, "generator_ckpt_sha256") |
| actual_generator_sha256 = ( |
| self.model._distributed_file_sha256( |
| str(config.generator_ckpt))) |
| if actual_generator_sha256 != expected_generator_sha256: |
| raise ValueError( |
| "generator initialization SHA-256 mismatch") |
| if fixed_openvid_dual: |
| expected_generator_sha256 = str(getattr( |
| config, "generator_ckpt_sha256", "")).strip().lower() |
| if ( |
| len(expected_generator_sha256) != 64 |
| or any( |
| character not in "0123456789abcdef" |
| for character in expected_generator_sha256 |
| ) |
| ): |
| raise ValueError( |
| "fixed-only OpenVid requires generator_ckpt_sha256") |
| actual_generator_sha256 = ( |
| self.model._distributed_file_sha256( |
| str(config.generator_ckpt))) |
| if actual_generator_sha256 != expected_generator_sha256: |
| raise ValueError( |
| "selected 14B ODE generator checkpoint SHA-256 " |
| "mismatch") |
| print(f"Loading pretrained generator from {config.generator_ckpt}") |
| if str(config.generator_ckpt).endswith(".safetensors"): |
| state_dict = load_safetensors_file( |
| str(config.generator_ckpt), device="cpu") |
| else: |
| state_dict = torch.load(config.generator_ckpt, map_location="cpu") |
| if fixed_openvid_dual and ( |
| not isinstance(state_dict, dict) |
| or state_dict.get("model_name") != "Wan2.1-T2V-14B" |
| or state_dict.get("ode_recipe") != "causvid_14b_16k" |
| or int(state_dict.get("seed", -1)) != 0 |
| or state_dict.get("model_name") |
| != str(config.generator_ckpt_model_name) |
| or state_dict.get("ode_recipe") |
| != str(config.generator_ckpt_ode_recipe) |
| or int(state_dict.get("step", -1)) |
| != int(config.generator_ckpt_step) |
| or state_dict.get("prompt_manifest_sha256") |
| != _strict_sha256( |
| config.generator_ckpt_prompt_manifest_sha256, |
| "generator_ckpt_prompt_manifest_sha256", |
| ) |
| ): |
| raise ValueError( |
| "fixed-only OpenVid initialization must be a selected " |
| "Wan-14B CausVid ODE checkpoint") |
| 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=True |
| ) |
| 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.optimizer_state_offload_report = None |
| if self.optimizer_state_cpu_offload_during_backward: |
| torch.cuda.synchronize(self.device) |
| self.optimizer_state_offload_report = move_adamw_moment_state_( |
| self.generator_optimizer, "cpu") |
| torch.cuda.empty_cache() |
| if self.is_main_process: |
| print( |
| "[memory] VACE optimizer-state CPU offload ready: " |
| f"{self.optimizer_state_offload_report}", |
| flush=True, |
| ) |
| self.current_learning_rate = self._set_generator_learning_rate( |
| self.step) |
|
|
| def _set_generator_learning_rate(self, update_index): |
| schedule = self.learning_rate_schedule |
| if schedule["name"] == "constant": |
| value = float(schedule["value"]) |
| else: |
| value = _cosine_learning_rate( |
| int(update_index), |
| start=float(schedule["start"]), |
| end=float(schedule["end"]), |
| total_updates=int(schedule["total_updates"]), |
| ) |
| for parameter_group in self.generator_optimizer.param_groups: |
| parameter_group["lr"] = value |
| self.current_learning_rate = value |
| return value |
|
|
| 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 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 |
| )): |
| unexpected = [ |
| key for key in ( |
| "current_teacher_features", |
| "current_teacher_video01", |
| "current_teacher_latent", |
| "ode_latent", |
| ) |
| if batch.get(key) is not None |
| ] |
| if unexpected: |
| raise ValueError( |
| "fixed-only MMD rejects current-teacher payloads: " |
| f"{unexpected}") |
| 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) |
| 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 _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(): |
| conditional_dict = self.model.text_encoder( |
| text_prompts=text_prompts) |
| 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 |
|
|
| features, replay_descriptor = ( |
| self.model.collect_rollout_features( |
| image_or_video_shape=image_or_video_shape, |
| conditional_dict=conditional_dict, |
| unconditional_dict=unconditional_dict, |
| initial_latent=None, |
| vjepa2_window_sample_keys=( |
| vjepa2_window_sample_keys), |
| require_current_teacher=bool( |
| gradient_selected |
| and getattr( |
| self.model, |
| "current_teacher_attraction_enabled", |
| False, |
| ) |
| ), |
| 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 |
|
|
| @staticmethod |
| def _slice_fresh_video_batch_mapping(mapping, start, stop, batch_size): |
| sliced = {} |
| for key, value in mapping.items(): |
| if ( |
| isinstance(value, torch.Tensor) |
| and value.ndim > 0 |
| and int(value.shape[0]) == batch_size |
| ): |
| sliced[key] = value[start:stop] |
| else: |
| sliced[key] = value |
| return sliced |
|
|
| def _aggregate_fresh_video_forward_entries( |
| self, |
| physical_entries, |
| *, |
| gradient_selected, |
| ): |
| if not physical_entries: |
| raise ValueError("logical fresh-video microbatch is empty") |
| first_features = physical_entries[0]["features"] |
| if isinstance(first_features, Mapping): |
| features = { |
| name: torch.cat([ |
| entry["features"][name] for entry in physical_entries |
| ], dim=0) |
| for name in first_features |
| } |
| else: |
| features = torch.cat([ |
| entry["features"] for entry in physical_entries |
| ], dim=0) |
| sample_id_parts = [ |
| entry.get("sample_ids") for entry in physical_entries |
| ] |
| if all(value is None for value in sample_id_parts): |
| sample_ids = None |
| elif any(value is None for value in sample_id_parts): |
| raise RuntimeError( |
| "sample identities are incomplete within a logical batch") |
| else: |
| sample_ids = tuple( |
| sample_id |
| for values in sample_id_parts |
| for sample_id in values |
| ) |
| if len(sample_ids) != self.fresh_video_logical_micro_batch_size: |
| raise RuntimeError( |
| "logical fresh-video sample identity count mismatch") |
| aggregate = { |
| "features": features, |
| "gradient_selected": bool(gradient_selected), |
| "sample_ids": sample_ids, |
| "physical_entries": tuple(physical_entries), |
| } |
| teacher_values = [ |
| entry.get("current_teacher_features") |
| for entry in physical_entries |
| ] |
| if any(value is not None for value in teacher_values): |
| if not all(value is not None for value in teacher_values): |
| raise RuntimeError( |
| "matched teacher features are incomplete within a " |
| "logical microbatch") |
| if isinstance(teacher_values[0], Mapping): |
| aggregate["current_teacher_features"] = { |
| name: torch.cat([ |
| value[name] for value in teacher_values |
| ], dim=0) |
| for name in teacher_values[0] |
| } |
| else: |
| aggregate["current_teacher_features"] = torch.cat( |
| teacher_values, dim=0) |
| for entry in physical_entries: |
| entry.pop("features", None) |
| entry.pop("current_teacher_features", None) |
| return aggregate |
|
|
| def _split_fresh_video_replay_units( |
| self, |
| logical_entry, |
| logical_cotangent, |
| ): |
| logical_features = logical_entry["features"] |
| units = [] |
| logical_offset = 0 |
| replay_batch = self.fresh_video_physical_replay_batch_size |
| for physical_entry in logical_entry["physical_entries"]: |
| descriptor = physical_entry["replay_descriptor"] |
| physical_batch = int(descriptor["latent_shape"][0]) |
| for start in range(0, physical_batch, replay_batch): |
| stop = start + replay_batch |
| replay_descriptor = dict(descriptor) |
| replay_descriptor["latent_shape"] = ( |
| replay_batch, *tuple(descriptor["latent_shape"][1:])) |
| replay_descriptor["student_latent"] = descriptor[ |
| "student_latent"][start:stop] |
| generation_descriptor = dict( |
| descriptor["generation_descriptor"]) |
| generation_descriptor["noise"] = generation_descriptor[ |
| "noise"][start:stop] |
| generation_descriptor["trajectory_seeds"] = ( |
| generation_descriptor["trajectory_seeds"][start:stop]) |
| replay_descriptor["generation_descriptor"] = ( |
| generation_descriptor) |
| for key, value in tuple(replay_descriptor.items()): |
| if key in { |
| "rng_state", "latent_shape", "generation_descriptor", |
| "student_latent", |
| }: |
| continue |
| if ( |
| isinstance(value, torch.Tensor) |
| and value.ndim > 0 |
| and int(value.shape[0]) == physical_batch |
| ): |
| replay_descriptor[key] = value[start:stop] |
| unit_slice = slice( |
| logical_offset + start, logical_offset + stop) |
| unit_features = ( |
| { |
| name: value[unit_slice] |
| for name, value in logical_features.items() |
| } |
| if isinstance(logical_features, Mapping) |
| else logical_features[unit_slice] |
| ) |
| unit_cotangent = ( |
| { |
| name: value[unit_slice] |
| for name, value in logical_cotangent.items() |
| } |
| if isinstance(logical_cotangent, Mapping) |
| else logical_cotangent[unit_slice] |
| ) |
| units.append({ |
| "replay_descriptor": replay_descriptor, |
| "features": unit_features, |
| "cotangent": unit_cotangent, |
| "image_or_video_shape": [ |
| replay_batch, |
| *physical_entry["image_or_video_shape"][1:], |
| ], |
| "conditional_dict": ( |
| self._slice_fresh_video_batch_mapping( |
| physical_entry["conditional_dict"], |
| start, stop, physical_batch)), |
| "unconditional_dict": ( |
| self._slice_fresh_video_batch_mapping( |
| physical_entry["unconditional_dict"], |
| start, stop, physical_batch)), |
| }) |
| logical_offset += physical_batch |
| if logical_offset != self.fresh_video_logical_micro_batch_size: |
| raise RuntimeError( |
| "logical replay partition does not cover its full batch") |
| return units |
|
|
| 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: |
| for physical_entry in entry["physical_entries"]: |
| indices = physical_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): |
| expected_forward = self.fresh_video_total_samples |
| expected_gradient = self.fresh_video_generator_grad_global_samples |
| expected_teacher = ( |
| expected_gradient |
| if bool(getattr( |
| self.model, |
| "current_teacher_attraction_enabled", |
| False, |
| )) |
| else 0 |
| ) |
| full_student_mode = self.fresh_video_full_student_partial_gradient |
| expected_kss = expected_forward * ( |
| expected_forward if full_student_mode else expected_gradient) |
| expected_fixed = expected_forward if full_student_mode else expected_gradient |
| expected_live = expected_fixed * expected_teacher |
| expected_teacher_self = expected_teacher * expected_teacher |
| metadata = result.metadata |
| branch_metadata = metadata.get("branches") |
| if isinstance(branch_metadata, Mapping): |
| registry = getattr( |
| self.model, "dual_branch_registry", None) |
| if isinstance(registry, Mapping) and registry: |
| expected_branches = set(registry) |
| else: |
| feature_type = str(getattr( |
| self.model, "feature_type", "")) |
| expected_branches = ( |
| {"dinov3", "vjepa2", "videomae_v2"} |
| if feature_type |
| == "triple_dinov3_vjepa2_videomae_v2_dense_latent" |
| else ( |
| {"dinov3", "vjepa2"} |
| if feature_type |
| == "dual_dinov3_vjepa2_dense_latent" |
| else {"vjepa2", "videomae_v2"} |
| ) |
| ) |
| if set(branch_metadata) != expected_branches: |
| raise RuntimeError( |
| "multi-representation MMD metadata branch set mismatch: " |
| f"{set(branch_metadata)} != {expected_branches}") |
| denominator_keys = [ |
| "repulsion_denominator", |
| "fixed_teacher_attraction_denominator", |
| ] |
| expected_branch_values = [ |
| expected_kss, |
| expected_fixed, |
| ] |
| if expected_teacher: |
| denominator_keys.extend([ |
| "current_teacher_attraction_denominator", |
| "current_teacher_self_kernel_denominator", |
| ]) |
| expected_branch_values.extend([ |
| expected_live, |
| expected_teacher_self, |
| ]) |
| for name, branch in branch_metadata.items(): |
| observed_branch_values = tuple( |
| int(torch.as_tensor(branch[key]).item()) |
| for key in denominator_keys |
| ) |
| if observed_branch_values != tuple(expected_branch_values): |
| raise RuntimeError( |
| f"multi-representation branch {name} denominator mismatch: " |
| f"{observed_branch_values} != " |
| f"{tuple(expected_branch_values)}") |
| if int(torch.as_tensor( |
| branch.get("full_student_partial_gradient", 0) |
| ).item()) != 1: |
| raise RuntimeError( |
| f"multi-representation branch {name} omitted full-student injection") |
| metadata = {**metadata, **branch_metadata["vjepa2"]} |
| observed = ( |
| int(result.global_samples), |
| int(result.gradient_global_samples), |
| int(result.current_teacher_global_samples), |
| int(torch.as_tensor( |
| metadata["gradcache_kss_denominator"]).item()), |
| int(torch.as_tensor( |
| metadata["fixed_teacher_attraction_denominator"]).item()), |
| int(torch.as_tensor( |
| metadata.get("current_teacher_attraction_denominator", 0)).item()), |
| int(torch.as_tensor( |
| metadata.get("current_teacher_self_kernel_denominator", 0)).item()), |
| int(torch.as_tensor( |
| metadata["gradcache_full_student_partial_gradient"]).item()), |
| ) |
| expected = ( |
| expected_forward, |
| expected_gradient, |
| expected_teacher, |
| expected_kss, |
| expected_fixed, |
| expected_live, |
| expected_teacher_self, |
| int(full_student_mode), |
| ) |
| if observed != expected: |
| raise RuntimeError( |
| "partial-MMD geometry mismatch: " |
| f"observed={observed}, expected={expected}") |
| if int(torch.as_tensor( |
| metadata.get("repulsion_denominator", -1) |
| ).item()) != expected_kss: |
| raise RuntimeError("partial-MMD core Kss denominator mismatch") |
| if full_student_mode: |
| if int(torch.as_tensor( |
| metadata.get("full_student_partial_gradient", 0) |
| ).item()) != 1: |
| raise RuntimeError( |
| "loss did not activate full-student partial gradients") |
| configured_denominators = ( |
| str(getattr( |
| self.config, "partial_mmd_normalization_identity", "")), |
| int(getattr( |
| self.config, "partial_mmd_kss_denominator", -1)), |
| int(getattr( |
| self.config, |
| "partial_mmd_fixed_attraction_denominator", -1)), |
| int(getattr( |
| self.config, "partial_mmd_live_cross_denominator", -1)), |
| int(getattr( |
| self.config, |
| "partial_mmd_teacher_self_denominator", -1)), |
| ) |
| expected_config = ( |
| "full_student_population_denominators_v1", |
| expected_kss, |
| expected_fixed, |
| expected_live, |
| expected_teacher_self, |
| ) |
| if configured_denominators != expected_config: |
| raise RuntimeError( |
| "full-student MMD config normalization mismatch: " |
| f"{configured_denominators} != {expected_config}") |
| configured_compensation = getattr( |
| self.config, "partial_mmd_gradient_compensation", None) |
| if ( |
| configured_compensation is not None |
| and float(configured_compensation) |
| != self.fresh_video_generator_grad_sample_scale |
| ): |
| raise RuntimeError( |
| "full-student MMD gradient compensation mismatch: " |
| f"{configured_compensation} != " |
| f"{self.fresh_video_generator_grad_sample_scale}") |
| if bool(getattr( |
| self.model, "nystrom_repulsion_fifo_enabled", False |
| )): |
| raise RuntimeError("partial-MMD experiment requires FIFO=0") |
| if bool(getattr( |
| self.model, "nystrom_attraction_ema_enabled", False |
| )): |
| raise RuntimeError( |
| "partial-MMD experiment requires attraction EMA=0") |
| layers = tuple(getattr( |
| self.model, |
| "vjepa2_dense_layers", |
| (int(self.model.dense_layer),), |
| )) |
| if layers != (15,): |
| raise RuntimeError( |
| f"partial-MMD experiment requires only layer 15, got {layers}") |
| if self.is_main_process and int(self.step) == 0: |
| print( |
| "[partial-mmd-geometry] " |
| f"forward={expected_forward} " |
| f"gradient={expected_gradient} " |
| f"teacher={expected_teacher} " |
| f"Kss_denominator={expected_kss} " |
| f"fixed_attraction_denominator={expected_fixed} " |
| f"live_cross_denominator={expected_live} " |
| f"teacher_self_denominator={expected_teacher_self} " |
| f"full_student_partial_gradient={int(full_student_mode)} " |
| "gradient_compensation=" |
| f"{self.fresh_video_generator_grad_sample_scale:g} " |
| "fifo=0 attraction_ema=0 layer=15", |
| flush=True, |
| ) |
| if bool(getattr(self.model, "dual_video_mmd_enabled", False)): |
| registry = getattr(self.model, "dual_branch_registry", {}) |
| shapes = { |
| name: list(branch["sample_shape"]) |
| for name, branch in sorted(registry.items()) |
| } |
| expected_shapes = { |
| name: ([1280, 1024] if name == "videomae_v2" |
| else [1024, 1024]) |
| for name in registry |
| } |
| if shapes != expected_shapes: |
| raise RuntimeError( |
| f"multi-encoder branch sample shapes changed: {shapes}") |
| feature_type = str(getattr( |
| self.model, "feature_type", "")) |
| shape_tag = ( |
| "triple-sample-shapes" |
| if feature_type |
| == "triple_dinov3_vjepa2_videomae_v2_dense_latent" |
| else ( |
| "dinov3-vjepa2-sample-shapes" |
| if feature_type == "dual_dinov3_vjepa2_dense_latent" |
| else "dual-sample-shapes" |
| ) |
| ) |
| print( |
| f"[{shape_tag}] " |
| + " ".join( |
| f"{name}={shape[0]}x{shape[1]}" |
| for name, shape in sorted(shapes.items()) |
| ), |
| flush=True, |
| ) |
|
|
| 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._fresh_progress("feature_collection_start") |
| restore_error = None |
| try: |
| self.model.restore_frozen_modules_before_feature_replay() |
| except BaseException as exc: |
| restore_error = exc |
| if not self._all_ranks_true(restore_error is None): |
| message = ( |
| f": {type(restore_error).__name__}: {restore_error}" |
| if restore_error is not None |
| else " on another rank" |
| ) |
| raise RuntimeError( |
| "frozen replay module restore failed" + message |
| ) from restore_error |
| |
| 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: |
| physical_per_logical = ( |
| self.fresh_video_logical_micro_batch_size |
| // self.fresh_video_physical_forward_batch_size) |
| for microbatch_index in range(self.fresh_video_microbatches): |
| gradient_selected = bool( |
| microbatch_index >= ( |
| self.fresh_video_microbatches |
| - self.fresh_video_generator_grad_microbatches) |
| ) |
| physical_entries = [] |
| for physical_slot in range(physical_per_logical): |
| physical_index = ( |
| microbatch_index * physical_per_logical |
| + physical_slot) |
| window_sample_keys = None |
| if getattr( |
| self.model, |
| "vjepa2_dense_frame_sampling", |
| "center_contiguous", |
| ) == "random_contiguous_per_video": |
| local_start = ( |
| physical_index |
| * self.fresh_video_physical_forward_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( |
| self.fresh_video_physical_forward_batch_size)) |
| batch = ( |
| self._matched_batch_for_microbatch(physical_index) |
| if getattr( |
| self, "matched_counter_sampling_enabled", False) |
| else self._next_batch() |
| ) |
| physical_entries.append( |
| self._collect_rollout_microbatch( |
| batch, |
| gradient_selected=gradient_selected, |
| vjepa2_window_sample_keys=window_sample_keys, |
| ) |
| ) |
| entry = self._aggregate_fresh_video_forward_entries( |
| physical_entries, |
| gradient_selected=gradient_selected, |
| ) |
| 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: |
| for physical_entry in entry["physical_entries"]: |
| descriptor = physical_entry.get( |
| "replay_descriptor", {}) |
| descriptor.pop("cached_student_video01", None) |
| physical_entry.pop("conditional_dict", None) |
| physical_entry.pop("unconditional_dict", None) |
| physical_entry.pop("image_or_video_shape", None) |
| entries.append(entry) |
| |
| 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) |
| if ( |
| use_prepared_replay |
| and self.model.generator_replay_frozen_encoder_cpu_offload |
| ): |
| if not hasattr(cache, "offload_finalized_to_cpu"): |
| raise RuntimeError( |
| "prepared triple GradCache lacks CPU offload support") |
| result = cache.offload_finalized_to_cpu() |
| for entry in entries: |
| features = entry.pop("features") |
| if entry.get("gradient_selected", False): |
| if not isinstance(features, Mapping): |
| raise TypeError( |
| "triple replay features must be a mapping") |
| entry["features"] = { |
| name: value.detach().to( |
| device="cpu", non_blocking=False) |
| for name, value in features.items() |
| } |
| 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, |
| ) |
| replay_units = [] |
| for entry, cotangent in zip( |
| replay_entries, replay_video_cotangents |
| ): |
| replay_units.extend( |
| self._split_fresh_video_replay_units(entry, cotangent)) |
| if ( |
| len(replay_units) |
| * self.fresh_video_physical_replay_batch_size |
| != self.fresh_video_generator_grad_local_samples |
| ): |
| raise RuntimeError( |
| "physical replay units do not cover exactly the configured " |
| "gradient suffix") |
| prepared_replays = [None] * len(replay_units) |
| self._fresh_progress("encoder_vae_replay_start", units=len(replay_units)) |
| if use_prepared_replay: |
| prepared_replays = [ |
| self.model.prepare_latent_gradients( |
| replay_descriptor=unit["replay_descriptor"], |
| expected_features=unit["features"], |
| feature_cotangent=unit["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 unit in replay_units |
| ] |
|
|
| replay_offload_bytes = 0 |
| self._fresh_progress("encoder_vae_replay_complete") |
| if use_prepared_replay: |
| offload_error = None |
| try: |
| replay_offload_bytes = ( |
| self.model.offload_frozen_modules_for_generator_replay() |
| ) |
| except BaseException as exc: |
| offload_error = exc |
| if not self._all_ranks_true(offload_error is None): |
| message = ( |
| f": {type(offload_error).__name__}: {offload_error}" |
| if offload_error is not None |
| else " on another rank" |
| ) |
| raise RuntimeError( |
| "frozen replay module offload failed" + message |
| ) from offload_error |
|
|
| replay_logs = [] |
| self._fresh_progress( |
| "generator_backward_start", units=len(replay_units)) |
| for unit, prepared_replay in zip(replay_units, prepared_replays): |
| surrogate, replay_log = self.model.replay_generator_gradients( |
| replay_descriptor=unit["replay_descriptor"], |
| expected_features=unit["features"], |
| feature_cotangent=unit["cotangent"], |
| image_or_video_shape=unit["image_or_video_shape"], |
| conditional_dict=unit["conditional_dict"], |
| unconditional_dict=unit["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) |
| if ( |
| len(replay_logs) % 4 == 0 |
| or len(replay_logs) == len(replay_units) |
| ): |
| self._fresh_progress( |
| "generator_backward_progress", |
| completed=len(replay_logs), |
| total=len(replay_units), |
| ) |
|
|
| 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("physical_forward_calls"): result.loss.new_tensor( |
| len(entries) * physical_per_logical, dtype=torch.long), |
| dense_key("physical_replay_calls"): result.loss.new_tensor( |
| len(replay_units), 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, |
| ) |
| ), |
| "fresh_video_logical_micro_batch_size": ( |
| result.loss.new_tensor( |
| self.fresh_video_logical_micro_batch_size, |
| dtype=torch.long, |
| ) |
| ), |
| "fresh_video_physical_forward_batch_size": ( |
| result.loss.new_tensor( |
| self.fresh_video_physical_forward_batch_size, |
| dtype=torch.long, |
| ) |
| ), |
| "fresh_video_physical_replay_batch_size": ( |
| result.loss.new_tensor( |
| self.fresh_video_physical_replay_batch_size, |
| dtype=torch.long, |
| ) |
| ), |
| "generator_replay_frozen_module_offload_bytes": ( |
| result.loss.new_tensor( |
| replay_offload_bytes, dtype=torch.long) |
| ), |
| }) |
| if isinstance(meta.get("branches"), Mapping): |
| for branch_name, branch in meta["branches"].items(): |
| prefix = f"{branch_name}_dense" |
| branch_logs = { |
| 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.get( |
| "fixed_teacher_attraction", |
| branch["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"]), |
| } |
| if "current_teacher_attraction" in branch: |
| branch_logs.update({ |
| f"{prefix}_current_teacher_attraction": ( |
| branch["current_teacher_attraction"]), |
| f"{prefix}_teacher_cross_denominator": ( |
| branch[ |
| "current_teacher_attraction_denominator"]), |
| f"{prefix}_teacher_self_denominator": ( |
| branch[ |
| "current_teacher_self_kernel_denominator"]), |
| }) |
| log_dict.update(branch_logs) |
| log_dict["dual_dense_normalized_mmd_sum"] = meta[ |
| "normalized_mmd_sum"] |
| teacher_cache_logs = [ |
| physical_entry["current_teacher_cache_log"] |
| for entry in entries |
| for physical_entry in entry["physical_entries"] |
| if "current_teacher_cache_log" in physical_entry |
| ] |
| if teacher_cache_logs: |
| cache_stats = torch.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, device=torch.device( |
| "cuda", |
| self.device, |
| )) |
| 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: |
| for physical_entry in entry.get("physical_entries", ()): |
| descriptor = physical_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 _fresh_progress(self, phase, **fields): |
| if self.is_main_process and bool(getattr(self.config, "fresh_video_progress_logs", False)): |
| suffix = " ".join(f"{key}={value}" for key, value in fields.items()) |
| print(f"[fresh-progress] update={self.step + 1} phase={phase} {suffix}", flush=True) |
|
|
| 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: |
| current_learning_rate = self._set_generator_learning_rate(self.step) |
| 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: |
| if self.optimizer_state_cpu_offload_during_backward: |
| move_adamw_moment_state_( |
| self.generator_optimizer, self.device) |
| torch.cuda.synchronize(self.device) |
| |
| self.generator_optimizer.step() |
| if self.optimizer_state_cpu_offload_during_backward: |
| torch.cuda.synchronize(self.device) |
| self.optimizer_state_offload_report = ( |
| move_adamw_moment_state_( |
| self.generator_optimizer, "cpu")) |
| torch.cuda.empty_cache() |
| 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 |
| if self.is_main_process and bool(getattr(self.config, "fresh_video_progress_logs", False)): |
| print(f"[update-committed] step={self.step} committed={int(update_committed)}", flush=True) |
|
|
| |
| 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) |
|
|
| |
| checkpoint_due = self.step % self.config.log_iters == 0 |
| if (not self.config.no_save) and (self.step - start_step) > 0 and checkpoint_due: |
| if not update_committed: |
| raise RuntimeError("refusing to checkpoint an uncommitted optimizer update") |
| 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() |
| } |
| ) |
| 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() |
| wandb_loss_dict["learning_rate"] = current_learning_rate |
|
|
| 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 |
|
|
