Add training and inference source code

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.pytest_cache/
+*.zip

1.3B/.gitignore

@@ -0,0 +1,11 @@
+configs/local.yaml
+configs/local.env
+configs/path-remapping.json
+runs/
+outputs/
+cache/
+assets/
+.venv*/
+__pycache__/
+.pytest_cache/
+*.pyc

1.3B/README.md

@@ -0,0 +1,183 @@
+# Anonymous Video Generation — 1.3B
+
+[Project overview](../README.md) · [14B release](../14B/README.md)
+
+---
+
+For a guided implementation walkthrough, see [Reading the code](../docs/paper-to-code.md).
+The YAML uses the paper's distribution batch, differentiation batch, and reference minibatch terminology; `code/utils/config.py` keeps legacy configurations compatible.
+
+Anonymous Video Generation post-trains few-step autoregressive video generators by matching generated and reference videos in frozen representation spaces. It combines a hybrid Nyström–Monte Carlo distribution estimator with gradient replay, allowing distribution estimation to use more samples than backpropagation. Only the generator is optimized; diffusion score models are not needed during post-training.
+
+This repository provides the **Wan2.1-T2V-1.3B three-encoder configuration**, using **DINOv3, V-JEPA 2, and VideoMAEv2**. This configuration achieves **84.25 VBench Total** in the paper.
+
+## Supported Configuration
+
+| Backbone | Feature encoders | Sampling | Training GPUs | Updates |
+|---|---|---|---|---:|
+| Wan2.1-T2V-1.3B | DINOv3 + V-JEPA 2 + VideoMAEv2 | Causal, 4 steps per chunk | 4 | 150 |
+
+Training uses a distribution batch of **256 generated videos**, a differentiation batch of **128 generated videos**, and a reference minibatch of **128 videos**. Inference produces **81-frame, 832 × 480 videos at 16 FPS**.
+
+## 1. Environment
+
+Use Linux, Python 3.10, and NVIDIA GPUs with BF16 support. Run the commands below from the `1.3B/` directory.
+
+```bash
+conda create -n video-generation python=3.10 -y
+conda activate video-generation
+
+python -m pip install torch==2.5.1 torchvision==0.20.1 \
+  --index-url https://download.pytorch.org/whl/cu121
+python -m pip install -r requirements.txt
+python -m pip install flash-attn==2.7.4.post1 --no-build-isolation
+
+export PYTHON_BIN="$(command -v python)"
+```
+
+## 2. Models and Data
+
+### Inference assets
+
+Inference needs the Wan2.1 base model, text encoder, tokenizer, VAE, and a trained Anonymous Video Generation checkpoint. Place the matching Wan2.1 assets in the local model directory:
+
+```bash
+export ASSET_ROOT=/path/to/video-generation-assets
+
+# Place the matching model files at the local paths described below.
+```
+
+The base model is stored under `$ASSET_ROOT/models/Wan2.1-T2V-1.3B/`. The trained generator checkpoint can be stored separately and selected with `CHECKPOINT` when running inference.
+
+### Training assets
+
+Training starts from the Self-Forcing ODE initialization, stored locally at `$ASSET_ROOT/checkpoints/ode_init.pt`.
+
+It additionally requires:
+
+- **Frozen feature encoders:** DINOv3 ViT-L/16, V-JEPA 2 ViT-L, and VideoMAEv2-Large.
+- **Reference features:** precomputed features for the 8,682-video reference collection, with the aligned prompt/video manifest.
+- **Nyström references:** one fixed summary per encoder, built from the same collection.
+
+Reference preparation consists of extracting the three representations and building their Nyström summaries. During training, reference features are read directly from disk. See [Asset Preparation](docs/assets.md) for model sources, the directory layout, and preprocessing tools.
+
+**The trained Anonymous Video Generation weights and reference collection are not distributed with this source release.** Reproducing the reported training run requires the corresponding reference assets in addition to the public model downloads.
+
+### Configure asset paths
+
+Generate the local configuration once the asset root is set:
+
+```bash
+"$PYTHON_BIN" scripts/prepare_config.py --asset-root "$ASSET_ROOT"
+```
+
+This creates `configs/local.yaml` from [the three-encoder configuration](configs/three_encoder.yaml). Local asset paths are excluded from version control. Download dependencies before launching a run; the training and inference scripts use local models in offline mode.
+
+## 3. Quick Inference
+
+Create a text file with one prompt per line, then select a trained checkpoint:
+
+```bash
+export CHECKPOINT=/path/to/video-generation/model.pt
+
+GPUS=0 \
+  PROMPTS=/path/to/prompts.txt \
+  OUTPUT_DIR=outputs/demo \
+  bash scripts/infer.sh
+```
+
+The launcher generates one video per prompt with seed 0, using the generator weights without EMA. The four-step inference path has no additional classifier-free guidance branch. Example prompts are provided in [test_8.txt](code/prompts/test_8.txt).
+
+## 4. Training
+
+Prepare the environment and training assets above, then start the three-encoder run on four GPUs:
+
+```bash
+export LOGDIR=runs/video-generation
+GPUS=0,1,2,3 bash scripts/train.sh
+```
+
+The launcher uses FSDP and gradient checkpointing. Each worker generates 64 videos for distribution estimation and replays 32 with gradients. The feature encoders remain frozen, but gradients pass through them to the generated videos. W&B logging is disabled by the launcher.
+
+Use a separate `LOGDIR` for each run. Training starts from the ODE initialization and runs for **150 updates**, saving model weights, optimizer state, and per-worker random states every five updates. All checkpoints are retained. The final generator checkpoint is:
+
+```text
+runs/video-generation/checkpoint_model_000150/model.pt
+```
+
+Monitor progress with `tail -f "$LOGDIR/train.log"`. To generate videos from the trained model:
+
+```bash
+export CHECKPOINT="$LOGDIR/checkpoint_model_000150/model.pt"
+GPUS=0 PROMPTS=/path/to/prompts.txt \
+  OUTPUT_DIR=outputs/video-generation bash scripts/infer.sh
+```
+
+<details>
+<summary><b>Training configuration</b></summary>
+
+| Setting | Value |
+|---|---|
+| Initialization | Self-Forcing ODE checkpoint |
+| Distribution batch | 256 generated videos |
+| Differentiation batch | 128 generated videos |
+| Reference minibatch | 128 videos, sampled with replacement |
+| Encoder weights | DINOv3 : V-JEPA 2 : VideoMAEv2 = 10 : 1 : 1 |
+| Loss aggregation | Sum of weighted, normalized branch losses divided by 3 |
+| Reference estimator | Nyström + Monte Carlo, weights 2 : 1 |
+| Optimizer | AdamW, betas (0, 0.999), weight decay 0 |
+| Gradient clipping | Maximum norm 10 |
+| Learning rate | Cosine decay from 2.5e-6 to 5e-7 over 125 updates, then constant |
+| Training length | 150 updates |
+| Checkpoint interval | Every 5 updates |
+
+The full configuration is in [configs/three_encoder.yaml](configs/three_encoder.yaml).
+
+</details>
+
+### Short training run
+
+To check the complete training and checkpoint path before a full run:
+
+```bash
+export LOGDIR=runs/video-generation-smoke
+GPUS=0,1,2,3 bash scripts/train.sh --smoke
+```
+
+This runs **one optimizer update with the same batch sizes and all three encoders**. The resulting checkpoint is `$LOGDIR/checkpoint_model_000001/model.pt` and can be loaded with the inference command above.
+
+<details>
+<summary><b>Running in tmux and CPU tests</b></summary>
+
+On a machine with four allocated GPUs, start a persistent session:
+
+```bash
+tmux new -s video-generation
+```
+
+Inside the session, activate the environment, set `PYTHON_BIN` and `LOGDIR`, and run the training command. Detach with `Ctrl-b`, then `d`; reconnect with `tmux attach -t video-generation`.
+
+Run the CPU numerical tests separately:
+
+```bash
+"$PYTHON_BIN" -m pip install pytest
+"$PYTHON_BIN" scripts/test_cpu.py
+```
+
+See the [validation record](docs/validation.md) for revision-specific test results and the GPU environment used.
+
+</details>
+
+## 5. Evaluation
+
+The paper reports the following result for the three-encoder generator after 150 updates:
+
+| Model | VBench Total ↑ | Quality ↑ | Semantic ↑ |
+|---|---:|---:|---:|
+| Anonymous Video Generation, 1.3B, three encoders | **84.25** | **85.06** | **80.99** |
+
+The per-dimension scores are available in the [reference summary](results/reference/three-encoder-summary.json). Use the upstream [VBench](https://github.com/Vchitect/VBench) evaluator to evaluate generated videos; VBench and its evaluation models are maintained separately from this training and inference release. A short training run verifies execution and does not reproduce the reported benchmark score.
+
+## Acknowledgments
+
+This implementation builds on [Self-Forcing](https://github.com/guandeh17/Self-Forcing) and [Wan2.1](https://github.com/Wan-Video/Wan2.1). We also thank the authors of [DINOv3](https://github.com/facebookresearch/dinov3), [V-JEPA 2](https://github.com/facebookresearch/vjepa2), [VideoMAEv2](https://github.com/OpenGVLab/VideoMAEv2), and [VBench](https://github.com/Vchitect/VBench) for making their work available. Bundled source retains its upstream license notices; model weights and datasets follow their respective upstream terms.

1.3B/code/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

1.3B/code/configs/default_config.yaml

@@ -0,0 +1,11 @@
+independent_first_frame: false
+warp_denoising_step: false
+weight_decay: 0.01
+same_step_across_blocks: true
+discriminator_lr_multiplier: 1.0
+last_step_only: false
+i2v: false
+num_training_frames: 21
+gc_interval: 100
+context_noise: 0
+causal: true

1.3B/code/demo_utils/memory.py

@@ -0,0 +1,135 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import torch
+
+
+cpu = torch.device('cpu')
+gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(module: torch.nn.Module, **kwargs):
+        original_class = module.__class__
+        module.__dict__['forge_backup_original_class'] = original_class
+
+        def hacked_get_attr(self, name: str):
+            if '_parameters' in self.__dict__:
+                _parameters = self.__dict__['_parameters']
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if '_buffers' in self.__dict__:
+                _buffers = self.__dict__['_buffers']
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
+            '__getattr__': hacked_get_attr,
+        })
+
+        return
+
+    @staticmethod
+    def _uninstall_module(module: torch.nn.Module):
+        if 'forge_backup_original_class' in module.__dict__:
+            module.__class__ = module.__dict__.pop('forge_backup_original_class')
+        return
+
+    @staticmethod
+    def install_model(model: torch.nn.Module, **kwargs):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(model: torch.nn.Module):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+    if hasattr(model, 'scale_shift_table'):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, 'weight'):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(device=None):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats['active_bytes.all.current']
+    bytes_reserved = memory_stats['reserved_bytes.all.current']
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024 ** 3)
+
+
+def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(*args):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f'Unloaded {m.__class__.__name__} as complete.')
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(model, target_device, unload=True):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
+
+    gpu_complete_modules.append(model)
+    return

1.3B/code/inference.py

@@ -0,0 +1,238 @@
+import argparse
+import torch
+import os
+from omegaconf import OmegaConf
+from utils.config import load_config
+from tqdm import tqdm
+from torchvision import transforms
+import imageio.v3 as iio  # torchvision.io.write_video was removed in torchvision 0.26
+from einops import rearrange
+
+
+def write_video(path: str, frames: torch.Tensor, fps: int = 16) -> None:
+    """Drop-in replacement for ``torchvision.io.write_video``.
+
+    Accepts a ``[T, H, W, C]`` tensor in ``[0, 255]`` (float or uint8) and
+    writes it to ``path`` as an H.264 MP4 via the imageio-ffmpeg plugin.
+    """
+    if frames.dtype != torch.uint8:
+        frames = frames.clamp(0, 255).to(torch.uint8)
+    iio.imwrite(
+        path,
+        frames.cpu().numpy(),
+        fps=fps,
+        codec="libx264",
+        plugin="FFMPEG",
+    )
+import torch.distributed as dist
+from torch.utils.data import DataLoader, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from pipeline import CausalInferencePipeline
+from utils.dataset import TextDataset, TextImagePairDataset
+from utils.misc import set_seed
+
+import demo_utils.memory as _demo_memory
+from demo_utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller
+from pipeline import causal_inference as _causal_inference_module
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_path", type=str, help="Path to the config file")
+parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")
+parser.add_argument("--data_path", type=str, help="Path to the dataset")
+parser.add_argument("--extended_prompt_path", type=str, help="Path to the extended prompt")
+parser.add_argument("--output_folder", type=str, help="Output folder")
+parser.add_argument("--num_output_frames", type=int, default=21,
+                    help="Number of overlap frames between sliding windows")
+parser.add_argument("--i2v", action="store_true", help="Whether to perform I2V (or T2V by default)")
+parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")
+parser.add_argument("--seed", type=int, default=0, help="Random seed")
+parser.add_argument("--num_samples", type=int, default=1, help="Number of samples to generate per prompt")
+parser.add_argument("--save_with_index", action="store_true",
+                    help="Whether to save the video using the index or prompt as the filename")
+args = parser.parse_args()
+
+# Initialize distributed inference
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend='nccl')
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    world_size = dist.get_world_size()
+    set_seed(args.seed + local_rank)
+else:
+    device = torch.device("cuda")
+    local_rank = 0
+    world_size = 1
+    set_seed(args.seed)
+
+# ``demo_utils.memory.gpu`` is computed at import time as
+# ``cuda:<current_device>`` and is also re-exported into
+# ``pipeline.causal_inference``.  In a distributed launch every rank imports
+# before ``set_device`` runs, so they all see ``cuda:0`` and end up loading
+# every model onto rank 0's GPU.  Patch both module-level bindings so the
+# downstream ``.to(device=gpu)`` calls land on the per-rank device.
+_demo_memory.gpu = device
+_causal_inference_module.gpu = device
+gpu = device
+
+print(f'[rank {local_rank}] Free VRAM on {device}: {get_cuda_free_memory_gb(gpu)} GB')
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+torch.set_grad_enabled(False)
+
+config = load_config(args.config_path, "configs/default_config.yaml")
+
+# Initialize pipeline
+if not hasattr(config, 'denoising_step_list'):
+    raise ValueError("This release supports few-step causal inference only")
+pipeline = CausalInferencePipeline(config, device=device)
+
+def _strip_wrapper_prefixes(state_dict):
+    """Strip FSDP / activation-checkpointing / torch.compile wrapper prefixes
+    that may be embedded in saved state_dict keys.  ``EMA_FSDP.state_dict()``
+    in particular keeps ``_fsdp_wrapped_module.`` prefixes."""
+    out = {}
+    for k, v in state_dict.items():
+        new_k = (k.replace("_fsdp_wrapped_module.", "")
+                  .replace("_checkpoint_wrapped_module.", "")
+                  .replace("_orig_mod.", ""))
+        out[new_k] = v
+    return out
+
+
+if args.checkpoint_path:
+    raw_state_dict = torch.load(args.checkpoint_path, map_location="cpu", weights_only=False)
+    key = "generator_ema" if args.use_ema else "generator"
+    if key not in raw_state_dict:
+        available = list(raw_state_dict.keys())
+        raise KeyError(
+            f"Checkpoint {args.checkpoint_path!r} has no '{key}' entry "
+            f"(found: {available}).  Pass --use_ema only for DMD checkpoints saved "
+            f"after ema_start_step."
+        )
+    cleaned = _strip_wrapper_prefixes(raw_state_dict[key])
+    missing, unexpected = pipeline.generator.load_state_dict(cleaned, strict=False)
+    if local_rank == 0:
+        print(f"Loaded {args.checkpoint_path!r} (key='{key}', {len(cleaned)} tensors)")
+
+pipeline = pipeline.to(dtype=torch.bfloat16)
+if low_memory:
+    DynamicSwapInstaller.install_model(pipeline.text_encoder, device=gpu)
+else:
+    pipeline.text_encoder.to(device=gpu)
+pipeline.generator.to(device=gpu)
+pipeline.vae.to(device=gpu)
+
+
+# Create dataset
+if args.i2v:
+    assert not dist.is_initialized(), "I2V does not support distributed inference yet"
+    transform = transforms.Compose([
+        transforms.Resize((480, 832)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+    dataset = TextImagePairDataset(args.data_path, transform=transform)
+else:
+    dataset = TextDataset(prompt_path=args.data_path, extended_prompt_path=args.extended_prompt_path)
+num_prompts = len(dataset)
+print(f"Number of prompts: {num_prompts}")
+
+if dist.is_initialized():
+    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
+else:
+    sampler = SequentialSampler(dataset)
+dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+# Create output directory (only on main process to avoid race conditions)
+if local_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+if dist.is_initialized():
+    dist.barrier()
+
+
+def encode(self, videos: torch.Tensor) -> torch.Tensor:
+    device, dtype = videos[0].device, videos[0].dtype
+    scale = [self.mean.to(device=device, dtype=dtype),
+             1.0 / self.std.to(device=device, dtype=dtype)]
+    output = [
+        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+        for u in videos
+    ]
+
+    output = torch.stack(output, dim=0)
+    return output
+
+
+for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
+    idx = batch_data['idx'].item()
+
+    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
+    # Unpack the batch data for convenience
+    if isinstance(batch_data, dict):
+        batch = batch_data
+    elif isinstance(batch_data, list):
+        batch = batch_data[0]  # First (and only) item in the batch
+
+    all_video = []
+    num_generated_frames = 0  # Number of generated (latent) frames
+
+    if args.i2v:
+        # For image-to-video, batch contains image and caption
+        prompt = batch['prompts'][0]  # Get caption from batch
+        prompts = [prompt] * args.num_samples
+
+        # Process the image
+        image = batch['image'].squeeze(0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)
+
+        # Encode the input image as the first latent
+        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)
+        initial_latent = initial_latent.repeat(args.num_samples, 1, 1, 1, 1)
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames - 1, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+    else:
+        # For text-to-video, batch is just the text prompt
+        prompt = batch['prompts'][0]
+        extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None
+        if extended_prompt is not None:
+            prompts = [extended_prompt] * args.num_samples
+        else:
+            prompts = [prompt] * args.num_samples
+        initial_latent = None
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+
+    # Generate 81 frames
+    video, latents = pipeline.inference(
+        noise=sampled_noise,
+        text_prompts=prompts,
+        return_latents=True,
+        initial_latent=initial_latent,
+        low_memory=low_memory,
+    )
+    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
+    all_video.append(current_video)
+    num_generated_frames += latents.shape[1]
+
+    # Final output video
+    video = 255.0 * torch.cat(all_video, dim=1)
+
+    # Clear VAE cache
+    pipeline.vae.model.clear_cache()
+
+    # Save the video if the current prompt is not a dummy prompt
+    if idx < num_prompts:
+        model = "regular" if not args.use_ema else "ema"
+        for seed_idx in range(args.num_samples):
+            # All processes save their videos
+            if args.save_with_index:
+                output_path = os.path.join(args.output_folder, f'{idx}-{seed_idx}_{model}.mp4')
+            else:
+                output_path = os.path.join(args.output_folder, f'{prompt[:100]}-{seed_idx}.mp4')
+            write_video(output_path, video[seed_idx], fps=16)

1.3B/code/losses/__init__.py

@@ -0,0 +1 @@
+"""Distribution-matching losses used by distribution matching."""

1.3B/code/losses/multi_encoder_mmd.py

@@ -0,0 +1,235 @@
+"""Combine distribution-matching losses across frozen video encoders.
+
+Each encoder defines its own video feature space and Gaussian kernel. Compute
+one MMD per encoder, apply the release's detached per-encoder loss scaling,
+then apply the configured weights and aggregation divisor. ``VideoNystromMMD``
+implements generated repulsion and persistent/minibatch reference attraction.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import OrderedDict
+from typing import Any, Mapping, Optional, Sequence
+
+import torch
+from torch import nn
+
+from losses.video_nystrom_mmd import VideoNystromMMD
+
+
+def normalize_mmd_with_detached_scale(
+    raw_mmd: torch.Tensor,
+    normalization_eps: float = 1.0e-4,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return M / (stop_gradient(M) + eps) and its detached scale.
+
+    Detaching only the denominator preserves gradients through M while
+    balancing the different numerical scales of the encoder objectives.
+    This release-specific scaling is additional to the paper's MMD objective;
+    the expression deliberately has no absolute value or clamping.
+    """
+    denominator = raw_mmd.detach() + float(normalization_eps)
+    return raw_mmd / denominator, denominator
+
+
+class MultiEncoderMMD(nn.Module):
+    """Weighted distribution matching in multiple frozen feature spaces.
+
+    A branch is one encoder, its kernel bandwidth, and its persistent Nyström
+    reference. For raw branch losses M_e, the combined objective is
+    sum_e weight_e * M_e / (stop_gradient(M_e) + eps), divided by
+    ``aggregation_divisor``. That divisor is explicit; it need not equal the
+    number of branches or the sum of weights.
+    """
+
+    def __init__(
+        self,
+        branches: Mapping[str, VideoNystromMMD],
+        *,
+        weights: Optional[Mapping[str, float]] = None,
+        normalization_eps: float = 1.0e-4,
+        aggregation_divisor: float = 1.0,
+    ) -> None:
+        super().__init__()
+        ordered = OrderedDict()
+        for name, objective in branches.items():
+            ordered[name] = objective
+        self.branches = nn.ModuleDict(ordered)
+        configured_weights = weights or {name: 1.0 for name in ordered}
+        self.branch_weights = {}
+        for name in ordered:
+            weight = float(configured_weights[name])
+            self.branch_weights[name] = weight
+        self.normalization_eps = float(normalization_eps)
+        self.aggregation_divisor = float(aggregation_divisor)
+        self.current_teacher_attraction_enabled = True
+        self.attraction_ema_enabled = False
+        self.repulsion_fifo_enabled = False
+
+    @property
+    def branch_names(self) -> tuple[str, ...]:
+        return tuple(self.branches.keys())
+
+    @property
+    def sample_shapes(self) -> dict[str, tuple[int, int]]:
+        return {
+            name: objective.sample_shape
+            for name, objective in self.branches.items()
+        }
+
+    @property
+    def landmarks(self) -> torch.Tensor:
+        """Legacy diagnostics use the first branch; branch logs stay exact."""
+        return self.branches[self.branch_names[0]].landmarks
+
+    @property
+    def branch_registry(self) -> dict[str, dict[str, Any]]:
+        return {
+            name: {
+                "sample_shape": objective.sample_shape,
+                "reference_hash": objective.reference_hash,
+                "weight": self.branch_weights[name],
+            }
+            for name, objective in self.branches.items()
+        }
+
+    def _validate_mapping(
+        self,
+        value: Optional[Mapping[str, torch.Tensor]],
+        label: str,
+        *,
+        required: bool,
+    ) -> Optional[Mapping[str, torch.Tensor]]:
+        if value is None:
+            return None
+        return value
+
+    def forward(
+        self,
+        generated: Optional[Mapping[str, torch.Tensor]] = None,
+        current_teacher: Optional[Mapping[str, torch.Tensor]] = None,
+        *,
+        gradient_students: Optional[Mapping[str, torch.Tensor]] = None,
+        student_reference: Optional[Mapping[str, torch.Tensor]] = None,
+        gradient_student_indices: Optional[
+            torch.Tensor | Sequence[int]
+        ] = None,
+        stage: bool = True,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        """Evaluate each encoder loss, then balance and combine their scales.
+
+        ``gradient_students`` selects generated samples that retain gradients;
+        ``student_reference`` holds the full generated feature population.
+        ``current_teacher`` is the current frozen reference minibatch.
+        These input names are retained for existing callers and checkpoints.
+        """
+        # Select full-gradient or cached partial-gradient inputs. Every encoder
+        # receives the same sample selection.
+        partial = gradient_students is not None or student_reference is not None
+        generated = self._validate_mapping(
+            generated, "generated", required=not partial)
+        gradient_students = self._validate_mapping(
+            gradient_students, "gradient_students", required=partial)
+        student_reference = self._validate_mapping(
+            student_reference, "student_reference", required=partial)
+        current_teacher = self._validate_mapping(
+            current_teacher, "current_teacher", required=True)
+
+        # Each feature space has its own kernel objective; encoder features
+        # and persistent reference landmarks remain separate.
+        normalized_terms = []
+        branch_metadata: dict[str, dict[str, Any]] = {}
+        for name, objective in self.branches.items():
+            if partial:
+                raw, metadata = objective(
+                    gradient_students=gradient_students[name],
+                    student_reference=student_reference[name],
+                    current_teacher=current_teacher[name],
+                    gradient_student_indices=gradient_student_indices,
+                    stage=stage,
+                )
+            else:
+                raw, metadata = objective(
+                    generated[name],
+                    current_teacher=current_teacher[name],
+                    stage=stage,
+                )
+            # Balance the branch scale before applying its configured weight.
+            normalized, denominator = normalize_mmd_with_detached_scale(
+                raw, self.normalization_eps)
+            weighted = self.branch_weights[name] * normalized
+            normalized_terms.append(weighted)
+            branch_metadata[name] = {
+                **metadata,
+                "raw_mmd": raw.detach(),
+                "normalization_denominator": denominator.detach(),
+                "normalized_mmd": normalized.detach(),
+                "branch_weight": raw.new_tensor(self.branch_weights[name]),
+                "weighted_normalized_mmd": weighted.detach(),
+                "sample_shape": objective.sample_shape,
+                "reference_hash": objective.reference_hash,
+            }
+        # Preserve the explicit reduction order and configured divisor.
+        total = (
+            torch.stack(normalized_terms).sum() / self.aggregation_divisor
+        )
+        # Detached values below are diagnostics; ``total`` retains the graph.
+        first = branch_metadata[self.branch_names[0]]
+        metadata = {
+            "branches": branch_metadata,
+            "branch_names": self.branch_names,
+            "raw_mmd_sum": torch.stack([
+                branch_metadata[name]["raw_mmd"] for name in self.branch_names
+            ]).sum(),
+            "normalized_mmd_sum": total.detach(),
+            "normalization_eps": total.new_tensor(self.normalization_eps),
+            "normalization": "raw_div_stop_gradient_raw_plus_eps",
+            "aggregation": "weighted_sum_then_divide",
+            "aggregation_divisor": total.new_tensor(
+                self.aggregation_divisor),
+            # Shared batch counts are used by feature-gradient replay;
+            # per-encoder diagnostics remain available under ``branches``.
+            "global_samples": first["global_samples"],
+            "gradient_student_samples": first.get(
+                "gradient_student_samples", first["global_samples"]),
+            "repulsion_denominator": first.get("repulsion_denominator"),
+            "repulsion": torch.stack([
+                branch_metadata[name]["repulsion"]
+                for name in self.branch_names
+            ]).sum(),
+            "attraction": torch.stack([
+                branch_metadata[name]["attraction"]
+                for name in self.branch_names
+            ]).sum(),
+            "attraction_weight": total.new_ones(()),
+            "reference_constant": torch.stack([
+                branch_metadata[name]["reference_constant"]
+                for name in self.branch_names
+            ]).sum(),
+        }
+        return total, metadata
+
+    def prepare_commit(self) -> None:
+        for objective in self.branches.values():
+            objective.prepare_commit()
+
+    def commit(self) -> None:
+        for objective in self.branches.values():
+            objective.commit()
+
+    def discard(self) -> None:
+        for objective in self.branches.values():
+            objective.discard()
+
+
+# Historical imports and serialized class references use these aliases.
+MultiRepresentationVideoNystromMMD = MultiEncoderMMD
+fd_normalize_raw_mmd = normalize_mmd_with_detached_scale
+
+__all__ = [
+    "MultiEncoderMMD",
+    "normalize_mmd_with_detached_scale",
+    "MultiRepresentationVideoNystromMMD",
+    "fd_normalize_raw_mmd",
+]

1.3B/code/losses/multi_representation_video_nystrom_mmd.py

@@ -0,0 +1,7 @@
+"""Compatibility module; read ``losses.multi_encoder_mmd`` for the implementation."""
+
+import sys
+from losses import multi_encoder_mmd as _implementation
+
+# Reuse the module itself so existing imports and module-level patches agree.
+sys.modules[__name__] = _implementation

1.3B/code/losses/multilayer_video_nystrom_mmd.py

@@ -0,0 +1,381 @@
+"""Independent normalized Nyström MMD objectives over multiple video layers.
+
+The V-JEPA2 caller concatenates the sampled dense matrices from every selected
+layer along the token axis.  This module splits that transport tensor back into
+its layer-local matrices, evaluates one frozen-reference MMD per layer, applies
+detached scalar self-normalization independently, and sums the results.
+
+Keeping the objectives separate is important: concatenating layer features
+inside one kernel would change both the bandwidth and the teacher reference
+distribution instead of giving each layer equal normalized influence.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import torch
+from torch import nn
+
+from losses.video_nystrom_mmd import VideoNystromMMD
+
+
+class MultiLayerVideoNystromMMD(nn.Module):
+    """Sum independent, optionally self-normalized per-layer video MMDs."""
+
+    def __init__(
+        self,
+        objectives: Mapping[int, VideoNystromMMD],
+        *,
+        self_normalize: bool = True,
+        norm_eps: float = 1.0e-4,
+        layer_weights: Mapping[int, float] | None = None,
+    ) -> None:
+        super().__init__()
+        if not isinstance(objectives, Mapping) or not objectives:
+            raise ValueError("objectives must be a non-empty layer mapping")
+        ordered_layers = tuple(int(layer) for layer in objectives)
+        if len(set(ordered_layers)) != len(ordered_layers):
+            raise ValueError("objective layer indices must be unique")
+        if any(layer < 0 for layer in ordered_layers):
+            raise ValueError("objective layer indices must be non-negative")
+        if tuple(sorted(ordered_layers)) != ordered_layers:
+            raise ValueError("objective layer indices must be sorted")
+        if (
+            not math.isfinite(float(norm_eps))
+            or float(norm_eps) <= 0.0
+        ):
+            raise ValueError("norm_eps must be finite and positive")
+
+        weights = (
+            {layer: 1.0 for layer in ordered_layers}
+            if layer_weights is None
+            else {int(layer): float(value)
+                  for layer, value in layer_weights.items()}
+        )
+        if set(weights) != set(ordered_layers):
+            raise ValueError(
+                "layer_weights must contain exactly the objective layers")
+        if any(
+            not math.isfinite(value) or value < 0.0
+            for value in weights.values()
+        ):
+            raise ValueError(
+                "layer weights must be finite and non-negative")
+
+        modules = {}
+        sample_shapes = []
+        for layer in ordered_layers:
+            objective = objectives[layer]
+            if not isinstance(objective, VideoNystromMMD):
+                raise TypeError(
+                    "every layer objective must be VideoNystromMMD")
+            if any(
+                parameter.requires_grad
+                for parameter in objective.parameters()
+            ):
+                raise ValueError(
+                    "Nyström reference objectives must remain frozen")
+            modules[str(layer)] = objective
+            sample_shapes.append(objective.sample_shape)
+        feature_dims = {shape[1] for shape in sample_shapes}
+        if len(feature_dims) != 1:
+            raise ValueError(
+                "all layer objectives must share one feature dimension")
+
+        self.objectives = nn.ModuleDict(modules)
+        self.layer_indices = ordered_layers
+        self.layer_token_counts = tuple(
+            int(shape[0]) for shape in sample_shapes)
+        self.feature_dim = int(sample_shapes[0][1])
+        self.self_normalize = bool(self_normalize)
+        self.norm_eps = float(norm_eps)
+        self.layer_weights = tuple(
+            float(weights[layer]) for layer in ordered_layers)
+        ema_settings = {
+            (
+                objective.attraction_ema_enabled,
+                objective.attraction_ema_beta,
+            )
+            for objective in self.objectives.values()
+        }
+        if len(ema_settings) != 1:
+            raise ValueError(
+                "all layer objectives must use the same attraction EMA "
+                "enabled/beta setting")
+        (
+            self.attraction_ema_enabled,
+            self.attraction_ema_beta,
+        ) = next(iter(ema_settings))
+
+    @property
+    def sample_shape(self) -> tuple[int, int]:
+        """Shape of the concatenated GradCache transport matrix."""
+        return sum(self.layer_token_counts), self.feature_dim
+
+    @property
+    def landmarks(self) -> torch.Tensor:
+        """Compatibility device/landmark-count view used by GradCache logs."""
+        return self.objectives[str(self.layer_indices[0])].landmarks
+
+    @property
+    def reference_hashes(self) -> tuple[str, ...]:
+        return tuple(
+            str(self.objectives[str(layer)].reference_hash)
+            for layer in self.layer_indices
+        )
+
+    @property
+    def reference_metadata(self) -> tuple[Mapping[str, Any], ...]:
+        return tuple(
+            self.objectives[str(layer)].reference_metadata
+            for layer in self.layer_indices
+        )
+
+    @property
+    def attraction_ema_pending(self) -> bool:
+        pending = tuple(
+            objective.attraction_ema_pending
+            for objective in self.objectives.values()
+        )
+        if pending and any(value != pending[0] for value in pending[1:]):
+            raise RuntimeError(
+                "per-layer attraction EMA pending state is inconsistent")
+        return bool(pending[0]) if pending else False
+
+    def split_features(
+        self,
+        generated: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        if not isinstance(generated, torch.Tensor):
+            raise TypeError("generated must be a torch.Tensor")
+        if (
+            generated.ndim != 3
+            or min(generated.shape) <= 0
+            or not generated.is_floating_point()
+        ):
+            raise ValueError(
+                "generated must be a non-empty floating [B,N,D] tensor")
+        if not bool(torch.isfinite(generated.detach()).all().item()):
+            raise FloatingPointError(
+                "generated contains non-finite values")
+        expected_tokens, expected_dim = self.sample_shape
+        actual_shape = (
+            int(generated.shape[1]), int(generated.shape[2]))
+        if actual_shape != (expected_tokens, expected_dim):
+            raise ValueError(
+                "concatenated multi-layer video matrix shape does not match "
+                f"the frozen references: {actual_shape} != "
+                f"{(expected_tokens, expected_dim)}")
+        return tuple(torch.split(
+            generated, self.layer_token_counts, dim=1))
+
+    def forward(
+        self,
+        generated: torch.Tensor,
+        *,
+        stage: bool = True,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        layer_features = self.split_features(generated)
+        total = generated.new_zeros((), dtype=torch.float32)
+        raw_sum = total.clone()
+        repulsion_sum = total.clone()
+        attraction_sum = total.clone()
+        current_attraction_sum = total.clone()
+        reference_constant_sum = total.clone()
+        layer_metadata: dict[str, dict[str, torch.Tensor]] = {}
+        global_samples = None
+        ema_updates = None
+        attraction_current_coefficient = None
+
+        for layer, weight, features in zip(
+            self.layer_indices,
+            self.layer_weights,
+            layer_features,
+        ):
+            objective = self.objectives[str(layer)]
+            raw_mmd, metadata = (
+                objective(features, stage=stage)
+                if self.attraction_ema_enabled
+                else objective(features)
+            )
+            if self.self_normalize:
+                # Nyström attraction is approximate, so its finite estimate
+                # can be slightly negative even though exact squared MMD is
+                # non-negative. Normalize by detached magnitude to avoid both
+                # a singularity near -eps and an accidental gradient reversal.
+                denominator = raw_mmd.detach().abs() + self.norm_eps
+                if (
+                    not bool(torch.isfinite(denominator).item())
+                    or float(denominator) == 0.0
+                ):
+                    raise FloatingPointError(
+                        f"layer {layer} has an invalid MMD normalization "
+                        "denominator")
+                normalized_mmd = raw_mmd / denominator
+            else:
+                normalized_mmd = raw_mmd
+            weighted_mmd = float(weight) * normalized_mmd
+            total = total + weighted_mmd
+            raw_sum = raw_sum + raw_mmd.detach()
+            repulsion_sum = repulsion_sum + metadata["repulsion"]
+            attraction_sum = attraction_sum + metadata["attraction"]
+            if self.attraction_ema_enabled:
+                current_attraction_sum = (
+                    current_attraction_sum + metadata[
+                        "current_attraction"])
+            reference_constant_sum = (
+                reference_constant_sum + metadata["reference_constant"])
+            observed_global_samples = int(metadata["global_samples"])
+            if global_samples is None:
+                global_samples = observed_global_samples
+            elif observed_global_samples != global_samples:
+                raise RuntimeError(
+                    "per-layer MMD objectives disagree on global batch size")
+            if self.attraction_ema_enabled:
+                observed_updates = int(
+                    metadata["attraction_ema_updates"])
+                if ema_updates is None:
+                    ema_updates = observed_updates
+                elif observed_updates != ema_updates:
+                    raise RuntimeError(
+                        "per-layer attraction EMA update counts disagree")
+                observed_coefficient = float(
+                    metadata["attraction_current_coefficient"])
+                if attraction_current_coefficient is None:
+                    attraction_current_coefficient = observed_coefficient
+                elif observed_coefficient != attraction_current_coefficient:
+                    raise RuntimeError(
+                        "per-layer attraction EMA coefficients disagree")
+            layer_metadata[str(layer)] = {
+                **metadata,
+                "raw_mmd": raw_mmd.detach(),
+                "normalized_mmd": normalized_mmd.detach(),
+                "weighted_mmd": weighted_mmd.detach(),
+                "weight": raw_mmd.new_tensor(float(weight)),
+            }
+
+        if not bool(torch.isfinite(total.detach()).item()):
+            raise FloatingPointError(
+                "combined multi-layer MMD loss is non-finite")
+        combined_metadata = {
+            "layers": layer_metadata,
+            "layer_indices": torch.tensor(
+                self.layer_indices,
+                device=total.device,
+                dtype=torch.long,
+            ),
+            "num_layers": total.new_tensor(
+                len(self.layer_indices), dtype=torch.long),
+            "self_normalized": total.new_tensor(
+                int(self.self_normalize), dtype=torch.long),
+            "normalization_epsilon": total.new_tensor(self.norm_eps),
+            "raw_mmd_sum": raw_sum.detach(),
+            "normalized_mmd_sum": total.detach(),
+            # Aggregated raw-kernel terms preserve the existing generic log
+            # contract. Per-layer diagnostics above are authoritative.
+            "repulsion": repulsion_sum.detach(),
+            "attraction": attraction_sum.detach(),
+            "reference_constant": reference_constant_sum.detach(),
+            "global_samples": total.new_tensor(
+                int(global_samples), dtype=torch.long),
+        }
+        if self.attraction_ema_enabled:
+            combined_metadata.update({
+                "current_attraction": current_attraction_sum.detach(),
+                "attraction_ema_enabled": total.new_ones(
+                    (), dtype=torch.long),
+                "attraction_ema_updates": total.new_tensor(
+                    int(ema_updates), dtype=torch.long),
+                "attraction_current_coefficient": total.new_tensor(
+                    float(attraction_current_coefficient)),
+            })
+        return total, combined_metadata
+
+    def prepare_commit(self) -> None:
+        """Validate every layer before publishing any EMA scalar."""
+        if not self.attraction_ema_enabled:
+            return
+        for objective in self.objectives.values():
+            objective.prepare_commit()
+
+    def commit(self) -> None:
+        """Publish one attraction EMA update for every layer."""
+        if not self.attraction_ema_enabled:
+            return
+        if not self.attraction_ema_pending:
+            return
+        self.prepare_commit()
+        for objective in self.objectives.values():
+            objective.commit()
+
+    def discard(self) -> None:
+        for objective in self.objectives.values():
+            objective.discard()
+
+    def attraction_ema_state_dict(self) -> Mapping[str, Any] | None:
+        if not self.attraction_ema_enabled:
+            return None
+        if self.attraction_ema_pending:
+            raise RuntimeError(
+                "refusing to checkpoint pending per-layer attraction EMA")
+        return {
+            "version": 1,
+            "layers": list(self.layer_indices),
+            "states": {
+                str(layer): self.objectives[
+                    str(layer)].attraction_ema_state_dict()
+                for layer in self.layer_indices
+            },
+        }
+
+    def validate_attraction_ema_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> Mapping[str, Any]:
+        if not self.attraction_ema_enabled:
+            raise ValueError(
+                "cannot load attraction EMA state while the path is disabled")
+        if (
+            not isinstance(state, Mapping)
+            or set(state) != {"version", "layers", "states"}
+            or state.get("version") != 1
+            or tuple(state.get("layers", ())) != self.layer_indices
+            or not isinstance(state.get("states"), Mapping)
+            or set(state["states"])
+            != {str(layer) for layer in self.layer_indices}
+        ):
+            raise ValueError(
+                "multi-layer attraction EMA checkpoint schema mismatch")
+
+        # Validate every layer before the first live scalar is mutated.
+        validated = {}
+        for layer in self.layer_indices:
+            objective = self.objectives[str(layer)]
+            validated[str(layer)] = (
+                objective.validate_attraction_ema_state_dict(
+                    state["states"][str(layer)])
+            )
+        update_counts = {
+            int(layer_state["updates"].item())
+            for layer_state in validated.values()
+        }
+        if len(update_counts) != 1:
+            raise ValueError(
+                "multi-layer attraction EMA update counts disagree")
+        return {
+            "version": 1,
+            "layers": list(self.layer_indices),
+            "states": validated,
+        }
+
+    def load_attraction_ema_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> None:
+        validated = self.validate_attraction_ema_state_dict(state)
+        for layer in self.layer_indices:
+            self.objectives[str(layer)].load_attraction_ema_state_dict(
+                validated["states"][str(layer)])

1.3B/code/losses/video_nystrom_mmd.py

@@ -0,0 +1,1366 @@
+"""Single-encoder distribution matching: exact repulsion and reference attraction.
+
+Reading map: offline helpers build the persistent Nyström reference; training
+starts at ``VideoNystromMMD.forward``. Paper Sec. 3.2 calls its combination
+with the Monte Carlo reference minibatch the hybrid reference kernel mean.
+The released three-encoder path uses full-population partial gradients with
+FIFO and attraction EMA disabled; those optional modes are retained below.
+
+The released estimator includes diagonal pairs (B * B denominator) and the
+reference-reference constant. This differs from the paper's self-excluding
+repulsion and constant-free optimization expression; preserve both here,
+including their effect on the release's subsequent detached loss scaling.
+
+Every input sample is a complete ``[N, D]`` dense-video feature matrix.  For
+generated samples ``G``, frozen landmarks ``Z`` and a frozen reference pack,
+the legacy objective is
+
+    mean(K_gg) - 2 * mean(K_gz @ alpha) + k_rr.
+
+An opt-in matched-teacher path additionally accepts a fresh teacher batch
+``T``. It computes the complete biased Gaussian-kernel terms ``mean(K_gt)``
+and ``mean(K_tt)`` without sampling or approximation. The fixed-reference and
+fresh-teacher MMDs are combined with normalized non-negative group weights, so
+enabling a second teacher target does not duplicate the student repulsion or
+silently double the total attraction coefficient.
+
+The explicit partial-gradient interface uses gradient students ``G`` and a
+detached full student population ``S``. Without explicit student indices it
+preserves the rectangular estimator ``mean K(G,S)`` and evaluates attractions
+only for ``G``. Supplying the positions of ``G`` in ``S`` opts into the
+full-student estimator: selected differentiable values are reinserted into
+detached ``S``, then the loss uses ``mean K(S,S)``, fixed attraction over
+``S``, and matched-teacher attraction over ``S x T``. Thus only ``G`` gets
+gradients, including both sides of the full self-kernel. Both partial modes are defined without FIFO or attraction EMA, since those
+stateful transforms change their stated estimators.
+
+By default, ``K_gg`` is the exact, biased Gaussian-RBF Gram matrix (including
+its diagonal).  An optional transactional FIFO changes only this repulsion
+term to the current-current off-diagonal sum plus the current-history cross
+sum, normalized by the exact number of ordered pairs.  FIFO entries are
+detached FP32 feature matrices and are published only after optimizer success.
+Only the generated samples receive gradients.  The reference attraction is
+represented by a kernel-ridge/Nyström fit and ``k_rr`` is the exact biased
+reference-reference constant.
+
+An optional transactional path applies an EMA to the complete scalar
+attraction term only.  The exact current-batch repulsion is never averaged.
+The current attraction keeps its natural ``1-beta`` derivative after the
+first committed update; historical attraction values are detached.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import math
+import os
+import tempfile
+from collections.abc import Mapping, Sequence
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+from torch import nn
+
+
+REFERENCE_PACK_VERSION = 1
+_PACK_KEYS = {
+    "format_version",
+    "landmarks",
+    "alpha",
+    "k_rr",
+    "sigma",
+    "metadata",
+    "content_hash",
+}
+
+
+def canonicalize_metadata(metadata: Mapping[str, Any]) -> Dict[str, Any]:
+    """Return a deterministic, JSON-safe metadata mapping."""
+    try:
+        encoded = json.dumps(
+            dict(metadata),
+            sort_keys=True,
+            separators=(",", ":"),
+            ensure_ascii=True,
+            allow_nan=False,
+        )
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            "metadata must contain only finite JSON-serializable values"
+        ) from exc
+    decoded = json.loads(encoded)
+    return decoded
+
+
+def _canonical_json_sha256(value: Mapping[str, Any]) -> str:
+    payload = json.dumps(
+        dict(value),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+        allow_nan=False,
+    ).encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()
+
+
+def _canonicalize_shared_mount_aliases(value: Any) -> tuple[Any, bool]:
+    """Normalize deployment paths in DINOv3 contract metadata.
+
+    Only deployment paths are
+    normalized; encoder revision, weight digest, and every feature setting stay
+    in the comparison. The original reference pack and its content hash are
+    never rewritten.
+    """
+    if isinstance(value, Mapping):
+        original = dict(value)
+        normalized = {}
+        changed = False
+        for key, item in original.items():
+            normalized[key], item_changed = _canonicalize_shared_mount_aliases(item)
+            changed = changed or item_changed
+        if original.get("feature_type") == "dinov3_dense_latent" and (
+            "repo_path" in original or "weight_path" in original
+        ):
+            weight = str(original.get("weight_sha256", ""))
+            revision = str(original.get("source_commit", ""))
+            normalized["repo_path"] = "source:" + revision
+            normalized["weight_path"] = "sha256:" + weight
+            normalized.pop("feature_contract_sha256", None)
+            normalized["feature_contract_sha256"] = _canonical_json_sha256(normalized)
+            changed = True
+        return normalized, changed
+    if isinstance(value, list):
+        normalized = []
+        changed = False
+        for item in value:
+            result, item_changed = _canonicalize_shared_mount_aliases(item)
+            normalized.append(result)
+            changed = changed or item_changed
+        return normalized, changed
+    return value, False
+
+
+def _metadata_contains(actual: Any, expected: Any, path: str = "metadata") -> None:
+    return None
+
+
+def _validate_video_tensor(tensor: torch.Tensor, name: str) -> Tuple[int, int, int]:
+    return int(tensor.shape[0]), int(tensor.shape[1]), int(tensor.shape[2])
+
+
+def pairwise_frobenius_sq(
+    left: torch.Tensor,
+    right: torch.Tensor,
+) -> torch.Tensor:
+    """Squared Frobenius distances between ``[B,N,D]`` video matrices."""
+    _, left_n, left_d = _validate_video_tensor(left, "left")
+    _, right_n, right_d = _validate_video_tensor(right, "right")
+    left_flat = left.float().flatten(1)
+    right_flat = right.float().flatten(1)
+    # vector_norm performs the reduction without allocating an elementwise
+    # square tensor the size of a potentially 80-GiB FP32 reference.
+    left_sq = torch.linalg.vector_norm(
+        left_flat, dim=1, dtype=torch.float32).square().unsqueeze(1)
+    right_sq = torch.linalg.vector_norm(
+        right_flat, dim=1, dtype=torch.float32).square().unsqueeze(0)
+    # Scale only the small pairwise result; scaling left_flat first would allocate
+    # another full reference-sized tensor (about 80 GiB for the 20K bank).
+    cross = left_flat @ right_flat.t()
+    return (left_sq + right_sq - 2.0 * cross).clamp_min(0)
+
+
+def rbf_frobenius_kernel(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    sigma: torch.Tensor | float,
+) -> torch.Tensor:
+    """Gaussian RBF kernel using complete video-matrix samples."""
+    sigma_tensor = torch.as_tensor(
+        sigma, device=left.device, dtype=torch.float32)
+    distances_sq = pairwise_frobenius_sq(left, right)
+    return torch.exp(-distances_sq / (2.0 * sigma_tensor.square()))
+
+
+@torch.no_grad()
+def median_frobenius_sigma(
+    reference: torch.Tensor,
+    *,
+    max_samples: Optional[int] = 2048,
+    chunk_size: int = 256,
+    min_sigma: float = 1.0e-6,
+) -> torch.Tensor:
+    """Deterministic median-distance bandwidth over video samples.
+
+    A deterministic linspace subset bounds the number of pair distances.
+    Pair blocks are constructed in chunks; only the final one-dimensional
+    vector of upper-triangular distances is retained for the exact median of
+    that subset.
+    """
+    count, _, _ = _validate_video_tensor(reference, "reference")
+    samples = reference.detach().float()
+    if max_samples is not None:
+        if count > int(max_samples):
+            indices = torch.linspace(
+                0,
+                count - 1,
+                steps=int(max_samples),
+                device=samples.device,
+                dtype=torch.float64,
+            ).round().long()
+            samples = samples.index_select(0, indices)
+    count = int(samples.shape[0])
+    if count < 2:
+        return samples.new_tensor(1.0)
+
+    pieces = []
+    for row_start in range(0, count, chunk_size):
+        row_end = min(row_start + chunk_size, count)
+        rows = samples[row_start:row_end]
+        for col_start in range(row_start, count, chunk_size):
+            col_end = min(col_start + chunk_size, count)
+            block = pairwise_frobenius_sq(
+                rows, samples[col_start:col_end])
+            if col_start == row_start:
+                upper = torch.triu_indices(
+                    row_end - row_start,
+                    col_end - col_start,
+                    offset=1,
+                    device=block.device,
+                )
+                values = block[upper[0], upper[1]]
+            else:
+                values = block.reshape(-1)
+            positive = values[values > 0]
+            if positive.numel():
+                pieces.append(positive)
+    if not pieces:
+        return samples.new_tensor(1.0)
+    median_sq = torch.cat(pieces).median().clamp_min(float(min_sigma) ** 2)
+    return median_sq.sqrt()
+
+
+def _nearest_centroid_labels(
+    samples: torch.Tensor,
+    centroids: torch.Tensor,
+    chunk_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    labels = []
+    minimum_sq = []
+    for start in range(0, samples.shape[0], chunk_size):
+        distances = pairwise_frobenius_sq(
+            samples[start:start + chunk_size], centroids)
+        values, indices = distances.min(dim=1)
+        labels.append(indices)
+        minimum_sq.append(values)
+    return torch.cat(labels), torch.cat(minimum_sq)
+
+
+def _deterministic_cluster_sums(
+    flattened: torch.Tensor,
+    labels: torch.Tensor,
+    num_clusters: int,
+    *,
+    max_workspace_bytes: int = 1 << 30,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Segment-sum rows without CUDA atomics or an N-by-D temporary.
+
+    Rows are sorted once, while feature columns are processed in bounded
+    chunks. Each column retains the same fixed-order prefix sum as the legacy
+    implementation, but joint video/text references no longer require two
+    additional full-reference CUDA tensors.
+    """
+    counts = torch.bincount(labels, minlength=int(num_clusters))
+    order = torch.argsort(labels, stable=True)
+    sums = flattened.new_zeros(
+        (int(num_clusters), flattened.shape[1]))
+    nonempty = counts > 0
+    ends = counts.cumsum(0) - 1
+    selected_ends = ends[nonempty]
+    starts = selected_ends - counts[nonempty] + 1
+    has_predecessor = starts > 0
+    # ``index_select`` and ``cumsum`` coexist, so budget for both tensors.
+    bytes_per_column = max(1, 2 * flattened.shape[0] * flattened.element_size())
+    feature_chunk = max(
+        1, min(flattened.shape[1], int(max_workspace_bytes) // bytes_per_column))
+    for begin in range(0, flattened.shape[1], feature_chunk):
+        end = min(begin + feature_chunk, flattened.shape[1])
+        sorted_values = flattened[:, begin:end].index_select(0, order)
+        prefix = sorted_values.cumsum(dim=0)
+        cluster_sums = prefix.index_select(0, selected_ends)
+        if bool(has_predecessor.any()):
+            cluster_sums[has_predecessor] -= prefix.index_select(
+                0, starts[has_predecessor] - 1)
+        sums[nonempty, begin:end] = cluster_sums
+    return sums, counts
+
+
+@torch.no_grad()
+def deterministic_video_kmeans(
+    reference: torch.Tensor,
+    num_landmarks: int,
+    *,
+    iterations: int = 20,
+    chunk_size: int = 256,
+) -> torch.Tensor:
+    """Deterministic farthest-first initialized Lloyd k-means."""
+    count, num_tokens, feature_dim = _validate_video_tensor(
+        reference, "reference")
+    num_landmarks = int(num_landmarks)
+
+    samples = reference.detach().float()
+    global_mean = samples.mean(dim=0, keepdim=True)
+    first = int(
+        pairwise_frobenius_sq(samples, global_mean)
+        .squeeze(1)
+        .argmax()
+        .item()
+    )
+    selected = [first]
+    closest_sq = pairwise_frobenius_sq(
+        samples, samples[first:first + 1]).squeeze(1)
+    for _ in range(1, num_landmarks):
+        next_index = int(closest_sq.argmax().item())
+        selected.append(next_index)
+        next_sq = pairwise_frobenius_sq(
+            samples, samples[next_index:next_index + 1]).squeeze(1)
+        closest_sq = torch.minimum(closest_sq, next_sq)
+
+    centroids = samples.index_select(
+        0, torch.tensor(selected, device=samples.device)).clone()
+    previous_labels = None
+    for _ in range(int(iterations)):
+        labels, _ = _nearest_centroid_labels(
+            samples, centroids, int(chunk_size))
+        if previous_labels is not None and torch.equal(labels, previous_labels):
+            break
+        previous_labels = labels
+        flattened = samples.flatten(1)
+        sums, cluster_counts = _deterministic_cluster_sums(
+            flattened, labels, num_landmarks)
+        nonempty = cluster_counts > 0
+        updated = centroids.flatten(1).clone()
+        updated[nonempty] = (
+            sums[nonempty]
+            / cluster_counts[nonempty, None].to(sums.dtype)
+        )
+        centroids = updated.reshape(
+            num_landmarks, num_tokens, feature_dim)
+    return centroids
+
+
+@torch.no_grad()
+def compute_nystrom_alpha(
+    reference: torch.Tensor,
+    landmarks: torch.Tensor,
+    sigma: torch.Tensor | float,
+    *,
+    ridge: float = 1.0e-6,
+    chunk_size: int = 256,
+) -> torch.Tensor:
+    """Fit coefficients for the persistent Nyström reference kernel mean.
+
+    Solve (K_ZZ + ridge * I) alpha = mean_R K_ZR once, offline. Training
+    evaluates K(x, Z) @ alpha without revisiting the full reference bank R.
+    Code ``alpha`` is the paper's coefficient vector beta (Appendix C.2),
+    not the scalar mixture weight alpha_e in the hybrid mean (Eq. 9).
+    """
+    ref_count, ref_n, ref_d = _validate_video_tensor(reference, "reference")
+    _, landmark_n, landmark_d = _validate_video_tensor(
+        landmarks, "landmarks")
+
+    reference = reference.detach().float()
+    landmarks = landmarks.detach().float()
+    target = landmarks.new_zeros(landmarks.shape[0], dtype=torch.float64)
+    for start in range(0, ref_count, int(chunk_size)):
+        block = rbf_frobenius_kernel(
+            landmarks, reference[start:start + int(chunk_size)], sigma)
+        target.add_(block.double().sum(dim=1))
+    target.div_(float(ref_count))
+
+    kernel_zz = rbf_frobenius_kernel(
+        landmarks, landmarks, sigma).double()
+    kernel_zz = 0.5 * (kernel_zz + kernel_zz.t())
+    kernel_zz.fill_diagonal_(1.0)
+    regularized = kernel_zz + float(ridge) * torch.eye(
+        landmarks.shape[0],
+        device=landmarks.device,
+        dtype=torch.float64,
+    )
+    try:
+        alpha = torch.linalg.solve(regularized, target)
+    except RuntimeError as exc:
+        raise ValueError(
+            "Nyström landmark kernel is singular; increase ridge") from exc
+    return alpha.float()
+
+
+@torch.no_grad()
+def compute_reference_krr(
+    reference: torch.Tensor,
+    sigma: torch.Tensor | float,
+    *,
+    chunk_size: int = 256,
+) -> torch.Tensor:
+    """Exact biased ``mean(K_rr)``, including the reference diagonal."""
+    count, _, _ = _validate_video_tensor(reference, "reference")
+    reference = reference.detach().float()
+    total = reference.new_zeros((), dtype=torch.float64)
+    for row_start in range(0, count, int(chunk_size)):
+        rows = reference[row_start:row_start + int(chunk_size)]
+        for col_start in range(0, count, int(chunk_size)):
+            cols = reference[col_start:col_start + int(chunk_size)]
+            kernel = rbf_frobenius_kernel(rows, cols, sigma)
+            block_sum = kernel.double().sum()
+            if row_start == col_start:
+                block_sum.add_(
+                    (1.0 - kernel.diagonal()).double().sum())
+            total.add_(block_sum)
+    return (total / float(count * count)).float()
+
+
+def _hash_tensor(hasher: Any, name: str, tensor: torch.Tensor) -> None:
+    value = tensor.detach().cpu().contiguous()
+    descriptor = {
+        "name": name,
+        "dtype": str(value.dtype),
+        "shape": list(value.shape),
+    }
+    hasher.update(json.dumps(
+        descriptor, sort_keys=True, separators=(",", ":")).encode("utf-8"))
+    raw = value.reshape(-1).view(torch.uint8).numpy().tobytes()
+    hasher.update(raw)
+
+
+def compute_reference_pack_hash(pack: Mapping[str, Any]) -> str:
+    """Hash pack contents, deliberately excluding ``content_hash`` itself."""
+    required_without_hash = _PACK_KEYS - {"content_hash"}
+    missing = required_without_hash - set(pack)
+    metadata = canonicalize_metadata(pack["metadata"])
+    hasher = hashlib.sha256()
+    hasher.update(b"video-nystrom-reference-pack\0")
+    hasher.update(str(int(pack["format_version"])).encode("ascii"))
+    hasher.update(json.dumps(
+        metadata,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+    ).encode("utf-8"))
+    for name in ("landmarks", "alpha", "k_rr", "sigma"):
+        value = pack[name]
+        _hash_tensor(hasher, name, value)
+    return hasher.hexdigest()
+
+
+def validate_reference_pack(
+    pack: Mapping[str, Any],
+    *,
+    expected_metadata: Optional[Mapping[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Prepare reference tensors without enforcing metadata or content identity."""
+    return {
+        "format_version": REFERENCE_PACK_VERSION,
+        "landmarks": pack["landmarks"],
+        "alpha": pack["alpha"],
+        "k_rr": pack["k_rr"].reshape(()),
+        "sigma": pack["sigma"].reshape(()),
+        "metadata": canonicalize_metadata(pack["metadata"]),
+        "content_hash": pack["content_hash"],
+    }
+
+
+@torch.no_grad()
+def build_video_nystrom_reference(
+    reference: torch.Tensor,
+    *,
+    num_landmarks: int,
+    metadata: Mapping[str, Any],
+    bandwidth: Optional[float] = None,
+    median_max_samples: Optional[int] = 2048,
+    kmeans_iterations: int = 20,
+    ridge: float = 1.0e-6,
+    chunk_size: int = 256,
+    seed: int = 0,
+) -> Dict[str, Any]:
+    """Construct a frozen, portable Nyström reference pack."""
+    reference_count, num_tokens, feature_dim = _validate_video_tensor(
+        reference, "reference")
+    seed = int(seed)
+    reference = reference.detach().float()
+    if bandwidth is None:
+        sigma = median_frobenius_sigma(
+            reference,
+            max_samples=median_max_samples,
+            chunk_size=int(chunk_size),
+        )
+        bandwidth_source = "median_frobenius_distance"
+    else:
+        sigma = reference.new_tensor(float(bandwidth))
+        bandwidth_source = "provided"
+
+    landmarks = deterministic_video_kmeans(
+        reference,
+        int(num_landmarks),
+        iterations=int(kmeans_iterations),
+        chunk_size=int(chunk_size),
+    )
+    alpha = compute_nystrom_alpha(
+        reference,
+        landmarks,
+        sigma,
+        ridge=float(ridge),
+        chunk_size=int(chunk_size),
+    )
+    k_rr = compute_reference_krr(
+        reference, sigma, chunk_size=int(chunk_size))
+
+    canonical_metadata = canonicalize_metadata(metadata)
+    technical_metadata = {
+        "sample_axis": "video",
+        "sample_shape": [num_tokens, feature_dim],
+        "kernel": "gaussian_rbf",
+        "kernel_metric": "frobenius",
+        "generated_gram": "biased_exact_including_diagonal",
+        "reference_count": reference_count,
+        "num_landmarks": int(num_landmarks),
+        "bandwidth_source": bandwidth_source,
+        "median_max_samples": (
+            None if median_max_samples is None else int(median_max_samples)
+        ),
+        "kmeans": "deterministic_farthest_first_lloyd",
+        "seed": seed,
+        "kmeans_iterations": int(kmeans_iterations),
+        "ridge": float(ridge),
+    }
+    canonical_metadata["nystrom"] = technical_metadata
+    canonical_metadata = canonicalize_metadata(canonical_metadata)
+
+    pack: Dict[str, Any] = {
+        "format_version": REFERENCE_PACK_VERSION,
+        "landmarks": landmarks.detach().cpu().contiguous(),
+        "alpha": alpha.detach().cpu().contiguous(),
+        "k_rr": k_rr.detach().cpu().reshape(()),
+        "sigma": sigma.detach().cpu().reshape(()),
+        "metadata": canonical_metadata,
+    }
+    pack["content_hash"] = compute_reference_pack_hash(pack)
+    return validate_reference_pack(pack)
+
+
+def save_reference_pack(
+    pack: Mapping[str, Any],
+    path: str | os.PathLike,
+    *,
+    atomic: bool = True,
+) -> None:
+    """Validate and save a portable pack, atomically by default."""
+    validated = validate_reference_pack(pack)
+    portable = {
+        **validated,
+        "landmarks": validated["landmarks"].detach().cpu().contiguous(),
+        "alpha": validated["alpha"].detach().cpu().contiguous(),
+        "k_rr": validated["k_rr"].detach().cpu().reshape(()),
+        "sigma": validated["sigma"].detach().cpu().reshape(()),
+    }
+    # Moving tensors to CPU is part of the serialized representation, so hash
+    # the exact portable payload rather than copying the incoming hash blindly.
+    portable["content_hash"] = compute_reference_pack_hash(portable)
+    destination = os.path.abspath(os.fspath(path))
+    if not atomic:
+        torch.save(portable, destination)
+        return
+    directory = os.path.dirname(destination)
+    descriptor, temporary = tempfile.mkstemp(
+        prefix=f".{os.path.basename(destination)}.",
+        suffix=".tmp",
+        dir=directory,
+    )
+    os.close(descriptor)
+    try:
+        torch.save(portable, temporary)
+        os.replace(temporary, destination)
+    finally:
+        if os.path.exists(temporary):
+            os.unlink(temporary)
+
+
+def load_reference_pack(
+    path: str | os.PathLike,
+    *,
+    expected_metadata: Optional[Mapping[str, Any]] = None,
+    map_location: str | torch.device = "cpu",
+) -> Dict[str, Any]:
+    """Load a reference pack and prepare its tensors."""
+    try:
+        pack = torch.load(
+            os.fspath(path),
+            map_location=map_location,
+            weights_only=True,
+        )
+    except TypeError:
+        pack = torch.load(os.fspath(path), map_location=map_location)
+    return validate_reference_pack(
+        pack, expected_metadata=expected_metadata)
+
+
+def _kernel_sum_chunked(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    sigma: torch.Tensor,
+    chunk_size: int,
+) -> torch.Tensor:
+    total = left.new_zeros(())
+    self_kernel = left is right
+    for row_start in range(0, left.shape[0], chunk_size):
+        rows = left[row_start:row_start + chunk_size]
+        for col_start in range(0, right.shape[0], chunk_size):
+            cols = right[col_start:col_start + chunk_size]
+            kernel = rbf_frobenius_kernel(rows, cols, sigma)
+            block_sum = kernel.sum()
+            if self_kernel and row_start == col_start:
+                # The RBF self-kernel is analytically one.  Explicitly replace
+                # it to avoid norm/Gram cancellation error in FP32.
+                diagonal = kernel.diagonal()
+                block_sum = block_sum + (1.0 - diagonal).sum()
+            total = total + block_sum
+    return total
+
+
+def _kernel_matvec_chunked(
+    samples: torch.Tensor,
+    landmarks: torch.Tensor,
+    alpha: torch.Tensor,
+    sigma: torch.Tensor,
+    chunk_size: int,
+) -> torch.Tensor:
+    outputs = []
+    for row_start in range(0, samples.shape[0], chunk_size):
+        rows = samples[row_start:row_start + chunk_size]
+        values = rows.new_zeros(rows.shape[0])
+        for col_start in range(0, landmarks.shape[0], chunk_size):
+            cols = landmarks[col_start:col_start + chunk_size]
+            kernel = rbf_frobenius_kernel(rows, cols, sigma)
+            values = values + kernel @ alpha[
+                col_start:col_start + chunk_size]
+        outputs.append(values)
+    return torch.cat(outputs)
+
+
+class VideoNystromMMD(nn.Module):
+    """One encoder's MMD: repulsion - 2 * attraction + reference constant.
+
+    The generated-sample kernel is evaluated exactly. Nyström approximates
+    only attraction to the persistent reference bank; attraction to the
+    current reference minibatch is also exact when enabled.
+    """
+
+    _STATE_VERSION = 1
+    _ATTRACTION_EMA_STATE_VERSION = 1
+    _REPULSION_FIFO_STATE_VERSION = 1
+
+    def __init__(
+        self,
+        reference_pack: Mapping[str, Any],
+        *,
+        expected_metadata: Optional[Mapping[str, Any]] = None,
+        kernel_chunk_size: int = 256,
+        attraction_ema_enabled: bool = False,
+        attraction_ema_beta: float = 0.99,
+        attraction_weight: float = 1.0,
+        current_teacher_attraction_enabled: bool = False,
+        current_teacher_attraction_weight: float = 1.0,
+        repulsion_fifo_size: int = 0,
+    ):
+        super().__init__()
+        current_teacher_attraction_enabled = bool(
+            current_teacher_attraction_enabled)
+        repulsion_fifo_size = int(repulsion_fifo_size)
+        pack = validate_reference_pack(
+            reference_pack, expected_metadata=expected_metadata)
+        self.kernel_chunk_size = int(kernel_chunk_size)
+        self.attraction_ema_enabled = bool(attraction_ema_enabled)
+        self.attraction_ema_beta = float(attraction_ema_beta)
+        self.attraction_weight = float(attraction_weight)
+        self.current_teacher_attraction_enabled = (
+            current_teacher_attraction_enabled)
+        self.current_teacher_attraction_weight = float(
+            current_teacher_attraction_weight)
+        self.repulsion_fifo_size = repulsion_fifo_size
+        self.repulsion_fifo_enabled = bool(repulsion_fifo_size > 0)
+        self.reference_metadata = pack["metadata"]
+        self.reference_hash = pack["content_hash"]
+        self.register_buffer(
+            "landmarks", pack["landmarks"].detach().float().clone(),
+            persistent=True)
+        self.register_buffer(
+            "alpha", pack["alpha"].detach().float().clone(),
+            persistent=True)
+        self.register_buffer(
+            "k_rr", pack["k_rr"].detach().float().reshape(()).clone(),
+            persistent=True)
+        self.register_buffer(
+            "sigma", pack["sigma"].detach().float().reshape(()).clone(),
+            persistent=True)
+        # FIFO contents have their own strict trainer-checkpoint payload.  A
+        # non-persistent buffer still follows ``module.to(device)`` without
+        # silently duplicating the large history in the module state_dict.
+        self.register_buffer(
+            "repulsion_fifo",
+            torch.empty(
+                (0, int(self.landmarks.shape[1]), int(self.landmarks.shape[2])),
+                dtype=torch.float32,
+                device=self.landmarks.device,
+            ),
+            persistent=False,
+        )
+        self.register_buffer(
+            "repulsion_fifo_commits",
+            torch.zeros((), dtype=torch.long, device=self.landmarks.device),
+            persistent=False,
+        )
+        # Keep the disabled module's state_dict byte-schema compatible with
+        # the original exact objective.  Enabled modules persist these two
+        # scalars and advertise the changed schema through extra state.
+        self.register_buffer(
+            "attraction_ema",
+            torch.zeros((), dtype=torch.float32),
+            persistent=self.attraction_ema_enabled,
+        )
+        self.register_buffer(
+            "attraction_ema_updates",
+            torch.zeros((), dtype=torch.long),
+            persistent=self.attraction_ema_enabled,
+        )
+        # Pending state is deliberately not a buffer: an uncommitted optimizer
+        # update must never enter a checkpoint.
+        self._pending_attraction: Optional[torch.Tensor] = None
+        self._pending_repulsion_fifo: Optional[torch.Tensor] = None
+        # The final FIFO allocation is prepared before optimizer.step().
+        # Post-step commit only publishes this already-owned tensor.
+        self._prepared_repulsion_fifo: Optional[torch.Tensor] = None
+
+    @property
+    def sample_shape(self) -> Tuple[int, int]:
+        return int(self.landmarks.shape[1]), int(self.landmarks.shape[2])
+
+    @property
+    def attraction_ema_pending(self) -> bool:
+        return self._pending_attraction is not None
+
+    @property
+    def repulsion_fifo_pending(self) -> bool:
+        return self._pending_repulsion_fifo is not None
+
+    @property
+    def repulsion_fifo_fill(self) -> int:
+        return int(self.repulsion_fifo.shape[0])
+
+    def forward(
+        self,
+        generated: Optional[torch.Tensor] = None,
+        current_teacher: Optional[torch.Tensor] = None,
+        *,
+        gradient_students: Optional[torch.Tensor] = None,
+        student_reference: Optional[torch.Tensor] = None,
+        gradient_student_indices: Optional[
+            torch.Tensor | Sequence[int]
+        ] = None,
+        stage: bool = True,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Evaluate the legacy full-batch or explicit partial-gradient loss.
+
+        Legacy callers pass ``generated`` and retain the original two-sided
+        ``K(generated, generated)`` derivative exactly. Partial-gradient
+        callers instead pass both ``gradient_students`` (the row-side samples
+        whose derivatives are required) and ``student_reference`` (the full
+        detached column-side student population). The latter computes
+
+        ``mean K(G,S) - 2 * attraction(G) + reference_constant``
+
+        with a rectangular ``G x S`` student population. The matched
+        teacher population must have exactly ``|G|`` samples. Neither
+        ``student_reference`` nor ``current_teacher`` can receive gradients.
+
+        Passing ``gradient_student_indices`` opts into the full-student
+        partial-gradient estimator. The indices identify the positions of
+        ``gradient_students`` in ``student_reference``. Those differentiable
+        values are functionally inserted into the otherwise detached student
+        population, yielding complete ``mean K(S,S)``, fixed attraction over
+        all ``S``, and exact current-teacher attraction over ``S x T``. Only
+        indexed students retain gradients; ``|T|`` must still equal ``|G|``.
+        """
+        # 1. Choose the generated population and samples that retain a graph.
+        # ``student_reference`` is generated data, not the persistent reference
+        # bank; these input names remain compatible with existing callers.
+        partial_mode = (
+            gradient_students is not None or student_reference is not None)
+        full_student_partial_mode = gradient_student_indices is not None
+        if partial_mode:
+            gradient_input = gradient_students
+        else:
+            gradient_input = generated
+
+        gradient_batch, num_tokens, feature_dim = _validate_video_tensor(
+            gradient_input,
+            "gradient_students" if partial_mode else "generated",
+        )
+
+        samples = gradient_input.float()
+        selected_indices = None
+        if partial_mode:
+            reference_batch, reference_tokens, reference_dim = (
+                _validate_video_tensor(
+                    student_reference, "student_reference"))
+            reference_samples = student_reference.detach().float()
+            if full_student_partial_mode:
+                try:
+                    selected_indices = torch.as_tensor(
+                        gradient_student_indices,
+                        device=gradient_input.device,
+                    )
+                except (TypeError, ValueError) as exc:
+                    raise TypeError(
+                        "gradient_student_indices must be an integer "
+                        "tensor or sequence") from exc
+                integer_dtypes = (
+                    torch.uint8,
+                    torch.int8,
+                    torch.int16,
+                    torch.int32,
+                    torch.int64,
+                )
+                selected_indices = selected_indices.to(dtype=torch.long)
+                selected_reference = reference_samples.index_select(
+                    0, selected_indices)
+                # Functionally replace only selected detached values. Reusing
+                # this same object on both sides of K(S,S) gives selected
+                # students both row- and column-side derivatives.
+                objective_samples = reference_samples.index_copy(
+                    0, selected_indices, samples)
+                repulsion_reference_samples = objective_samples
+                attraction_samples = objective_samples
+                attraction_batch = reference_batch
+            else:
+                objective_samples = samples
+                repulsion_reference_samples = reference_samples
+                attraction_samples = samples
+                attraction_batch = gradient_batch
+        else:
+            reference_batch = gradient_batch
+            # Preserve object identity here. The legacy self-kernel helper
+            # relies on it both for its exact diagonal replacement and for the
+            # original two-sided K_ss derivative.
+            reference_samples = samples
+            objective_samples = samples
+            repulsion_reference_samples = reference_samples
+            attraction_samples = samples
+            attraction_batch = gradient_batch
+
+        # 2. Read the frozen reference minibatch for the Monte Carlo term.
+        # These features are targets; gradients flow through generated samples.
+        teacher_samples = None
+        teacher_batch = None
+        if self.current_teacher_attraction_enabled:
+            teacher_batch, teacher_tokens, teacher_dim = (
+                _validate_video_tensor(current_teacher, "current_teacher"))
+            if (
+                teacher_batch != gradient_batch
+                or teacher_tokens != num_tokens
+                or teacher_dim != feature_dim
+            ):
+                student_name = (
+                    "gradient_students" if partial_mode else "generated")
+                raise ValueError(
+                    "current teacher features must exactly match "
+                    f"{student_name} "
+                    "[B,N,D]: "
+                    f"{(teacher_batch, teacher_tokens, teacher_dim)} != "
+                    f"{(gradient_batch, num_tokens, feature_dim)}")
+            # The current teacher is a matched frozen target. Detach before
+            # every kernel operation so callers can never accidentally create
+            # a teacher-side gradient graph.
+            teacher_samples = current_teacher.detach().float()
+
+        # 3. Exact generated-sample repulsion: sum every kernel pair in
+        # chunks. Chunking controls memory without approximating the sum.
+        exact_kgg_sum = _kernel_sum_chunked(
+            objective_samples,
+            repulsion_reference_samples,
+            self.sigma,
+            self.kernel_chunk_size,
+        )
+        if self.repulsion_fifo_enabled:
+            history_size = self.repulsion_fifo_fill
+            denominator = gradient_batch * (
+                gradient_batch - 1 + history_size)
+            # ``_kernel_sum_chunked`` analytically replaces each self-kernel
+            # by one.  Subtracting B therefore removes the complete diagonal,
+            # including its autograd path, without forming a dense mask.
+            current_offdiag_sum = exact_kgg_sum - float(gradient_batch)
+            if history_size:
+                history_cross_sum = _kernel_sum_chunked(
+                    samples,
+                    self.repulsion_fifo,
+                    self.sigma,
+                    self.kernel_chunk_size,
+                )
+            else:
+                history_cross_sum = samples.new_zeros(())
+            repulsion = (
+                current_offdiag_sum + history_cross_sum
+            ) / float(denominator)
+        else:
+            history_size = 0
+            # Full-population replay keeps B * B pairs even when only a
+            # subset is differentiated. The rectangular mode uses |G| * B.
+            # Gradient-subset compensation is handled by the replay caller.
+            if full_student_partial_mode:
+                denominator = reference_batch * reference_batch
+            else:
+                denominator = gradient_batch * reference_batch
+            current_offdiag_sum = None
+            history_cross_sum = None
+            repulsion = exact_kgg_sum / float(denominator)
+        # 4. Persistent Nyström attraction (paper Eq. 8): K(x, Z) @ alpha.
+        # The saved ``alpha`` vector is paper beta, not the mixture weight.
+        fixed_attraction_values = _kernel_matvec_chunked(
+            attraction_samples,
+            self.landmarks,
+            self.alpha,
+            self.sigma,
+            self.kernel_chunk_size,
+        )
+        fixed_current_attraction = fixed_attraction_values.mean()
+        current_teacher_attraction = None
+        current_teacher_self_kernel = None
+        current_teacher_exact_mmd = None
+        # 5. Hybrid reference mean (paper Eq. 9): mix the persistent term
+        # with exact minibatch attraction over all generated x reference pairs.
+        # Group weights 2:1 correspond to paper alpha_e = 1/3.
+        if self.current_teacher_attraction_enabled:
+            fixed_weight = self.attraction_weight
+            current_weight = self.current_teacher_attraction_weight
+            teacher_weight_sum = fixed_weight + current_weight
+            fixed_normalized_weight = fixed_weight / teacher_weight_sum
+            current_normalized_weight = current_weight / teacher_weight_sum
+            current_teacher_attraction = _kernel_sum_chunked(
+                attraction_samples,
+                teacher_samples,
+                self.sigma,
+                self.kernel_chunk_size,
+            ) / float(attraction_batch * teacher_batch)
+            # The reference self-kernel has no generator gradient. Its value
+            # still affects the release's later detached-scale division.
+            with torch.no_grad():
+                current_teacher_self_kernel = _kernel_sum_chunked(
+                    teacher_samples,
+                    teacher_samples,
+                    self.sigma,
+                    self.kernel_chunk_size,
+                ) / float(teacher_batch * teacher_batch)
+            current_attraction = (
+                fixed_normalized_weight * fixed_current_attraction
+                + current_normalized_weight * current_teacher_attraction
+            )
+            reference_constant = (
+                fixed_normalized_weight * self.k_rr
+                + current_normalized_weight * current_teacher_self_kernel
+            )
+            current_teacher_exact_mmd = (
+                repulsion
+                - 2.0 * current_teacher_attraction
+                + current_teacher_self_kernel
+            )
+        else:
+            fixed_normalized_weight = None
+            current_normalized_weight = None
+            teacher_weight_sum = None
+            current_attraction = fixed_current_attraction
+            reference_constant = self.k_rr
+        # 6. Optional historical attraction/FIFO state. These branches are
+        # inactive in the released three-encoder configuration.
+        history_updates = int(self.attraction_ema_updates.item())
+        if self.attraction_ema_enabled and history_updates > 0:
+            ema_blend_current_coefficient = 1.0 - self.attraction_ema_beta
+            ema_attraction = (
+                self.attraction_ema_beta * self.attraction_ema.detach()
+                + ema_blend_current_coefficient * current_attraction
+            )
+            if self.current_teacher_attraction_enabled:
+                # Keep the transactional EMA's smoothed forward value while
+                # preserving the full derivative of this update's normalized
+                # fixed/current teacher mixture. Without this straight-through
+                # form beta=0.99 would silently shrink all teacher gradients
+                # to one percent relative to K_ss.
+                attraction = current_attraction + (
+                    ema_attraction - current_attraction).detach()
+                current_coefficient = 1.0
+            else:
+                attraction = ema_attraction
+                current_coefficient = ema_blend_current_coefficient
+        else:
+            ema_blend_current_coefficient = 1.0
+            current_coefficient = 1.0
+            attraction = current_attraction
+
+        if self.attraction_ema_enabled and stage:
+            if self._pending_attraction is not None:
+                raise RuntimeError(
+                    "an attraction EMA update is already staged; "
+                    "commit or discard it before another staged forward")
+            self._pending_attraction = current_attraction.detach().clone()
+
+        if self.repulsion_fifo_enabled and stage:
+            if self._pending_repulsion_fifo is not None:
+                raise RuntimeError(
+                    "a repulsion FIFO update is already staged; "
+                    "commit or discard it before another staged forward")
+            # The fresh GradCache supplies the identical global batch on every
+            # rank.  Cloning it here gives each rank the same independently
+            # owned FIFO candidate and prevents all gradients/aliasing.
+            self._pending_repulsion_fifo = (
+                samples.detach().to(dtype=torch.float32).clone())
+
+        # 7. Assemble the raw single-encoder loss. MultiEncoderMMD applies
+        # detached-scale normalization and encoder weighting after this.
+        if self.current_teacher_attraction_enabled:
+            # This is the weighted average of the fixed-reference MMD and the
+            # current exact MMD. K_ss is shared, so it appears once and the
+            # total attraction coefficient remains -2 for every raw scale.
+            loss = repulsion - 2.0 * attraction + reference_constant
+            effective_attraction_weight = 1.0
+        else:
+            # Preserve the legacy value and gradient exactly when the opt-in
+            # current teacher path is disabled.
+            if self.attraction_weight == 1.0:
+                loss = repulsion - 2.0 * attraction + self.k_rr
+            else:
+                loss = (
+                    repulsion
+                    - 2.0 * self.attraction_weight * attraction
+                    + self.k_rr
+                )
+            effective_attraction_weight = self.attraction_weight
+        metadata = {
+            "repulsion": repulsion.detach(),
+            "attraction": attraction.detach(),
+            "attraction_weight": loss.new_tensor(
+                effective_attraction_weight),
+            "reference_constant": reference_constant.detach(),
+            "sigma": self.sigma.detach(),
+            "global_samples": torch.tensor(
+                reference_batch, device=loss.device, dtype=torch.long),
+        }
+        if partial_mode:
+            metadata.update({
+                "partial_student_reference_enabled": loss.new_ones(
+                    (), dtype=torch.long),
+                "gradient_student_samples": loss.new_tensor(
+                    gradient_batch, dtype=torch.long),
+                "student_reference_samples": loss.new_tensor(
+                    reference_batch, dtype=torch.long),
+                "student_reference_detached": loss.new_ones(
+                    (), dtype=torch.long),
+                "repulsion_denominator": loss.new_tensor(
+                    denominator, dtype=torch.long),
+                "student_repulsion_denominator": loss.new_tensor(
+                    denominator, dtype=torch.long),
+                "fixed_teacher_attraction_denominator": loss.new_tensor(
+                    attraction_batch, dtype=torch.long),
+            })
+            if full_student_partial_mode:
+                metadata.update({
+                    "full_student_partial_gradient": loss.new_ones(
+                        (), dtype=torch.long),
+                    "student_repulsion_gradient_sides": loss.new_tensor(
+                        2, dtype=torch.long),
+                })
+        if self.current_teacher_attraction_enabled:
+            metadata.update({
+                "current_teacher_attraction_enabled": loss.new_ones(
+                    (), dtype=torch.long),
+                "fixed_teacher_attraction": (
+                    fixed_current_attraction.detach()),
+                "current_teacher_attraction": (
+                    current_teacher_attraction.detach()),
+                "fixed_teacher_attraction_weight": loss.new_tensor(
+                    self.attraction_weight),
+                "current_teacher_attraction_weight": loss.new_tensor(
+                    self.current_teacher_attraction_weight),
+                "teacher_attraction_weight_sum": loss.new_tensor(
+                    teacher_weight_sum),
+                "fixed_teacher_attraction_normalized_weight": (
+                    loss.new_tensor(fixed_normalized_weight)),
+                "current_teacher_attraction_normalized_weight": (
+                    loss.new_tensor(current_normalized_weight)),
+                "mixed_current_attraction": current_attraction.detach(),
+                "fixed_reference_constant": self.k_rr.detach(),
+                "current_teacher_self_kernel": (
+                    current_teacher_self_kernel.detach()),
+                "current_teacher_exact_mmd": (
+                    current_teacher_exact_mmd.detach()),
+                "current_teacher_samples": loss.new_tensor(
+                    teacher_batch, dtype=torch.long),
+            })
+            if partial_mode:
+                metadata.update({
+                    "current_teacher_attraction_denominator": (
+                        loss.new_tensor(
+                            attraction_batch * teacher_batch,
+                            dtype=torch.long,
+                        )
+                    ),
+                    "current_teacher_self_kernel_denominator": (
+                        loss.new_tensor(
+                            teacher_batch * teacher_batch,
+                            dtype=torch.long,
+                        )
+                    ),
+                })
+        if self.attraction_ema_enabled:
+            metadata.update({
+                "current_attraction": current_attraction.detach(),
+                "attraction_ema_enabled": loss.new_ones(
+                    (), dtype=torch.long),
+                "attraction_ema_updates": (
+                    self.attraction_ema_updates.detach().clone()),
+                "attraction_current_coefficient": loss.new_tensor(
+                    current_coefficient),
+            })
+            if self.current_teacher_attraction_enabled:
+                metadata.update({
+                    "attraction_ema_blend_current_coefficient": (
+                        loss.new_tensor(
+                            ema_blend_current_coefficient)),
+                    "attraction_gradient_current_coefficient": (
+                        loss.new_ones(())),
+                    "attraction_ema_smoothed_value": attraction.detach(),
+                    "effective_ema_loss": loss.detach(),
+                    "attraction_gradient_mode": (
+                        "straight_through_full_current_mix"),
+                })
+        if self.repulsion_fifo_enabled:
+            fill_after_commit = min(
+                self.repulsion_fifo_size,
+                history_size + gradient_batch,
+            )
+            metadata.update({
+                "repulsion_fifo_enabled": loss.new_ones(
+                    (), dtype=torch.long),
+                "repulsion_fifo_capacity": loss.new_tensor(
+                    self.repulsion_fifo_size, dtype=torch.long),
+                "repulsion_fifo_fill_before": loss.new_tensor(
+                    history_size, dtype=torch.long),
+                # Compatibility spelling used by the trainer logger.
+                "repulsion_fifo_filled_before": loss.new_tensor(
+                    history_size, dtype=torch.long),
+                "repulsion_fifo_fill_after_commit": loss.new_tensor(
+                    fill_after_commit, dtype=torch.long),
+                "repulsion_fifo_commits": (
+                    self.repulsion_fifo_commits.detach().clone()),
+                "repulsion_fifo_committed_updates": (
+                    self.repulsion_fifo_commits.detach().clone()),
+                "repulsion_current_offdiag_kernel_sum": (
+                    current_offdiag_sum.detach()),
+                "repulsion_fifo_current_offdiag_sum": (
+                    current_offdiag_sum.detach()),
+                "repulsion_history_kernel_sum": (
+                    history_cross_sum.detach()),
+                "repulsion_fifo_history_cross_sum": (
+                    history_cross_sum.detach()),
+                "repulsion_denominator": loss.new_tensor(
+                    denominator, dtype=torch.long),
+                "repulsion_fifo_denominator": loss.new_tensor(
+                    denominator, dtype=torch.long),
+            })
+        return loss, metadata
+
+    @torch.no_grad()
+    def prepare_commit(self) -> None:
+        """Validate all staged distribution state before optimizer mutation."""
+        if self.attraction_ema_enabled:
+            if self._pending_attraction is None:
+                raise RuntimeError("no attraction EMA update is staged")
+        if self.repulsion_fifo_enabled:
+            pending = self._pending_repulsion_fifo
+            if pending is None:
+                raise RuntimeError("no repulsion FIFO update is staged")
+            if self._prepared_repulsion_fifo is None:
+                combined = torch.cat((self.repulsion_fifo, pending), dim=0)
+                if combined.shape[0] > self.repulsion_fifo_size:
+                    combined = combined[-self.repulsion_fifo_size:]
+                # Allocate the right-sized, independently owned result while
+                # the trainer can still abort before optimizer mutation.
+                self._prepared_repulsion_fifo = (
+                    combined.detach().contiguous().clone())
+
+    @torch.no_grad()
+    def commit(self) -> None:
+        """Publish staged state only after a successful optimizer update."""
+        if (
+            self._pending_attraction is None
+            and self._pending_repulsion_fifo is None
+        ):
+            return
+        self.prepare_commit()
+        if self._pending_attraction is not None:
+            pending_attraction = self._pending_attraction
+            if int(self.attraction_ema_updates.item()) == 0:
+                self.attraction_ema.copy_(pending_attraction)
+            else:
+                self.attraction_ema.mul_(self.attraction_ema_beta).add_(
+                    pending_attraction,
+                    alpha=1.0 - self.attraction_ema_beta,
+                )
+            self.attraction_ema_updates.add_(1)
+        if self._pending_repulsion_fifo is not None:
+            if self._prepared_repulsion_fifo is None:
+                raise RuntimeError("repulsion FIFO commit was not prepared")
+            self.repulsion_fifo = self._prepared_repulsion_fifo
+            self.repulsion_fifo_commits.add_(1)
+        self.discard()
+
+    def discard(self) -> None:
+        """Drop all staged state after a failed or skipped update."""
+        self._pending_attraction = None
+        self._pending_repulsion_fifo = None
+        self._prepared_repulsion_fifo = None
+
+    def repulsion_fifo_state_dict(self) -> Optional[Dict[str, Any]]:
+        """Return the committed FIFO in oldest-to-newest order."""
+        if not self.repulsion_fifo_enabled:
+            return None
+        if self.repulsion_fifo_pending:
+            raise RuntimeError(
+                "refusing to checkpoint a pending repulsion FIFO update")
+        return {
+            "version": self._REPULSION_FIFO_STATE_VERSION,
+            "capacity": self.repulsion_fifo_size,
+            "sample_shape": list(self.sample_shape),
+            "dtype": "float32",
+            "count": self.repulsion_fifo_fill,
+            "commits": int(self.repulsion_fifo_commits.item()),
+            "values": self.repulsion_fifo.detach().cpu().clone(),
+        }
+    def validate_repulsion_fifo_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> Dict[str, Any]:
+        """Validate and normalize a portable FIFO checkpoint payload."""
+        if not self.repulsion_fifo_enabled:
+            raise ValueError(
+                "cannot load repulsion FIFO state while it is disabled")
+        expected_keys = {
+            "version", "capacity", "sample_shape", "dtype",
+            "count", "commits", "values",
+        }
+        count = int(state["count"])
+        commits = int(state["commits"])
+        values = state["values"]
+        return {
+            "version": self._REPULSION_FIFO_STATE_VERSION,
+            "capacity": self.repulsion_fifo_size,
+            "sample_shape": list(self.sample_shape),
+            "dtype": "float32",
+            "count": count,
+            "commits": commits,
+            "values": values.detach().cpu().clone(),
+        }
+
+    @torch.no_grad()
+    def load_repulsion_fifo_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> None:
+        validated = self.validate_repulsion_fifo_state_dict(state)
+        self.repulsion_fifo = validated["values"].to(
+            device=self.landmarks.device,
+            dtype=torch.float32,
+        ).contiguous().clone()
+        self.repulsion_fifo_commits.fill_(validated["commits"])
+        self.discard()
+
+    def attraction_ema_state_dict(self) -> Optional[Dict[str, Any]]:
+        """Return the small trainer checkpoint payload for the opt-in EMA."""
+        if not self.attraction_ema_enabled:
+            return None
+        if self.attraction_ema_pending:
+            raise RuntimeError(
+                "refusing to checkpoint a pending attraction EMA update")
+        return {
+            "version": self._ATTRACTION_EMA_STATE_VERSION,
+            "beta": self.attraction_ema_beta,
+            "value": self.attraction_ema.detach().cpu().clone(),
+            "updates": self.attraction_ema_updates.detach().cpu().clone(),
+        }
+
+    def validate_attraction_ema_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> Dict[str, Any]:
+        """Validate and normalize a small attraction EMA checkpoint payload."""
+        if not self.attraction_ema_enabled:
+            raise ValueError(
+                "cannot load attraction EMA state while the path is disabled")
+        expected_keys = {"version", "beta", "value", "updates"}
+        value = state["value"]
+        updates = state["updates"]
+        return {
+            "version": self._ATTRACTION_EMA_STATE_VERSION,
+            "beta": self.attraction_ema_beta,
+            "value": value.detach().cpu().clone(),
+            "updates": updates.detach().cpu().clone(),
+        }
+
+    @torch.no_grad()
+    def load_attraction_ema_state_dict(
+        self,
+        state: Mapping[str, Any],
+    ) -> None:
+        """Strictly restore only the transactional attraction EMA scalars."""
+        validated = self.validate_attraction_ema_state_dict(state)
+        self.attraction_ema.copy_(
+            validated["value"].to(
+                device=self.attraction_ema.device,
+                dtype=self.attraction_ema.dtype,
+            )
+        )
+        self.attraction_ema_updates.copy_(
+            validated["updates"].to(
+                device=self.attraction_ema_updates.device)
+        )
+        self.discard()
+
+    def get_extra_state(self) -> Dict[str, Any]:
+        state = {
+            "version": self._STATE_VERSION,
+            "reference_metadata": self.reference_metadata,
+            "reference_hash": self.reference_hash,
+        }
+        if self.attraction_ema_enabled:
+            state = {
+                **state,
+                "version": self._STATE_VERSION + 1,
+                "attraction_ema": {
+                    "version": self._ATTRACTION_EMA_STATE_VERSION,
+                    "beta": self.attraction_ema_beta,
+                },
+            }
+        if self.current_teacher_attraction_enabled:
+            state = {
+                **state,
+                "version": (4 if self.attraction_ema_enabled else 3),
+                "current_teacher_attraction": {
+                    "enabled": True,
+                    "fixed_weight": self.attraction_weight,
+                    "current_weight": (
+                        self.current_teacher_attraction_weight),
+                    "normalization": "divide_by_group_weight_sum",
+                    "estimator": "exact_biased_gaussian_mmd",
+                },
+            }
+        return state
+
+    def set_extra_state(self, state: Mapping[str, Any]) -> None:
+        expected = self.get_extra_state()

1.3B/code/losses/wan_kv_mmd.py

@@ -0,0 +1,507 @@
+"""RBF-MMD on Wan KV matrices with a mathematically valid EMA distribution.
+
+An RBF kernel has an (infinite-dimensional) feature map ``phi`` and
+
+    MMD²(P, Q) = || E_P[phi(x)] - E_Q[phi(y)] ||².
+
+We approximate ``phi`` with fixed random Fourier features (RFF), then keep an
+EMA of the *kernel mean embeddings*.  This represents an EMA mixture of whole
+historical distributions; averaging raw KV rows across unrelated videos does
+not.  Historical embeddings are detached, so gradients only reach the latest
+student KV matrix.  The natural ``1-beta`` factor is intentionally preserved.
+
+History updates are transactional: ``forward(stage=True)`` only stages detached
+current embeddings.  ``commit()`` must be called after a successful optimizer
+step; ``discard()`` is used when the optimizer step is skipped.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Iterable, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.distributed.nn.functional as dist_nn
+from torch import nn
+
+
+def _sq_dist(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    x2 = x.square().sum(-1, keepdim=True)
+    y2 = y.square().sum(-1).unsqueeze(0)
+    return (x2 + y2 - 2.0 * x @ y.t()).clamp_min(0)
+
+
+def exact_rbf_mmd(
+    student: torch.Tensor,
+    teacher: torch.Tensor,
+    *,
+    bandwidth: float,
+    bandwidth_multipliers: Iterable[float] = (0.5, 1.0, 2.0),
+) -> torch.Tensor:
+    """Biased multi-bandwidth RBF-MMD² reference implementation."""
+    if student.ndim != 2 or teacher.ndim != 2:
+        raise ValueError("exact_rbf_mmd expects [samples, features] matrices")
+    if student.shape[1] != teacher.shape[1] or min(student.shape[0], teacher.shape[0]) < 1:
+        raise ValueError("student/teacher matrices need matching feature dims and samples")
+    if not math.isfinite(float(bandwidth)) or float(bandwidth) <= 0:
+        raise ValueError("bandwidth must be finite and positive")
+    multipliers = tuple(float(v) for v in bandwidth_multipliers)
+    if not multipliers or any((not math.isfinite(v) or v <= 0) for v in multipliers):
+        raise ValueError("bandwidth multipliers must be finite and positive")
+
+    x, y = student.float(), teacher.float()
+    dxx, dyy, dxy = _sq_dist(x, x), _sq_dist(y, y), _sq_dist(x, y)
+    bandwidth_sq = x.new_tensor(float(bandwidth) ** 2)
+    terms = []
+    for multiplier in multipliers:
+        denom = 2.0 * bandwidth_sq * multiplier * multiplier
+        terms.append(
+            torch.exp(-dxx / denom).mean()
+            + torch.exp(-dyy / denom).mean()
+            - 2.0 * torch.exp(-dxy / denom).mean())
+    return torch.stack(terms).mean()
+
+
+class EMARBFMMD(nn.Module):
+    """EMA RBF kernel-mean MMD using fixed random Fourier features.
+
+    ``sample_mode="rows"`` preserves the original behavior, where every
+    ``[D]`` row in a ``[B, N, D]`` input is one sample.  In
+    ``sample_mode="video_matrix"``, the only sample axis is ``B`` and every
+    position-sensitive ``[N, D]`` matrix is one sample.  Its RBF distance is
+    the Frobenius distance, equivalently the Euclidean distance after
+    vectorizing the matrix.
+    """
+
+    _STATE_VERSION = 2
+    _SAMPLE_MODES = ("rows", "video_matrix")
+
+    def __init__(
+        self,
+        beta: float = 0.99,
+        num_rff_features: int = 256,
+        bandwidth: Optional[float] = None,
+        bandwidth_multipliers: Iterable[float] = (0.5, 1.0, 2.0),
+        seed: int = 0,
+        row_normalize: bool = False,
+        sample_mode: str = "rows",
+    ):
+        super().__init__()
+        if not 0.0 <= float(beta) < 1.0:
+            raise ValueError("EMA beta must be in [0, 1)")
+        if int(num_rff_features) <= 0:
+            raise ValueError("num_rff_features must be positive")
+        multipliers = tuple(float(v) for v in bandwidth_multipliers)
+        if not multipliers or any((not math.isfinite(v) or v <= 0) for v in multipliers):
+            raise ValueError("bandwidth multipliers must be finite and positive")
+        if bandwidth is not None and (
+            not math.isfinite(float(bandwidth)) or float(bandwidth) <= 0
+        ):
+            raise ValueError("bandwidth must be finite and positive")
+        if sample_mode not in self._SAMPLE_MODES:
+            raise ValueError(
+                f"sample_mode must be one of {self._SAMPLE_MODES}, "
+                f"got {sample_mode!r}")
+
+        self.beta = float(beta)
+        self.num_rff_features = int(num_rff_features)
+        self.bandwidth_multipliers = multipliers
+        self.seed = int(seed)
+        self.row_normalize = bool(row_normalize)
+        self.sample_mode = sample_mode
+
+        initial_bw_sq = (
+            float(bandwidth) ** 2 if bandwidth is not None else float("nan"))
+        self.register_buffer(
+            "bandwidth_sq", torch.tensor(initial_bw_sq, dtype=torch.float32),
+            persistent=True)
+        self.register_buffer("rff_weight", torch.empty(0), persistent=True)
+        self.register_buffer("rff_phase", torch.empty(0), persistent=True)
+        self.register_buffer("student_ema", torch.empty(0), persistent=True)
+        self.register_buffer("teacher_ema", torch.empty(0), persistent=True)
+        self.register_buffer(
+            "num_updates", torch.zeros((), dtype=torch.long), persistent=True)
+
+        # Deliberately not buffers: an uncommitted optimizer step must never enter
+        # a checkpoint.
+        self._pending_student: Optional[torch.Tensor] = None
+        self._pending_teacher: Optional[torch.Tensor] = None
+
+    @property
+    def embedding_dim(self) -> int:
+        return self.num_rff_features * len(self.bandwidth_multipliers)
+
+    @staticmethod
+    def _distributed() -> bool:
+        return (
+            dist.is_available()
+            and dist.is_initialized()
+            and dist.get_world_size() > 1
+        )
+
+    @classmethod
+    @torch.no_grad()
+    def _all_ranks_true(cls, local_value: bool, device: torch.device) -> bool:
+        """Return one consensus bit before entering later collectives.
+
+        A rank-local exception here would leave the other ranks blocked in the
+        RFF broadcast/all-reduce.  Every input-validation failure therefore
+        becomes an all-rank failure.
+        """
+        if not cls._distributed():
+            return bool(local_value)
+        flag = torch.tensor(
+            int(bool(local_value)), device=device, dtype=torch.int32)
+        dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+        return bool(flag.item())
+
+    @classmethod
+    def _global_mean(
+        cls,
+        local_sum: torch.Tensor,
+        local_count: int,
+        *,
+        return_count: bool = False,
+    ):
+        """Autograd-aware, sample-weighted mean across all ranks.
+
+        The differentiable all-reduce is intentional.  Since every rank
+        backpropagates the same global loss, its backward supplies the right
+        scaling before FSDP averages the replicated parameter gradients.
+        """
+        if not cls._distributed():
+            count = local_sum.new_tensor(float(local_count))
+            mean = local_sum / float(local_count)
+        else:
+            global_sum = dist_nn.all_reduce(
+                local_sum, op=dist.ReduceOp.SUM)
+            count = local_sum.new_tensor(float(local_count))
+            dist.all_reduce(count, op=dist.ReduceOp.SUM)
+            mean = global_sum / count
+        if return_count:
+            return mean, count
+        return mean
+
+    @torch.no_grad()
+    def _estimate_bandwidth_sq(
+        self, student: torch.Tensor, teacher: torch.Tensor, max_samples: int = 256
+    ) -> torch.Tensor:
+        if self.sample_mode == "rows":
+            merged = torch.cat([
+                student.detach().flatten(0, 1),
+                teacher.detach().flatten(0, 1),
+            ], dim=0).float()
+        else:
+            # Keep each [N, D] video matrix intact: B is the sole sample axis.
+            merged = torch.cat([
+                student.detach(),
+                teacher.detach(),
+            ], dim=0).float()
+        if merged.shape[0] > max_samples:
+            indices = torch.linspace(
+                0, merged.shape[0] - 1, max_samples, device=merged.device
+            ).round().long()
+            merged = merged.index_select(0, indices)
+        if self._distributed():
+            # Use all ranks for the one-time median heuristic.  Rows are
+            # padded only for the collective and unpadded immediately after.
+            local_size = torch.tensor(
+                [merged.shape[0]], device=merged.device, dtype=torch.long)
+            gathered_sizes = [
+                torch.zeros_like(local_size) for _ in range(dist.get_world_size())
+            ]
+            dist.all_gather(gathered_sizes, local_size)
+            sizes = [int(size.item()) for size in gathered_sizes]
+            max_size = max(sizes)
+            padded = merged.new_zeros((max_size, *merged.shape[1:]))
+            padded[:merged.shape[0]].copy_(merged)
+            gathered_samples = [
+                torch.empty_like(padded) for _ in range(dist.get_world_size())
+            ]
+            dist.all_gather(gathered_samples, padded)
+            merged = torch.cat([
+                samples[:size]
+                for samples, size in zip(gathered_samples, sizes)
+            ], dim=0)
+
+        if merged.shape[0] < 2:
+            estimate = merged.new_tensor(1.0)
+        else:
+            if not self._distributed() or dist.get_rank() == 0:
+                if self.sample_mode == "rows":
+                    distances_sq = torch.pdist(merged, p=2).square()
+                else:
+                    # ||X-Y||_F^2 without collapsing [N, D] into a token axis.
+                    squared_norm = merged.square().sum(dim=(1, 2))
+                    gram = torch.einsum(
+                        "ind,jnd->ij", merged, merged)
+                    pairwise_sq = (
+                        squared_norm[:, None]
+                        + squared_norm[None, :]
+                        - 2.0 * gram
+                    ).clamp_min(0)
+                    upper = torch.triu_indices(
+                        merged.shape[0],
+                        merged.shape[0],
+                        offset=1,
+                        device=merged.device,
+                    )
+                    distances_sq = pairwise_sq[upper[0], upper[1]]
+                positive = distances_sq[distances_sq > 0]
+                estimate = (
+                    positive.median()
+                    if positive.numel()
+                    else distances_sq.new_tensor(1.0)
+                ).clamp_min(1e-6)
+            else:
+                estimate = merged.new_tensor(1.0)
+        if self._distributed():
+            # All ranks must use the same fixed, checkpointed RFF map.
+            dist.broadcast(estimate, src=0)
+        return estimate
+
+    @torch.no_grad()
+    def _initialize_rff(self, student: torch.Tensor, teacher: torch.Tensor) -> None:
+        if self.sample_mode == "rows":
+            sample_shape = (int(student.shape[-1]),)
+        else:
+            sample_shape = (
+                int(student.shape[-2]),
+                int(student.shape[-1]),
+            )
+        weight_shape = (*sample_shape, self.num_rff_features)
+        if self.rff_weight.numel() > 0:
+            if self.rff_weight.shape != weight_shape:
+                raise ValueError(
+                    "KV sample shape changed: RFF map expects "
+                    f"{tuple(self.rff_weight.shape[:-1])}, got {sample_shape}")
+            return
+        if not bool(torch.isfinite(self.bandwidth_sq).item()):
+            self.bandwidth_sq.copy_(self._estimate_bandwidth_sq(student, teacher))
+
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(self.seed)
+        weight = torch.randn(
+            weight_shape, generator=generator,
+            dtype=torch.float32)
+        phase = 2.0 * math.pi * torch.rand(
+            self.num_rff_features, generator=generator, dtype=torch.float32)
+        self.rff_weight.resize_(weight.shape).copy_(
+            weight.to(device=student.device))
+        self.rff_phase.resize_(phase.shape).copy_(
+            phase.to(device=student.device))
+
+    def _kernel_sum(self, matrix: torch.Tensor) -> Tuple[torch.Tensor, int]:
+        samples = matrix.float()
+        if self.row_normalize:
+            samples = torch.nn.functional.normalize(samples, dim=-1)
+        if self.sample_mode == "rows":
+            samples = samples.flatten(0, 1)
+            base_projection = samples @ self.rff_weight
+        else:
+            # This is exactly
+            # samples.flatten(1) @ rff_weight.flatten(0, 1), while retaining
+            # the [N, D] axes and their positional identity.
+            base_projection = torch.einsum(
+                "bnd,ndr->br", samples, self.rff_weight)
+        bandwidth = self.bandwidth_sq.sqrt().clamp_min(1e-6)
+        scale = math.sqrt(
+            2.0 / (self.num_rff_features * len(self.bandwidth_multipliers)))
+        features = [
+            torch.cos(
+                base_projection / (bandwidth * multiplier) + self.rff_phase
+            ) * scale
+            for multiplier in self.bandwidth_multipliers
+        ]
+        return torch.cat(features, dim=-1).sum(dim=0), int(samples.shape[0])
+
+    def forward(
+        self,
+        student: torch.Tensor,
+        teacher: torch.Tensor,
+        *,
+        stage: bool = True,
+    ) -> Tuple[torch.Tensor, dict]:
+        local_shape_ok = (
+            student.ndim == 3
+            and teacher.ndim == 3
+            and student.shape == teacher.shape
+            and student.device == teacher.device
+            and min(student.shape[0], student.shape[1], student.shape[2]) > 0
+        )
+        if not self._all_ranks_true(local_shape_ok, student.device):
+            raise ValueError(
+                "invalid Wan KV shape/device on at least one rank; each local "
+                "student/teacher pair must match [B, tokens, features]")
+
+        # A mismatched feature width would initialize incompatible RFF maps.
+        # Matrix samples additionally require a shared N because every token
+        # position has its own slice of the position-sensitive random weight.
+        if self._distributed():
+            rff_sample_shape = (
+                student.shape[-1:]
+                if self.sample_mode == "rows"
+                else student.shape[1:]
+            )
+            shape_bounds = torch.tensor(
+                [*rff_sample_shape, *(-value for value in rff_sample_shape)],
+                device=student.device,
+                dtype=torch.long,
+            )
+            dist.all_reduce(shape_bounds, op=dist.ReduceOp.MIN)
+            num_dimensions = len(rff_sample_shape)
+            if any(
+                int(shape_bounds[index].item())
+                != -int(shape_bounds[index + num_dimensions].item())
+                for index in range(num_dimensions)
+            ):
+                raise ValueError(
+                    "Wan KV RFF sample shape differs across distributed ranks")
+
+        rff_initialized = self.rff_weight.numel() > 0
+        expected_rff_shape = (
+            (student.shape[-1], self.num_rff_features)
+            if self.sample_mode == "rows"
+            else (
+                student.shape[-2],
+                student.shape[-1],
+                self.num_rff_features,
+            )
+        )
+        local_rff_shape_ok = (
+            not rff_initialized
+            or self.rff_weight.shape == expected_rff_shape
+        )
+        if not self._all_ranks_true(local_rff_shape_ok, student.device):
+            raise ValueError(
+                "RFF map shape is incompatible on at least one rank")
+        if self._distributed():
+            bandwidth_initialized = bool(
+                torch.isfinite(self.bandwidth_sq).item())
+            state_bounds = torch.tensor(
+                [
+                    int(rff_initialized),
+                    -int(rff_initialized),
+                    int(bandwidth_initialized),
+                    -int(bandwidth_initialized),
+                ],
+                device=student.device,
+                dtype=torch.int32,
+            )
+            dist.all_reduce(state_bounds, op=dist.ReduceOp.MIN)
+            if (
+                int(state_bounds[0].item()) != -int(state_bounds[1].item())
+                or int(state_bounds[2].item()) != -int(state_bounds[3].item())
+            ):
+                raise RuntimeError(
+                    "RFF/bandwidth initialization state differs across ranks")
+
+        local_finite = bool(
+            torch.isfinite(student).all().item()
+            and torch.isfinite(teacher).all().item()
+        )
+        if not self._all_ranks_true(local_finite, student.device):
+            raise FloatingPointError(
+                "non-finite Wan KV feature on at least one rank")
+
+        teacher = teacher.detach()
+        self._initialize_rff(student, teacher)
+        student_sum, student_count = self._kernel_sum(student)
+        teacher_sum, teacher_count = self._kernel_sum(teacher)
+        current_student, global_count = self._global_mean(
+            student_sum, student_count, return_count=True)
+        current_teacher = self._global_mean(teacher_sum, teacher_count)
+
+        if self.student_ema.numel() == 0:
+            current_coefficient = 1.0
+            effective_student = current_student
+            effective_teacher = current_teacher.detach()
+        else:
+            current_coefficient = 1.0 - self.beta
+            # History has no graph. Consequently only current_student receives
+            # gradients, with the mathematically correct (1-beta) coefficient.
+            effective_student = (
+                self.beta * self.student_ema.detach()
+                + (1.0 - self.beta) * current_student)
+            effective_teacher = (
+                self.beta * self.teacher_ema.detach()
+                + (1.0 - self.beta) * current_teacher.detach())
+
+        loss = (effective_student - effective_teacher).square().sum()
+        if stage:
+            self._pending_student = current_student.detach()
+            self._pending_teacher = current_teacher.detach()
+        return loss, {
+            "bandwidth": self.bandwidth_sq.detach().sqrt(),
+            "history_updates": self.num_updates.detach().clone(),
+            "current_coefficient": loss.new_tensor(current_coefficient),
+            "local_samples": torch.tensor(
+                student_count, device=loss.device, dtype=torch.long),
+            "global_samples": global_count.detach().to(dtype=torch.long),
+        }
+
+    @torch.no_grad()
+    def commit(self) -> None:
+        """Commit staged kernel means after a successful optimizer step."""
+        if self._pending_student is None or self._pending_teacher is None:
+            return
+        if not (
+            torch.isfinite(self._pending_student).all()
+            and torch.isfinite(self._pending_teacher).all()
+        ):
+            self.discard()
+            raise FloatingPointError("refusing to commit non-finite KV kernel mean")
+        if self.student_ema.numel() == 0:
+            self.student_ema.resize_as_(self._pending_student).copy_(
+                self._pending_student)
+            self.teacher_ema.resize_as_(self._pending_teacher).copy_(
+                self._pending_teacher)
+        else:
+            self.student_ema.mul_(self.beta).add_(
+                self._pending_student, alpha=1.0 - self.beta)
+            self.teacher_ema.mul_(self.beta).add_(
+                self._pending_teacher, alpha=1.0 - self.beta)
+        self.num_updates.add_(1)
+        self.discard()
+
+    def discard(self) -> None:
+        """Drop staged state when backward/optimizer is skipped."""
+        self._pending_student = None
+        self._pending_teacher = None
+
+    def get_extra_state(self):
+        state = {
+            "version": self._STATE_VERSION,
+            "beta": self.beta,
+            "num_rff_features": self.num_rff_features,
+            "bandwidth_multipliers": self.bandwidth_multipliers,
+            "seed": self.seed,
+            "row_normalize": self.row_normalize,
+        }
+        # Preserve the exact legacy rows-mode payload.  Wan v3 checkpoints
+        # contain this dict and must remain loadable without migration.
+        if self.sample_mode != "rows":
+            state["sample_mode"] = self.sample_mode
+        return state
+
+    def set_extra_state(self, state):
+        expected = self.get_extra_state()
+        if state != expected:
+            raise ValueError(
+                f"Wan KV MMD checkpoint/config mismatch: saved={state}, "
+                f"current={expected}")
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys,
+        unexpected_keys, error_msgs,
+    ):
+        for name in (
+            "rff_weight", "rff_phase", "student_ema", "teacher_ema",
+        ):
+            key = prefix + name
+            if key in state_dict:
+                getattr(self, name).resize_as_(state_dict[key])
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys,
+            unexpected_keys, error_msgs)

1.3B/code/model/__init__.py

@@ -0,0 +1,4 @@
+"""distribution matching distribution-matching model."""
+from .distribution_matching import DistributionMatchingModel, WanKVMMD
+
+__all__ = ["DistributionMatchingModel", "WanKVMMD"]

1.3B/code/model/base.py

@@ -0,0 +1,343 @@
+from typing import Tuple
+from einops import rearrange
+from torch import nn
+import torch.distributed as dist
+import torch
+
+from pipeline import SelfForcingTrainingPipeline
+from utils.loss import get_denoising_loss
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BaseModel(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self._initialize_models(args, device)
+
+        self.device = device
+        self.args = args
+        self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
+        if hasattr(args, "denoising_step_list"):
+            self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
+            if args.warp_denoising_step:
+                timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+                self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+    def _initialize_models(self, args, device):
+        self.real_model_name = getattr(args, "real_name", "Wan2.1-T2V-1.3B")
+        self.fake_model_name = getattr(args, "fake_name", "Wan2.1-T2V-1.3B")
+
+        gen_kwargs = dict(getattr(args, "model_kwargs", {}))
+        if getattr(args, "attention_type", "softmax") != "softmax":
+            raise ValueError("Only causal softmax attention is supported")
+        gen_kwargs.setdefault("attention_type", "softmax")
+        self.generator = WanDiffusionWrapper(
+            **gen_kwargs, is_causal=bool(getattr(args, "causal", True)))
+
+        self.generator.model.requires_grad_(True)
+
+        self.uses_score_models = bool(
+            getattr(args, "uses_score_models", getattr(args, "distribution_loss", None) != "fd")
+        )
+        self.uses_real_score = bool(getattr(args, "uses_real_score", self.uses_score_models))
+        self.uses_fake_score = bool(getattr(args, "uses_fake_score", self.uses_score_models))
+        if self.uses_real_score:
+            self.real_score = WanDiffusionWrapper(model_name=self.real_model_name, is_causal=False)
+            self.real_score.model.requires_grad_(False)
+        else:
+            self.real_score = None
+        if self.uses_fake_score:
+            self.fake_score = WanDiffusionWrapper(model_name=self.fake_model_name, is_causal=False)
+            self.fake_score.model.requires_grad_(True)
+        else:
+            self.fake_score = None
+
+        # Construct the frozen 5.7B T5 encoder directly in the training dtype.
+        # Building four temporary fp32 copies can otherwise exceed host memory
+        # before FSDP gets a chance to shard them.
+        text_dtype = (
+            torch.bfloat16 if bool(getattr(args, "mixed_precision", False))
+            else torch.float32
+        )
+        self.text_encoder = WanTextEncoder(dtype=text_dtype)
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+        self.scheduler = self.generator.get_scheduler()
+        self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+    def _get_timestep(
+            self,
+            min_timestep: int,
+            max_timestep: int,
+            batch_size: int,
+            num_frame: int,
+            num_frame_per_block: int,
+            uniform_timestep: bool = False
+    ) -> torch.Tensor:
+        """
+        Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
+        from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
+        - If uniform_timestep, it will use the same timestep for all frames.
+        - If not uniform_timestep, it will use a different timestep for each block.
+        """
+        if uniform_timestep:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, 1],
+                device=self.device,
+                dtype=torch.long
+            ).repeat(1, num_frame)
+            return timestep
+        else:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, num_frame],
+                device=self.device,
+                dtype=torch.long
+            )
+            # make the noise level the same within every block
+            if self.independent_first_frame:
+                # the first frame is always kept the same
+                timestep_from_second = timestep[:, 1:]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1, num_frame_per_block)
+                timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1)
+                timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
+            else:
+                timestep = timestep.reshape(
+                    timestep.shape[0], -1, num_frame_per_block)
+                timestep[:, :, 1:] = timestep[:, :, 0:1]
+                timestep = timestep.reshape(timestep.shape[0], -1)
+            return timestep
+
+
+class SelfForcingModel(BaseModel):
+    def __init__(self, args, device):
+        super().__init__(args, device)
+        self.denoising_loss_func = get_denoising_loss(args.denoising_loss_type)()
+
+    def _sample_generator_input_descriptor(
+        self,
+        image_or_video_shape,
+        initial_latent: torch.Tensor = None,
+    ) -> dict:
+        """Explicitly sample `_run_generator`'s block count and initial noise.
+
+        The returned descriptor is reusable: passing it back to
+        ``_run_generator`` avoids consulting the ambient RNG for those direct
+        inputs.  A caller replaying the full rollout must additionally capture
+        the post-sampling RNG state used by stochastic pipeline internals.
+        """
+        if self.args.i2v:
+            noise_shape = [
+                image_or_video_shape[0],
+                image_or_video_shape[1] - 1,
+                *image_or_video_shape[2:],
+            ]
+        else:
+            noise_shape = list(image_or_video_shape)
+
+        min_num_frames = 20 if self.args.independent_first_frame else 21
+        max_num_frames = (
+            self.num_training_frames - 1
+            if self.args.independent_first_frame
+            else self.num_training_frames
+        )
+        if max_num_frames % self.num_frame_per_block != 0:
+            raise ValueError(
+                "num_training_frames is not divisible by "
+                "num_frame_per_block")
+        if min_num_frames % self.num_frame_per_block != 0:
+            raise ValueError(
+                "minimum generated frames is not divisible by "
+                "num_frame_per_block")
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        block_count = torch.randint(
+            min_num_blocks,
+            max_num_blocks + 1,
+            (1,),
+            device=self.device,
+        )
+        if dist.is_available() and dist.is_initialized():
+            dist.broadcast(block_count, src=0)
+        num_generated_blocks = int(block_count.item())
+        num_generated_frames = (
+            num_generated_blocks * self.num_frame_per_block)
+        if self.args.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+        noise_shape[1] = num_generated_frames
+        return {
+            "version": 1,
+            "num_generated_blocks": num_generated_blocks,
+            "noise": torch.randn(
+                noise_shape, device=self.device, dtype=self.dtype),
+        }
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        initial_latent: torch.tensor = None,
+        generation_descriptor: dict = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+            - denoised_timestep: an integer
+        """
+        # Step 1: Sample noise and backward simulate the generator's input
+        assert getattr(self.args, "backward_simulation", True), "Backward simulation needs to be enabled"
+        if initial_latent is not None:
+            conditional_dict["initial_latent"] = initial_latent
+        # During training, the number of generated frames should be uniformly sampled from
+        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
+        min_num_frames = 20 if self.args.independent_first_frame else 21
+        max_num_frames = self.num_training_frames - 1 if self.args.independent_first_frame else self.num_training_frames
+        assert max_num_frames % self.num_frame_per_block == 0
+        assert min_num_frames % self.num_frame_per_block == 0
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        if generation_descriptor is None:
+            generation_descriptor = (
+                self._sample_generator_input_descriptor(
+                    image_or_video_shape,
+                    initial_latent=initial_latent,
+                )
+            )
+        if (
+            not isinstance(generation_descriptor, dict)
+            or set(generation_descriptor)
+            != {"version", "num_generated_blocks", "noise"}
+            or int(generation_descriptor["version"]) != 1
+        ):
+            raise ValueError("invalid explicit generator input descriptor")
+        num_generated_blocks = int(
+            generation_descriptor["num_generated_blocks"])
+        if not min_num_blocks <= num_generated_blocks <= max_num_blocks:
+            raise ValueError(
+                "generator descriptor block count is outside the configured "
+                f"range [{min_num_blocks}, {max_num_blocks}]")
+        if dist.is_available() and dist.is_initialized():
+            block_min = torch.tensor(
+                num_generated_blocks, device=self.device, dtype=torch.long)
+            block_max = block_min.clone()
+            dist.all_reduce(block_min, op=dist.ReduceOp.MIN)
+            dist.all_reduce(block_max, op=dist.ReduceOp.MAX)
+            if int(block_min.item()) != int(block_max.item()):
+                raise ValueError(
+                    "generator descriptor block count differs across ranks")
+        num_generated_frames = num_generated_blocks * self.num_frame_per_block
+        if self.args.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+            min_num_frames += 1
+        expected_noise_shape = list(image_or_video_shape)
+        if self.args.i2v:
+            expected_noise_shape[1] -= 1
+        expected_noise_shape[1] = num_generated_frames
+        generation_noise = generation_descriptor["noise"]
+        expected_device = (
+            torch.device("cuda", int(self.device))
+            if isinstance(self.device, int)
+            else torch.device(self.device)
+        )
+        if (
+            not isinstance(generation_noise, torch.Tensor)
+            or tuple(generation_noise.shape) != tuple(expected_noise_shape)
+            or generation_noise.device != expected_device
+            or generation_noise.dtype != self.dtype
+            or not bool(torch.isfinite(generation_noise).all().item())
+        ):
+            raise ValueError(
+                "generator descriptor noise is incompatible with the "
+                f"expected {tuple(expected_noise_shape)} {self.dtype} on "
+                f"{expected_device}")
+
+        pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(
+            noise=generation_noise,
+            **conditional_dict,
+        )
+        # Slice last 21 frames
+        if pred_image_or_video.shape[1] > 21:
+            with torch.no_grad():
+                # Reencode to get image latent
+                latent_to_decode = pred_image_or_video[:, :-20, ...]
+                # Deccode to video
+                pixels = self.vae.decode_to_pixel(latent_to_decode)
+                frame = pixels[:, -1:, ...].to(self.dtype)
+                frame = rearrange(frame, "b t c h w -> b c t h w")
+                # Encode frame to get image latent
+                image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
+            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
+        else:
+            pred_image_or_video_last_21 = pred_image_or_video
+
+        if num_generated_frames != min_num_frames:
+            # Currently, we do not use gradient for the first chunk, since it contains image latents
+            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
+            if self.args.independent_first_frame:
+                gradient_mask[:, :1] = False
+            else:
+                gradient_mask[:, :self.num_frame_per_block] = False
+        else:
+            gradient_mask = None
+
+        pred_image_or_video_last_21 = pred_image_or_video_last_21.to(self.dtype)
+        return pred_image_or_video_last_21, gradient_mask, denoised_timestep_from, denoised_timestep_to
+
+    def _consistency_backward_simulation(
+        self,
+        noise: torch.Tensor,
+        **conditional_dict: dict
+    ) -> torch.Tensor:
+        """
+        Simulate the generator's input from noise to avoid training/inference mismatch.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
+        Input:
+            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - output: a tensor with shape [B, T, F, C, H, W].
+            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
+            represents the x0 prediction at each timestep.
+        """
+        if self.inference_pipeline is None:
+            self._initialize_inference_pipeline()
+
+        return self.inference_pipeline.inference_with_trajectory(
+            noise=noise, **conditional_dict
+        )
+
+    def _initialize_inference_pipeline(self):
+        """
+        Lazy initialize the inference pipeline during the first backward simulation run.
+        Here we encapsulate the inference code with a model-dependent outside function.
+        We pass our FSDP-wrapped modules into the pipeline to save memory.
+        """
+        self.inference_pipeline = SelfForcingTrainingPipeline(
+            denoising_step_list=self.denoising_step_list,
+            scheduler=self.scheduler,
+            generator=self.generator,
+            num_frame_per_block=self.num_frame_per_block,
+            independent_first_frame=self.args.independent_first_frame,
+            same_step_across_blocks=self.args.same_step_across_blocks,
+            last_step_only=self.args.last_step_only,
+            num_max_frames=self.num_training_frames,
+            context_noise=self.args.context_noise
+        )

1.3B/code/model/generation.py

@@ -0,0 +1,142 @@
+"""Shared generator setup and few-step rollout utilities.
+
+This module contains only generator behavior used by distribution matching.  Loss
+construction lives in :mod:`model.distribution_matching`.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+
+from model.base import SelfForcingModel
+
+
+class FewStepGenerationModel(SelfForcingModel):
+    """Initialize the causal generator and expose deterministic replay helpers."""
+
+    def __init__(self, args, device):
+        super().__init__(args, device)
+        self.causal_generator = bool(getattr(args, "causal", True))
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(
+            args, "same_step_across_blocks", True)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.causal_generator and self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(
+            args, "independent_first_frame", False)
+        if self.causal_generator and self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        self.inference_pipeline = None
+        self.guidance_scale = float(getattr(args, "guidance_scale", 1.0))
+
+        if not hasattr(self, "denoising_step_list"):
+            raise ValueError(
+                "few-step generation requires denoising_step_list")
+        denoising_steps = self.denoising_step_list.to(device=device)
+        if denoising_steps.numel() > 0 and denoising_steps[-1].item() == 0:
+            denoising_steps = denoising_steps[:-1]
+        self.generation_denoising_step_list = denoising_steps
+
+    @staticmethod
+    def _apply_gradient_mask(
+        latents: torch.Tensor,
+        gradient_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if gradient_mask is None:
+            return latents
+        frame_mask = gradient_mask.flatten(2).any(dim=2).view(
+            latents.shape[0], latents.shape[1], 1, 1, 1)
+        return torch.where(frame_mask, latents, latents.detach())
+
+    def _predict_x0_with_cfg(
+        self,
+        noisy_latent: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        _, x0_cond = self.generator(
+            noisy_image_or_video=noisy_latent,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+        if self.guidance_scale == 1.0:
+            return x0_cond
+
+        _, x0_uncond = self.generator(
+            noisy_image_or_video=noisy_latent,
+            conditional_dict=unconditional_dict,
+            timestep=timestep,
+        )
+        return x0_uncond + self.guidance_scale * (x0_cond - x0_uncond)
+
+    def _run_few_step_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        exit_index: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Run detached rollout steps followed by one differentiated exit step."""
+        latent = torch.randn(
+            image_or_video_shape,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        batch_size, num_frames = latent.shape[:2]
+        denoising_steps = self.generation_denoising_step_list.to(
+            device=self.device)
+        num_steps = len(denoising_steps)
+
+        if exit_index is None:
+            if num_steps <= 1:
+                exit_index = 0
+            else:
+                index_tensor = torch.randint(
+                    0, num_steps, (1,), device=self.device)
+                if dist.is_available() and dist.is_initialized():
+                    dist.broadcast(index_tensor, src=0)
+                exit_index = int(index_tensor.item())
+
+        for index in range(num_steps):
+            current_timestep = denoising_steps[index]
+            timestep = torch.ones(
+                [batch_size, num_frames],
+                device=self.device,
+                dtype=current_timestep.dtype,
+            ) * current_timestep
+            if index < exit_index:
+                with torch.no_grad():
+                    x0_pred = self._predict_x0_with_cfg(
+                        noisy_latent=latent,
+                        conditional_dict=conditional_dict,
+                        unconditional_dict=unconditional_dict,
+                        timestep=timestep,
+                    )
+                    next_timestep = denoising_steps[index + 1]
+                    latent = self.scheduler.add_noise(
+                        x0_pred.flatten(0, 1),
+                        torch.randn_like(x0_pred.flatten(0, 1)),
+                        next_timestep * torch.ones(
+                            [batch_size * num_frames],
+                            device=self.device,
+                            dtype=next_timestep.dtype,
+                        ),
+                    ).unflatten(0, x0_pred.shape[:2]).to(self.dtype)
+            else:
+                return self._predict_x0_with_cfg(
+                    noisy_latent=latent,
+                    conditional_dict=conditional_dict,
+                    unconditional_dict=unconditional_dict,
+                    timestep=timestep,
+                )
+        return latent

1.3B/code/model/wan_kv_mmd.py

@@ -0,0 +1,6 @@
+"""Compatibility import; read model/distribution_matching.py for the implementation."""
+import sys
+from . import distribution_matching as _implementation
+# 这是一个历史遗留问题的转发窗口
+# Share the module object so existing monkeypatches and private imports still work.
+sys.modules[__name__] = _implementation

1.3B/code/pipeline/__init__.py

@@ -0,0 +1,4 @@
+from .causal_inference import CausalInferencePipeline
+from .self_forcing_training import SelfForcingTrainingPipeline
+
+__all__ = ["CausalInferencePipeline", "SelfForcingTrainingPipeline"]

1.3B/code/pipeline/causal_inference.py

@@ -0,0 +1,472 @@
+from typing import List, Optional
+import torch
+
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation
+
+
+class CausalInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        if generator is None:
+            gen_kwargs = dict(getattr(args, "model_kwargs", {}))
+            if getattr(args, "attention_type", "softmax") != "softmax":
+                raise ValueError("Only causal softmax attention is supported")
+            gen_kwargs.setdefault("attention_type", "softmax")
+            self.generator = WanDiffusionWrapper(**gen_kwargs, is_causal=True)
+        else:
+            self.generator = generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize all causal hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.denoising_step_list = torch.tensor(
+            args.denoising_step_list, dtype=torch.long)
+        if args.warp_denoising_step:
+            timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+
+        self.attention_type = getattr(self.generator, "attention_type", "softmax")
+        self.kv_cache1 = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        # Linear-attention student does not have a (sliding) ``local_attn_size``;
+        # treat it as unbounded so the recurrent state is allocated for the full
+        # video.
+        self.local_attn_size = getattr(self.generator.model, "local_attn_size", -1)
+
+        print(
+            f"KV inference with {self.num_frame_per_block} frames per block "
+            f"(attention_type={self.attention_type})"
+        )
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def _set_advance_state(self, value: bool) -> None:
+        """Toggle ``advance_state`` on every linear/hybrid kv-cache dict.
+
+        See ``SelfForcingTrainingPipeline._set_advance_state``: the flag
+        controls whether the linear / hybrid recurrent state is updated in
+        place on a given forward.  Pipelines set ``False`` for every
+        denoising-loop pass and ``True`` only on the clean re-run so each
+        block contributes EXACTLY ONCE to the persistent state, matching
+        the chunkwise training semantics and aligning train / inference
+        cache distributions.  No-op for softmax caches.
+        """
+        if self.kv_cache1 is None:
+            return
+        for cache in self.kv_cache1:
+            if "state" in cache:
+                cache["advance_state"] = value
+            # ARL2 HybridBlockAttention layers: clean-pass GDN update flag.
+            if "gdn_state" in cache:
+                cache["update_state"] = value
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        initial_latent: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        profile: bool = False,
+        low_memory: bool = False,
+        decode: bool = True,
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+                It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        conditional_dict = self.text_encoder(
+            text_prompts=text_prompts
+        )
+
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(self.text_encoder, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Set up profiling if requested
+        if profile:
+            init_start = torch.cuda.Event(enable_timing=True)
+            init_end = torch.cuda.Event(enable_timing=True)
+            diffusion_start = torch.cuda.Event(enable_timing=True)
+            diffusion_end = torch.cuda.Event(enable_timing=True)
+            vae_start = torch.cuda.Event(enable_timing=True)
+            vae_end = torch.cuda.Event(enable_timing=True)
+            block_times = []
+            block_start = torch.cuda.Event(enable_timing=True)
+            block_end = torch.cuda.Event(enable_timing=True)
+            init_start.record()
+
+        # Step 1: Initialize KV cache to all zeros
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache1)):
+                cache = self.kv_cache1[block_index]
+                cache["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                cache["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                # Linear/hybrid layers carry a recurrent state; softmax
+                # layers carry a K/V cache.  In a MIXED model each layer
+                # has only one of the two -- zero whichever is present.
+                if "state" in cache:
+                    cache["state"].zero_()
+                if "normalizer" in cache:
+                    cache["normalizer"].zero_()
+                if "k" in cache:
+                    cache["k"].zero_()
+                if "v" in cache:
+                    cache["v"].zero_()
+                # ARL2 GDN layers: drop the leaked recurrent state so the next
+                # video starts from a clean (lazily re-allocated) state.
+                if "gdn_state" in cache:
+                    cache["gdn_state"] = None
+                    cache["update_state"] = True
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            # Initial-latent context forwards write clean K/V into the cache
+            # exactly once per block -> advance_state=True.
+            self._set_advance_state(True)
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = \
+                    initial_latent[:, current_start_frame:current_start_frame + self.num_frame_per_block]
+                output[:, current_start_frame:current_start_frame + self.num_frame_per_block] = current_ref_latents
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += self.num_frame_per_block
+
+        if profile:
+            init_end.record()
+            torch.cuda.synchronize()
+            diffusion_start.record()
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_num_frames in all_num_frames:
+            if profile:
+                block_start.record()
+
+            noisy_input = noise[
+                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # All denoising-loop forwards on this block read the same cross-block
+            # cache and do NOT advance it; only the clean re-run (Step 3.3)
+            # advances the cache.  No-op for softmax caches.
+            self._set_advance_state(False)
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                print(f"current_timestep: {current_timestep}")
+                # set current timestep
+                timestep = torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if index < len(self.denoising_step_list) - 1:
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+                    next_timestep = self.denoising_step_list[index + 1]
+                    noisy_input = self.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep * torch.ones(
+                            [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                    ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    # for getting real output
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            context_timestep = torch.ones_like(timestep) * self.args.context_noise
+            # This is the single write the linear/hybrid cache gets for this
+            # block -- restore the default ``advance_state=True`` before the
+            # forward so the clean K/V actually lands in the cache.
+            self._set_advance_state(True)
+            self.generator(
+                noisy_image_or_video=denoised_pred,
+                conditional_dict=conditional_dict,
+                timestep=context_timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=current_start_frame * self.frame_seq_length,
+            )
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        if profile:
+            # End diffusion timing and synchronize CUDA
+            diffusion_end.record()
+            torch.cuda.synchronize()
+            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
+            init_time = init_start.elapsed_time(init_end)
+            vae_start.record()
+
+        # Step 4: Decode the output (skippable for latent-only generation, e.g.
+        # ARL² Stage-2 teacher-video latent synthesis where pixels are not needed)
+        if not decode:
+            return (None, output) if return_latents else output
+        video = self.vae.decode_to_pixel(output, use_cache=False)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        if profile:
+            # End VAE timing and synchronize CUDA
+            vae_end.record()
+            torch.cuda.synchronize()
+            vae_time = vae_start.elapsed_time(vae_end)
+            total_time = init_time + diffusion_time + vae_time
+
+            print("Profiling results:")
+            print(f"  - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)")
+            print(f"  - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)")
+            for i, block_time in enumerate(block_times):
+                print(f"    - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)")
+            print(f"  - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)")
+            print(f"  - Total time: {total_time:.2f} ms")
+
+        if return_latents:
+            return video, output
+        else:
+            return video
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+
+        Per-layer cache shape depends on that layer's attention type:
+
+        * ``softmax`` layers carry a sliding-window K/V tensor
+          ``[B, kv_cache_size, H, Dv]`` in the activation dtype (bf16).
+          The cache size matches the all-softmax baseline.
+        * ``linear`` / ``hybrid`` layers carry a recurrent state
+          ``[B, H, M, Dv]`` plus denominator ``[B, H, M]`` in fp32 for
+          numerical stability over the long causal rollout.
+
+        For a MIXED model (some softmax, some linear/hybrid) we walk the
+        ``per_layer_self_attn_types`` list and pick the right shape per
+        layer.  Pure-linear / pure-hybrid models use the same code path
+        with every layer set to ``"linear_or_hybrid"``.
+        """
+        kv_cache1 = []
+        model = self.generator.model
+        num_heads = getattr(model, "num_heads", 12)
+        head_dim = model.dim // num_heads
+
+        # ── ARL2 path: softmax CausalWanModel with some blocks swapped to
+        # HybridBlockAttention (GDN).  GDN layers use a lazily-allocated
+        # gdn_state (+ update_state flag); other layers keep K/V.
+        def is_gdn_block(block):
+            cls = type(block.self_attn)
+            return (
+                cls.__name__ == "HybridBlockAttention"
+                and cls.__module__ == "wan.modules.hybrid_attention"
+            )
+
+        if any(is_gdn_block(b) for b in model.blocks):
+            sm_size = (self.local_attn_size * self.frame_seq_length
+                       if self.local_attn_size != -1 else 32760)
+            for b in model.blocks:
+                entry = {
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+                if is_gdn_block(b):
+                    entry["gdn_state"] = None
+                    entry["update_state"] = True
+                else:
+                    entry["k"] = torch.zeros([batch_size, sm_size, num_heads, head_dim], dtype=dtype, device=device)
+                    entry["v"] = torch.zeros([batch_size, sm_size, num_heads, head_dim], dtype=dtype, device=device)
+                kv_cache1.append(entry)
+            self.kv_cache1 = kv_cache1
+            return
+
+        if self.attention_type in ("linear", "hybrid"):
+            # Default: every layer is linear/hybrid.  A mixed model exposes
+            # an explicit per-layer types list that overrides this.
+            per_layer_types = getattr(
+                model, "per_layer_self_attn_types",
+                ["linear_or_hybrid"] * self.num_transformer_blocks,
+            )
+            # Softmax-layer KV cache size mirrors the all-softmax baseline.
+            softmax_cache_size = (
+                self.local_attn_size * self.frame_seq_length
+                if self.local_attn_size != -1
+                else 32760
+            )
+            for layer_idx in range(self.num_transformer_blocks):
+                ltype = per_layer_types[layer_idx] if layer_idx < len(per_layer_types) else "linear_or_hybrid"
+                if ltype == "softmax":
+                    kv_cache1.append({
+                        "k": torch.zeros(
+                            [batch_size, softmax_cache_size, num_heads, head_dim],
+                            dtype=dtype, device=device,
+                        ),
+                        "v": torch.zeros(
+                            [batch_size, softmax_cache_size, num_heads, head_dim],
+                            dtype=dtype, device=device,
+                        ),
+                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    })
+                else:
+                    # linear / hybrid recurrent state.  Kernel dim defaults
+                    # to head_dim but may be overridden via num_kernel_features.
+                    kernel_dim = head_dim
+                    try:
+                        first_block = model.blocks[layer_idx]
+                        kernel_dim = first_block.self_attn.kernel_dim
+                    except (AttributeError, IndexError):
+                        pass
+                    kv_cache1.append({
+                        "state": torch.zeros(
+                            [batch_size, num_heads, kernel_dim, head_dim],
+                            dtype=torch.float32, device=device,
+                        ),
+                        "normalizer": torch.zeros(
+                            [batch_size, num_heads, kernel_dim],
+                            dtype=torch.float32, device=device,
+                        ),
+                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    })
+        else:
+            # All-softmax model (CausalWanModel path; no mixing).
+            if self.local_attn_size != -1:
+                kv_cache_size = self.local_attn_size * self.frame_seq_length
+            else:
+                kv_cache_size = 32760
+            for _ in range(self.num_transformer_blocks):
+                kv_cache1.append({
+                    "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+                })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+        self.crossattn_cache = crossattn_cache

1.3B/code/pipeline/self_forcing_training.py

@@ -0,0 +1,406 @@
+"""Autoregressive training rollouts with a sampled denoising exit step.
+
+Each block conditions on earlier generated blocks. Training stops at a sampled
+exit and differentiates that denoising call when gradients are enabled. Released
+four-step inference runs the complete schedule in ``causal_inference.py``.
+"""
+
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class SelfForcingTrainingPipeline:
+    """Generate on-policy training videos and reproduce their rollout trajectory."""
+
+    def __init__(self,
+                 denoising_step_list: List[int],
+                 scheduler: SchedulerInterface,
+                 generator: WanDiffusionWrapper,
+                 num_frame_per_block=3,
+                 independent_first_frame: bool = False,
+                 same_step_across_blocks: bool = False,
+                 last_step_only: bool = False,
+                 num_max_frames: int = 21,
+                 context_noise: int = 0,
+                 **kwargs):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+
+        # Wan specific hyperparameters
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+
+    def _set_advance_state(self, value: bool) -> None:
+        """Toggle the ``advance_state`` flag on every layer's kv-cache dict.
+
+        Only the linear / hybrid attention recurrent caches read this flag
+        (their kv-cache dicts have ``"state"`` and ``"normalizer"`` fields);
+        the softmax cache uses overwrite-style writes, so this is a no-op
+        for those layers.  See ``wan/modules/linear_attention.py`` for the
+        rationale (avoid per-pass over-accumulation across the denoising
+        loop's multiple forwards on the same block).
+        """
+        if self.kv_cache1 is None:
+            return
+        for cache in self.kv_cache1:
+            if "state" in cache:
+                cache["advance_state"] = value
+            # ARL2 HybridBlockAttention layers use a GDN recurrent state and
+            # read ``update_state`` instead of ``advance_state`` (clean-pass
+            # update: skip writes during denoising, write on the clean rerun).
+            if "gdn_state" in cache:
+                cache["update_state"] = value
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        """Choose denoising exit indices on rank zero and share them across ranks."""
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Sampling an exit bounds the training graph to one denoising call
+            # per block; last_step_only optionally fixes it to the final call.
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def inference_with_trajectory(
+            self,
+            noise: torch.Tensor,
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        """Return a training rollout plus its selected denoising interval.
+
+        During feature collection the caller disables gradients. During replay,
+        the same random choices reconstruct the selected calls with gradients.
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        self._initialize_kv_cache(
+            batch_size=batch_size, dtype=noise.dtype, device=noise.device
+        )
+        self._initialize_crossattn_cache(
+            batch_size=batch_size, dtype=noise.dtype, device=noise.device
+        )
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+            output[:, :1] = initial_latent
+            # Initial-latent context forward is a legit "clean K/V" write -> advance cache.
+            self._set_advance_state(True)
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=initial_latent,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+            current_start_frame += 1
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        # Retain generator gradients only for the final 21 latent frames.
+        # Earlier context still affects the rollout through cached values.
+        start_gradient_frame_index = num_output_frames - 21
+
+        for block_index, current_num_frames in enumerate(all_num_frames):
+            noisy_input = noise[
+                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # All denoising-loop forwards on this block read the same cross-block
+            # cache and do NOT advance it; only the clean re-run (Step 3.3)
+            # advances the cache.  No-op for the softmax attention path.
+            self._set_advance_state(False)
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                if self.same_step_across_blocks:
+                    exit_flag = (index == exit_flags[0])
+                else:
+                    exit_flag = (index == exit_flags[block_index])  # Only backprop at the randomly selected timestep (consistent across all ranks)
+                timestep = torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if not exit_flag:
+                    # Earlier denoising calls construct the selected step's
+                    # input; their graphs are deliberately not retained.
+                    with torch.no_grad():
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        )
+                        next_timestep = self.denoising_step_list[index + 1]
+                        noisy_input = self.scheduler.add_noise(
+                            denoised_pred.flatten(0, 1),
+                            torch.randn_like(denoised_pred.flatten(0, 1)),
+                            next_timestep * torch.ones(
+                                [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                        ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    # The selected exit produces the block's clean prediction.
+                    # Replay differentiates this call within the gradient window.
+                    if current_start_frame < start_gradient_frame_index:
+                        with torch.no_grad():
+                            _, denoised_pred = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=conditional_dict,
+                                timestep=timestep,
+                                kv_cache=self.kv_cache1,
+                                crossattn_cache=self.crossattn_cache,
+                                current_start=current_start_frame * self.frame_seq_length
+                            )
+                    else:
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        )
+                    break
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: publish this generated block as context for the next.
+            # This cache write has no graph, preserving the training truncation.
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            # add context noise
+            denoised_pred = self.scheduler.add_noise(
+                denoised_pred.flatten(0, 1),
+                torch.randn_like(denoised_pred.flatten(0, 1)),
+                context_timestep * torch.ones(
+                    [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+            ).unflatten(0, denoised_pred.shape[:2])
+            # This is the single write the linear/hybrid cache gets for this
+            # block -- restore the default ``advance_state=True`` before the
+            # forward so the clean K/V actually lands in the cache.
+            self._set_advance_state(True)
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        # Step 3.5: Return the denoised timestep
+        if not self.same_step_across_blocks:
+            denoised_timestep_from, denoised_timestep_to = None, None
+        elif exit_flags[0] == len(self.denoising_step_list) - 1:
+            denoised_timestep_to = 0
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        else:
+            denoised_timestep_to = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+
+        if return_sim_step:
+            return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+
+        Per-layer cache shape depends on that layer's attention type
+        (mirrors :meth:`pipeline.causal_inference.CausalInferencePipeline._initialize_kv_cache`):
+
+          * ``softmax`` layers          -> ``k`` / ``v`` in bf16, shape
+                                           ``[B, cache_size, H, Dv]``.
+          * ``linear`` / ``hybrid``     -> ``state`` / ``normalizer`` in fp32,
+                                           shape ``[B, H, M, Dv]`` / ``[B, H, M]``.
+
+        For a MIXED model (some softmax + some linear/hybrid layers) we
+        walk the model's ``per_layer_self_attn_types`` list and allocate
+        the matching shape per layer.  Pure-linear / pure-hybrid /
+        pure-softmax models still work because the type list is uniform
+        in those cases.
+
+        Crashing on the wrong key (the previous "always-linear" allocator
+        handed a softmax layer a dict without ``"k"``, raising
+        ``KeyError: 'k'`` deep inside the attention forward) was exactly
+        the failure mode this fix addresses.
+        """
+        attention_type = getattr(self.generator, "attention_type", "softmax")
+        model = self.generator.model
+        num_heads = getattr(model, "num_heads", 12)
+        head_dim = model.dim // num_heads
+        kv_cache1 = []
+
+        # ── ARL2 path: a softmax CausalWanModel with some blocks' self_attn
+        # replaced by HybridBlockAttention.  Those blocks use a lazily-
+        # allocated GDN state (gdn_state/update_state); the rest keep K/V.
+        def is_gdn_block(block):
+            cls = type(block.self_attn)
+            return (
+                cls.__name__ == "HybridBlockAttention"
+                and cls.__module__ == "wan.modules.hybrid_attention"
+            )
+
+        if any(is_gdn_block(b) for b in model.blocks):
+            for b in model.blocks:
+                entry = {
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+                if is_gdn_block(b):
+                    entry["gdn_state"] = None        # lazily allocated H x D x D
+                    entry["update_state"] = True
+                else:
+                    entry["k"] = torch.zeros(
+                        [batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device)
+                    entry["v"] = torch.zeros(
+                        [batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device)
+                kv_cache1.append(entry)
+            self.kv_cache1 = kv_cache1
+            return
+
+        if attention_type in ("linear", "hybrid"):
+            # Default to "every layer is linear/hybrid"; mixed models
+            # override via ``per_layer_self_attn_types``.
+            per_layer_types = getattr(
+                model, "per_layer_self_attn_types",
+                ["linear_or_hybrid"] * self.num_transformer_blocks,
+            )
+            # Softmax-layer cache size.  Use the model's
+            # ``softmax_layer_local_attn_size`` knob so a windowed config
+            # only allocates as much storage as it actually uses.
+            sm_local = int(getattr(model, "softmax_layer_local_attn_size", -1))
+            if sm_local != -1:
+                softmax_cache_size = sm_local * self.frame_seq_length
+            else:
+                softmax_cache_size = self.kv_cache_size
+            for layer_idx in range(self.num_transformer_blocks):
+                ltype = (
+                    per_layer_types[layer_idx]
+                    if layer_idx < len(per_layer_types)
+                    else "linear_or_hybrid"
+                )
+                if ltype == "softmax":
+                    kv_cache1.append({
+                        "k": torch.zeros(
+                            [batch_size, softmax_cache_size, num_heads, head_dim],
+                            dtype=dtype, device=device,
+                        ),
+                        "v": torch.zeros(
+                            [batch_size, softmax_cache_size, num_heads, head_dim],
+                            dtype=dtype, device=device,
+                        ),
+                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    })
+                else:
+                    # Linear/hybrid recurrent state.  ``kernel_dim`` may
+                    # be overridden via ``num_kernel_features``; default
+                    # is ``head_dim`` (= 128 for Wan-1.3B).
+                    kernel_dim = head_dim
+                    try:
+                        first_linhyb_block = model.blocks[layer_idx]
+                        kernel_dim = first_linhyb_block.self_attn.kernel_dim
+                    except (AttributeError, IndexError):
+                        pass
+                    kv_cache1.append({
+                        "state": torch.zeros(
+                            [batch_size, num_heads, kernel_dim, head_dim],
+                            dtype=torch.float32, device=device,
+                        ),
+                        "normalizer": torch.zeros(
+                            [batch_size, num_heads, kernel_dim],
+                            dtype=torch.float32, device=device,
+                        ),
+                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    })
+        else:
+            # All-softmax (Wan baseline) path.
+            for _ in range(self.num_transformer_blocks):
+                kv_cache1.append({
+                    "k": torch.zeros([batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+                })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+        self.crossattn_cache = crossattn_cache

1.3B/code/prompts/test_8.txt

@@ -0,0 +1,8 @@
+A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.
+Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees and dramatic snow capped mountains in the distance, mid afternoon light with wispy clouds and a sun high in the distance creates a warm glow, the low camera view is stunning capturing the large furry mammal with beautiful photography, depth of field.
+A movie trailer featuring the adventures of the 30 year old space man wearing a red wool knitted motorcycle helmet, blue sky, salt desert, cinematic style, shot on 35mm film, vivid colors.
+Drone view of waves crashing against the rugged cliffs along Big Sur’s garay point beach. The crashing blue waters create white-tipped waves, while the golden light of the setting sun illuminates the rocky shore. A small island with a lighthouse sits in the distance, and green shrubbery covers the cliff’s edge. The steep drop from the road down to the beach is a dramatic feat, with the cliff’s edges jutting out over the sea. This is a view that captures the raw beauty of the coast and the rugged landscape of the Pacific Coast Highway.
+Animated scene features a close-up of a short fluffy monster kneeling beside a melting red candle. The art style is 3D and realistic, with a focus on lighting and texture. The mood of the painting is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image.
+A gorgeously rendered papercraft world of a coral reef, rife with colorful fish and sea creatures.
+This close-up shot of a Victoria crowned pigeon showcases its striking blue plumage and red chest. Its crest is made of delicate, lacy feathers, while its eye is a striking red color. The bird’s head is tilted slightly to the side, giving the impression of it looking regal and majestic. The background is blurred, drawing attention to the bird’s striking appearance.
+Photorealistic closeup video of two pirate ships battling each other as they sail inside a cup of coffee.

1.3B/code/tests/test_multi_representation_video_nystrom_mmd.py

@@ -0,0 +1,238 @@
+import torch
+
+from losses.multi_representation_video_nystrom_mmd import (
+    MultiRepresentationVideoNystromMMD,
+    fd_normalize_raw_mmd,
+)
+from losses.video_nystrom_mmd import (
+    VideoNystromMMD,
+    build_video_nystrom_reference,
+)
+from trainer.multi_representation_fresh_video_gradcache import (
+    MultiRepresentationFreshVideoGradCache,
+)
+
+
+def _objective(name: str, shape: tuple[int, int], seed: int):
+    generator = torch.Generator().manual_seed(seed)
+    reference = torch.randn(12, *shape, generator=generator)
+    pack = build_video_nystrom_reference(
+        reference,
+        num_landmarks=5,
+        metadata={
+            "feature": {"name": name, "matrix_shape": list(shape)},
+            "reference": {"dataset": "triple-unit-test"},
+        },
+        bandwidth=1.7,
+        kmeans_iterations=3,
+        ridge=1.0e-5,
+        chunk_size=16,
+    )
+    return VideoNystromMMD(
+        pack,
+        kernel_chunk_size=64,
+        attraction_weight=2.0,
+        current_teacher_attraction_enabled=True,
+        current_teacher_attraction_weight=1.0,
+    )
+
+
+def _triple():
+    return MultiRepresentationVideoNystromMMD(
+        {
+            "dinov3": _objective("tiny-dinov3", (1, 3), 4199),
+            "vjepa2": _objective("tiny-vjepa2", (1, 2), 4200),
+            "videomae_v2": _objective("tiny-videomae", (2, 3), 4201),
+        },
+        weights={"dinov3": 15.0, "vjepa2": 1.0, "videomae_v2": 2.0},
+        normalization_eps=1.0e-4,
+        aggregation_divisor=10.0,
+    )
+
+
+def test_three_raw_mmds_are_normalized_weighted_then_divided_by_ten():
+    objective = _triple()
+    generator = torch.Generator().manual_seed(4202)
+    reference = {
+        "dinov3": torch.randn(256, 1, 3, generator=generator),
+        "vjepa2": torch.randn(256, 1, 2, generator=generator),
+        "videomae_v2": torch.randn(256, 2, 3, generator=generator),
+    }
+    indices = torch.arange(128, 256)
+    gradient = {
+        name: value.index_select(0, indices).clone().requires_grad_(True)
+        for name, value in reference.items()
+    }
+    teacher = {
+        "dinov3": torch.randn(
+            128, 1, 3, generator=generator, requires_grad=True),
+        "vjepa2": torch.randn(
+            128, 1, 2, generator=generator, requires_grad=True),
+        "videomae_v2": torch.randn(
+            128, 2, 3, generator=generator, requires_grad=True),
+    }
+    loss, metadata = objective(
+        gradient_students=gradient,
+        student_reference=reference,
+        current_teacher=teacher,
+        gradient_student_indices=indices,
+        stage=False,
+    )
+
+    expected_terms = []
+    for name, branch in objective.branches.items():
+        raw, branch_metadata = branch(
+            gradient_students=gradient[name],
+            student_reference=reference[name],
+            current_teacher=teacher[name],
+            gradient_student_indices=indices,
+            stage=False,
+        )
+        normalized = raw / (raw.detach() + 1.0e-4)
+        expected_terms.append(objective.branch_weights[name] * normalized)
+        observed = metadata["branches"][name]
+        torch.testing.assert_close(observed["raw_mmd"], raw.detach())
+        assert int(observed["repulsion_denominator"]) == 65536
+        assert int(observed["fixed_teacher_attraction_denominator"]) == 256
+        assert int(observed["current_teacher_attraction_denominator"]) == 32768
+        assert int(observed["current_teacher_self_kernel_denominator"]) == 16384
+        assert int(observed["full_student_partial_gradient"]) == 1
+        assert abs(float(observed[
+            "fixed_teacher_attraction_normalized_weight"]) - 2.0 / 3.0) < 1e-7
+        assert abs(float(observed[
+            "current_teacher_attraction_normalized_weight"]) - 1.0 / 3.0) < 1e-7
+        assert branch_metadata["repulsion_denominator"].item() == 65536
+    expected = torch.stack(expected_terms).sum() / 10.0
+    torch.testing.assert_close(loss, expected, rtol=2.0e-5, atol=2.0e-6)
+    # Aggregate exactly as (15 * DINOv3 + V-JEPA2 + 2 * VideoMAE-v2) / 10.
+    torch.testing.assert_close(
+        metadata["normalized_mmd_sum"], expected.detach(),
+        rtol=2.0e-5, atol=2.0e-6)
+    assert metadata["branches"]["dinov3"]["branch_weight"].item() == 15.0
+    assert metadata["branches"]["vjepa2"]["branch_weight"].item() == 1.0
+    assert metadata["branches"]["videomae_v2"]["branch_weight"].item() == 2.0
+    assert metadata["aggregation_divisor"].item() == 10.0
+
+    gradients = torch.autograd.grad(
+        loss,
+        tuple(gradient.values()) + tuple(teacher.values()),
+        allow_unused=True,
+    )
+    assert all(value is not None and value.norm() > 0 for value in gradients[:3])
+    assert all(value is None for value in gradients[3:])
+
+
+def test_dual_dinov3_weight_is_applied_after_fd_normalization():
+    objective = MultiRepresentationVideoNystromMMD(
+        {
+            "dinov3": _objective("tiny-dinov3-weighted", (1, 3), 4210),
+            "vjepa2": _objective("tiny-vjepa2-weighted", (1, 2), 4211),
+        },
+        weights={"dinov3": 10.0, "vjepa2": 1.0},
+        normalization_eps=1.0e-4,
+    )
+    generator = torch.Generator().manual_seed(4212)
+    generated = {
+        "dinov3": torch.randn(
+            8, 1, 3, generator=generator, requires_grad=True),
+        "vjepa2": torch.randn(
+            8, 1, 2, generator=generator, requires_grad=True),
+    }
+    teacher = {
+        "dinov3": torch.randn(8, 1, 3, generator=generator),
+        "vjepa2": torch.randn(8, 1, 2, generator=generator),
+    }
+
+    loss, metadata = objective(
+        generated,
+        current_teacher=teacher,
+        stage=False,
+    )
+
+    dino = metadata["branches"]["dinov3"]
+    vjepa = metadata["branches"]["vjepa2"]
+    expected = 10.0 * dino["normalized_mmd"] + vjepa["normalized_mmd"]
+    torch.testing.assert_close(loss.detach(), expected)
+    torch.testing.assert_close(
+        dino["weighted_normalized_mmd"],
+        10.0 * dino["normalized_mmd"],
+    )
+    assert dino["branch_weight"].item() == 10.0
+    assert vjepa["branch_weight"].item() == 1.0
+
+    gradients = torch.autograd.grad(loss, tuple(generated.values()))
+    assert all(torch.isfinite(value).all() and value.norm() > 0
+               for value in gradients)
+
+
+def test_mapping_gradcache_suffix_and_cotangent_compensation():
+    objective = _triple()
+    cache = MultiRepresentationFreshVideoGradCache(
+        objective,
+        target_local_samples=4,
+        gradient_local_samples=2,
+        full_student_partial_gradient=True,
+        require_sample_id_alignment=True,
+    )
+    generator = torch.Generator().manual_seed(4203)
+    batches = []
+    for index in range(4):
+        features = {
+            "dinov3": torch.randn(1, 1, 3, generator=generator),
+            "vjepa2": torch.randn(1, 1, 2, generator=generator),
+            "videomae_v2": torch.randn(1, 2, 3, generator=generator),
+        }
+        batches.append(features)
+        selected = index >= 2
+        teacher = ({
+            "dinov3": torch.randn(1, 1, 3, generator=generator),
+            "vjepa2": torch.randn(1, 1, 2, generator=generator),
+            "videomae_v2": torch.randn(1, 2, 3, generator=generator),
+        } if selected else None)
+        sample_id = (f"sample-{index}",)
+        cache.add(
+            features,
+            teacher,
+            gradient_selected=selected,
+            sample_ids=sample_id,
+            current_teacher_sample_ids=sample_id if selected else None,
+        )
+    result = cache.finalize()
+    assert result.replay_batch_indices == (2, 3)
+    assert len(result.fsdp_cotangents) == 2
+    assert set(result.fsdp_cotangents[0]) == {
+        "dinov3", "vjepa2", "videomae_v2"}
+    assert int(result.metadata["gradcache_kss_denominator"]) == 16
+    assert int(result.metadata["gradcache_gradient_global_samples"]) == 2
+    original_loss = result.loss.clone()
+    compensated = tuple({
+        name: value * 2.0 for name, value in cotangent.items()
+    } for cotangent in result.fsdp_cotangents)
+    for before, after in zip(result.fsdp_cotangents, compensated):
+        for name in before:
+            torch.testing.assert_close(after[name], before[name] * 2.0)
+    torch.testing.assert_close(result.loss, original_loss, rtol=0, atol=0)
+
+
+def test_fd_normalization_preserves_exact_signed_and_zero_neighborhood_behavior():
+    eps = 1.0e-4
+    cases = (
+        (-2.0e-4, 2.0, -1.0e4),
+        (0.0, 0.0, 1.0e4),
+        (1.0e-8, 1.0e-8 / (eps + 1.0e-8), 1.0 / (eps + 1.0e-8)),
+    )
+    for value, expected_forward, expected_gradient in cases:
+        raw = torch.tensor(value, dtype=torch.float64, requires_grad=True)
+        normalized, denominator = fd_normalize_raw_mmd(raw, eps)
+        torch.testing.assert_close(
+            normalized.detach(), torch.tensor(expected_forward, dtype=torch.float64))
+        (gradient,) = torch.autograd.grad(normalized, raw)
+        torch.testing.assert_close(
+            gradient, torch.tensor(expected_gradient, dtype=torch.float64))
+        torch.testing.assert_close(
+            denominator, raw.detach() + eps, rtol=0, atol=0)
+
+    singular = torch.tensor(-eps, dtype=torch.float64, requires_grad=True)
+    normalized, denominator = fd_normalize_raw_mmd(singular, eps)
+    assert denominator.item() == 0.0
+    assert torch.isinf(normalized) and normalized.item() < 0.0

1.3B/code/tests/test_readable_config.py

@@ -0,0 +1,98 @@
+"""Exercise readable config inputs through the actual public entry points."""
+import importlib
+import importlib.util
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+from types import SimpleNamespace
+
+import pytest
+import yaml
+from utils.config import load_config, resolve_config_aliases
+
+ROOT = Path(__file__).resolve().parents[2]
+RECIPE = "three_encoder.yaml" if ROOT.name == "1.3B" else "train.yaml"
+
+
+def test_readable_override_wins_over_legacy_defaults_without_losing_metadata(tmp_path):
+    defaults = tmp_path / "defaults.yaml"
+    defaults.write_text("total_batch_size: 8\nmax_steps: 99\nmetadata: {revision: frozen}\n")
+    path = tmp_path / "run.yaml"
+    path.write_text("distribution_batch_size: 256\ntraining_updates: 1\nreference_minibatch_size: 128\n")
+    cfg = load_config(path, defaults)
+    assert (cfg.total_batch_size, cfg.max_steps, cfg.expected_current_teacher_samples) == (256, 1, 128)
+    assert cfg.metadata.revision == "frozen"
+    # Existing runtime YAML and arbitrary metadata survive the boundary unchanged.
+    legacy = {"total_batch_size": 256, "metadata": {"shape": [1, 2, 3]}}
+    resolved = resolve_config_aliases(legacy)
+    resolved["metadata"]["shape"].append(4)
+    assert legacy["metadata"]["shape"] == [1, 2, 3]
+
+
+@pytest.mark.parametrize("step_flag", ["--max_steps", "--training-updates"])
+def test_training_entry_resolves_recipe_before_constructing_trainer(monkeypatch, tmp_path, step_flag):
+    captured = {}
+    class RecordingTrainer:
+        def __init__(self, config):
+            captured["config"] = config
+        def train(self):
+            captured["trained"] = True
+
+    monkeypatch.setitem(sys.modules, "trainer.distribution_matching", SimpleNamespace(DistributionMatchingTrainer=RecordingTrainer))
+    spec = importlib.util.spec_from_file_location("readable_training_entry", ROOT / "code/train.py")
+    entry = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(entry)
+    monkeypatch.setattr(entry.wandb, "finish", lambda: None)
+    monkeypatch.setattr(entry.dist, "is_initialized", lambda: False)
+    monkeypatch.chdir(ROOT / "code")
+    monkeypatch.setattr(sys, "argv", ["train.py", "--config_path", str(ROOT / "configs" / RECIPE),
+                                     "--logdir", str(tmp_path), "--disable-wandb", step_flag, "1"])
+    entry.main()
+    cfg = captured["config"]
+    assert captured["trained"]
+    assert cfg.trainer == "score_distillation" and cfg.distribution_loss == "wan_kv_mmd"
+    assert cfg.total_batch_size == 256 and cfg.expected_current_teacher_samples == 128
+    assert cfg.max_steps == 1 and cfg.lr == 2.5e-6
+    assert cfg.triple_aggregation_divisor == 3.0
+    # Stopping early must not shorten the scientific learning-rate schedule.
+    if ROOT.name == "1.3B":
+        assert cfg.generator_lr_cosine_end_step == 125
+    else:
+        assert cfg.lr_schedule_total_steps == 200
+
+
+def test_prepare_config_preserves_recipe_and_only_relocates_assets(tmp_path):
+    sandbox = tmp_path / "release"
+    for folder in ("configs", "scripts", "code/utils"):
+        (sandbox / folder).mkdir(parents=True, exist_ok=True)
+    for source in (Path("configs") / RECIPE, Path("scripts/prepare_config.py"), Path("code/utils/config.py")):
+        shutil.copy2(ROOT / source, sandbox / source)
+    assets = tmp_path / "assets with spaces"
+    command = [sys.executable, str(sandbox / "scripts/prepare_config.py"), "--asset-root", str(assets)]
+    if ROOT.name == "14B":
+        command.append("--inference-only")  # No external feature-bank files are needed to relocate paths.
+    subprocess.run(command, check=True, capture_output=True, text=True)
+    original = resolve_config_aliases(yaml.safe_load((ROOT / "configs" / RECIPE).read_text()))
+    prefix = "/path/to/video-generation-assets"
+    expected = {k: str(assets) + v[len(prefix):] if isinstance(v, str) and v.startswith(prefix + "/") else v
+                for k, v in original.items()}
+    actual = yaml.safe_load((sandbox / "configs/local.yaml").read_text())
+    assert actual == expected
+    assert list(actual) == list(expected)
+
+
+def test_legacy_modules_share_implementation_and_method_aliases():
+    for old, new, old_class, new_class in (
+        ("model.wan_kv_mmd", "model.distribution_matching", "WanKVMMD", "DistributionMatchingModel"),
+        ("trainer.distillation", "trainer.distribution_matching", "Trainer", "DistributionMatchingTrainer"),
+        ("losses.multi_representation_video_nystrom_mmd", "losses.multi_encoder_mmd", "MultiRepresentationVideoNystromMMD", "MultiEncoderMMD"),
+        ("trainer.multi_representation_fresh_video_gradcache", "trainer.feature_gradient_replay", "MultiRepresentationFreshVideoGradCache", "FeatureGradientReplay"),
+    ):
+        legacy, current = importlib.import_module(old), importlib.import_module(new)
+        assert legacy is current
+        assert getattr(legacy, old_class) is getattr(current, new_class)
+    from model.distribution_matching import DistributionMatchingModel
+    from trainer.distribution_matching import DistributionMatchingTrainer
+    assert DistributionMatchingModel.fresh_video_gradcache_prepare_replay is DistributionMatchingModel.prepare_latent_gradients
+    assert DistributionMatchingTrainer._fresh_video_generator_update is DistributionMatchingTrainer._run_distribution_matching_step

1.3B/code/tests/test_triple_feature_metadata_contract.py

@@ -0,0 +1,82 @@
+import torch
+
+from model.wan_kv_mmd import WanKVMMD
+from video_fd.encoders.vjepa2_dense_metadata import (
+    canonical_vjepa2_dense_feature_metadata,
+    uniform_dense_token_indices,
+)
+
+
+def test_triple_vjepa_metadata_matches_bank_last_64_contract():
+    model = object.__new__(WanKVMMD)
+    model.feature_type = "triple_dinov3_vjepa2_videomae_v2_dense_latent"
+    model.dense_feature_prefix = "triple"
+    model.dense_encoder_name = "triple"
+    model.dense_log_prefix = "triple_dense"
+    model.vjepa2_reference_model_id = "vjepa2-vitl-fpc64-256"
+    model.vjepa2_model_id = "/unused/local/model"
+    model.vjepa2_dense_layer = 15
+    model.vjepa2_dense_layers = (15,)
+    model.vjepa2_dense_num_frames = 64
+    model.vjepa2_dense_source_num_frames = 81
+    model.vjepa2_dense_frame_sampling = "last_contiguous"
+    model.vjepa2_dense_resize_short_side = 292
+    model.vjepa2_dense_crop_size = 256
+    model.vjepa2_dense_raw_tokens = 2048
+    model.vjepa2_dense_num_tokens = 1024
+    model.vjepa2_dense_feature_dim = 1024
+    indices = uniform_dense_token_indices(2048, 1024)
+    actual = model._vjepa2_feature_metadata(indices)
+    expected = canonical_vjepa2_dense_feature_metadata(
+        model_id="vjepa2-vitl-fpc64-256",
+        layer_index=15,
+        num_frames=64,
+        resize_short_side=292,
+        crop_size=256,
+        raw_tokens=2048,
+        num_tokens=1024,
+        feature_dim=1024,
+        token_indices=indices,
+        frame_sampling="last_contiguous",
+        source_num_frames=81,
+    )
+    assert actual == expected
+    assert actual["selected_frame_start"] == 17
+    assert actual["selected_frame_end_inclusive"] == 80
+
+
+def test_dual_dinov3_vjepa_metadata_uses_explicit_uniform_indices_by_default():
+    model = object.__new__(WanKVMMD)
+    model.feature_type = "dual_dinov3_vjepa2_dense_latent"
+    model.dense_feature_prefix = "dinov3_vjepa2"
+    model.dense_encoder_name = "dinov3_vjepa2"
+    model.dense_log_prefix = "dinov3_vjepa2_dense"
+    model.vjepa2_reference_model_id = "vjepa2-vitl-fpc64-256"
+    model.vjepa2_model_id = "/unused/local/model"
+    model.vjepa2_dense_layer = 15
+    model.vjepa2_dense_layers = (15,)
+    model.vjepa2_dense_num_frames = 64
+    model.vjepa2_dense_source_num_frames = 81
+    model.vjepa2_dense_frame_sampling = "center_contiguous"
+    model.vjepa2_dense_resize_short_side = 292
+    model.vjepa2_dense_crop_size = 256
+    model.vjepa2_dense_raw_tokens = 2048
+    model.vjepa2_dense_num_tokens = 1024
+    model.vjepa2_dense_feature_dim = 1024
+    actual = model._vjepa2_feature_metadata()
+    indices = uniform_dense_token_indices(2048, 1024)
+    expected = canonical_vjepa2_dense_feature_metadata(
+        model_id="vjepa2-vitl-fpc64-256",
+        layer_index=15,
+        num_frames=64,
+        resize_short_side=292,
+        crop_size=256,
+        raw_tokens=2048,
+        num_tokens=1024,
+        feature_dim=1024,
+        token_indices=indices,
+        frame_sampling="last_contiguous",
+        source_num_frames=81,
+    )
+    assert actual == expected
+    assert actual["token_alignment"] == "same_explicit_indices_and_order"

1.3B/code/tests/test_triple_partial_geometry.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+
+
+def test_released_recipe_uses_10_1_1_loss_divided_by_three():
+    import yaml
+    from utils.config import resolve_config_aliases
+    root = Path(__file__).resolve().parents[2]
+    config = resolve_config_aliases(yaml.safe_load((root / "configs/three_encoder.yaml").read_text()))
+    assert config["triple_dinov3_weight"] == 10.0
+    assert config["triple_vjepa2_weight"] == 1.0
+    assert config["triple_videomae_v2_weight"] == 1.0
+    assert config["triple_aggregation_divisor"] == 3.0
+    assert config["partial_mmd_geometry"] == [256, 128, 128]

1.3B/code/tests/test_triple_pixel_cotangent_accumulation.py

@@ -0,0 +1,34 @@
+import pytest
+import torch
+
+from model.wan_kv_mmd import _accumulate_fp32_pixel_cotangent
+
+
+def test_pixel_cotangents_accumulate_immediately_in_fp32():
+    accumulated = None
+    shape = None
+    expected = torch.zeros(2, 3, 4, dtype=torch.float32)
+    norms = []
+    for index, dtype in enumerate(
+        (torch.bfloat16, torch.float16, torch.float32), start=1
+    ):
+        value = torch.full((2, 3, 4), float(index), dtype=dtype)
+        expected.add_(value.float())
+        accumulated, shape, norm = _accumulate_fp32_pixel_cotangent(
+            accumulated, value, shape, branch_name=f"branch{index}")
+        norms.append(norm)
+    assert accumulated.dtype == torch.float32
+    torch.testing.assert_close(accumulated, expected)
+    assert shape == (2, 3, 4)
+    assert all(item.ndim == 0 for item in norms)
+
+
+def test_pixel_cotangent_accumulation_rejects_shape_and_nonfinite():
+    first, shape, _ = _accumulate_fp32_pixel_cotangent(
+        None, torch.ones(1, 2), None, branch_name="first")
+    with pytest.raises(ValueError, match="shapes differ"):
+        _accumulate_fp32_pixel_cotangent(
+            first, torch.ones(1, 3), shape, branch_name="second")
+    with pytest.raises(FloatingPointError, match="non-finite"):
+        _accumulate_fp32_pixel_cotangent(
+            None, torch.tensor([float("nan")]), None, branch_name="bad")

1.3B/code/tests/test_triple_video_feature_bank.py

@@ -0,0 +1,196 @@
+import copy
+import hashlib
+import json
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from utils.triple_video_feature_bank import (
+    BANK_KIND,
+    BANK_SCHEMA_VERSION,
+    TripleVideoFeatureBank,
+    canonical_sha256,
+    file_sha256,
+)
+from trainer.distillation import Trainer
+
+
+VJEPA_SHA = "1" * 64
+VIDEOMAE_REVISION = "9981a9c8f77118c421e5228e1b219468a4b0238d"
+VIDEOMAE_DOWNLOAD_HASH = "2" * 64
+SOURCE_SHA = "3" * 64
+DINOV3_SHA = "4" * 64
+DINOV3_COMMIT = "6876159a11b4df116f30f667f8c9888617df0751"
+
+
+def _write_bank(root: Path):
+    samples = [
+        {"index": index, "sample_id": f"sample-{index}", "prompt": f"p{index}"}
+        for index in range(3)
+    ]
+    branch_metadata = {
+        "dinov3": {
+            "sample_shape": [1024, 1024],
+            "dtype": "float16",
+            "feature_contract": {"name": "dinov3-test"},
+            "model_identity": {
+                "checkpoint_sha256": DINOV3_SHA,
+                "source_commit": DINOV3_COMMIT,
+            },
+        },
+        "vjepa2": {
+            "sample_shape": [1024, 1024],
+            "dtype": "float16",
+            "feature_contract": {"name": "vjepa2-test"},
+            "model_identity": {"sha256": VJEPA_SHA},
+        },
+        "videomae_v2": {
+            "sample_shape": [1280, 1024],
+            "dtype": "float16",
+            "feature_contract": {"name": "videomae-test"},
+            "model_identity": {
+                "revision": VIDEOMAE_REVISION,
+                "download_manifest_content_hash": VIDEOMAE_DOWNLOAD_HASH,
+            },
+        },
+    }
+    shard_dir = root / "shards"
+    shard_dir.mkdir(parents=True)
+    shards = []
+    for shard_index, (start, end) in enumerate(((0, 2), (2, 3))):
+        shard_content_hash = canonical_sha256({
+            "index": shard_index,
+            "sample_start": start,
+            "sample_end": end,
+        })
+        path = shard_dir / f"{shard_index:05d}.safetensors"
+        save_file(
+            {
+                "dinov3": torch.stack([
+                    torch.full((1024, 1024), float(row + 20), dtype=torch.float16)
+                    for row in range(start, end)
+                ]),
+                "vjepa2": torch.stack([
+                    torch.full((1024, 1024), float(row), dtype=torch.float16)
+                    for row in range(start, end)
+                ]),
+                "videomae_v2": torch.stack([
+                    torch.full((1280, 1024), float(row + 10), dtype=torch.float16)
+                    for row in range(start, end)
+                ]),
+            },
+            path,
+            metadata={"content_hash": shard_content_hash},
+        )
+        shards.append({
+            "index": shard_index,
+            "path": str(path.relative_to(root)),
+            "sample_start": start,
+            "sample_end": end,
+            "sample_ids_sha256": canonical_sha256([
+                item["sample_id"] for item in samples[start:end]
+            ]),
+            "sha256": file_sha256(path),
+            "content_hash": shard_content_hash,
+        })
+    payload = {
+        "schema_version": BANK_SCHEMA_VERSION,
+        "kind": BANK_KIND,
+        "sample_count": len(samples),
+        "source_manifest": {"sha256": SOURCE_SHA},
+        "branches": branch_metadata,
+        "samples": samples,
+        "shards": shards,
+    }
+    manifest = {**payload, "content_hash": canonical_sha256(payload)}
+    manifest_path = root / "manifest.json"
+    manifest_path.write_text(
+        json.dumps(manifest, sort_keys=True) + "\n", encoding="utf-8")
+    return manifest_path, manifest
+
+
+def _open(manifest_path: Path, manifest: dict, **overrides):
+    kwargs = {
+        "expected_manifest_sha256": file_sha256(manifest_path),
+        "expected_content_hash": manifest["content_hash"],
+        "expected_source_manifest_sha256": SOURCE_SHA,
+        "expected_sample_count": 3,
+        "expected_videomae_revision": VIDEOMAE_REVISION,
+        "expected_videomae_download_manifest_content_hash": (
+            VIDEOMAE_DOWNLOAD_HASH),
+        "expected_vjepa_model_sha256": VJEPA_SHA,
+        "expected_dinov3_checkpoint_sha256": DINOV3_SHA,
+        "expected_dinov3_source_commit": DINOV3_COMMIT,
+        "expected_shard_size": 2,
+        "validate_all_shards": True,
+    }
+    kwargs.update(overrides)
+    return TripleVideoFeatureBank(manifest_path, **kwargs)
+
+
+def test_random_access_preserves_order_and_duplicate_sample_ids(tmp_path):
+    manifest_path, manifest = _write_bank(tmp_path)
+    bank = _open(manifest_path, manifest)
+    assert [bank[index]["sample_ids"] for index in (2, 0, 2)] == [
+        "sample-2", "sample-0", "sample-2"]
+    loaded = bank.load_rows(["sample-2", "sample-0", "sample-2"])
+    assert tuple(loaded["dinov3"].shape) == (3, 1024, 1024)
+    assert tuple(loaded["vjepa2"].shape) == (3, 1024, 1024)
+    assert tuple(loaded["videomae_v2"].shape) == (3, 1280, 1024)
+    torch.testing.assert_close(
+        loaded["dinov3"][:, 0, 0], torch.tensor([22, 20, 22]).half())
+    torch.testing.assert_close(
+        loaded["vjepa2"][:, 0, 0], torch.tensor([2, 0, 2]).half())
+    torch.testing.assert_close(
+        loaded["videomae_v2"][:, 0, 0],
+        torch.tensor([12, 10, 12]).half(),
+    )
+
+
+
+
+
+
+def test_counter_sampler_is_rank_major_ordered_and_with_replacement():
+    class Dataset:
+        def __len__(self):
+            return 3
+
+        def __getitem__(self, index):
+            return {"sample_ids": f"sample-{index}", "prompts": f"p{index}"}
+
+        @staticmethod
+        def collate_fn(items):
+            return {
+                "sample_ids": [item["sample_ids"] for item in items],
+                "prompts": [item["prompts"] for item in items],
+            }
+
+    trainer = Trainer.__new__(Trainer)
+    trainer.matched_counter_sampling_enabled = True
+    trainer.dataset = Dataset()
+    trainer.config = SimpleNamespace(seed=7, batch_size=1)
+    trainer.step = 11
+    trainer.data_counter = 0
+    trainer.fresh_video_total_samples = 8
+    trainer.fresh_video_local_samples = 2
+    expected = tuple(
+        int.from_bytes(hashlib.sha256(
+            f"sf-matched-counter-v1\0{trainer.config.seed}\0{trainer.step}\0{position}"
+            .encode("ascii")
+        ).digest()[:8], "little") % len(trainer.dataset)
+        for position in range(trainer.fresh_video_total_samples)
+    )
+    assert trainer._matched_global_indices_for_step() == expected
+    # Eight with-replacement draws from a three-row bank must contain repeats.
+    assert len(set(expected)) < len(expected)
+    with mock.patch("trainer.distillation.dist.get_rank", return_value=2):
+        first = trainer._matched_batch_for_microbatch(0)
+        second = trainer._matched_batch_for_microbatch(1)
+    assert first["sample_ids"] == [f"sample-{expected[4]}"]
+    assert second["sample_ids"] == [f"sample-{expected[5]}"]
+    assert trainer.data_counter == 2

1.3B/code/tests/test_video_nystrom_partial_mmd.py

@@ -0,0 +1,372 @@
+import unittest
+
+import torch
+
+from losses.video_nystrom_mmd import (
+    VideoNystromMMD,
+    build_video_nystrom_reference,
+    rbf_frobenius_kernel,
+)
+
+
+class TestVideoNystromPartialMMD(unittest.TestCase):
+    def _pack(self):
+        reference = torch.randn(
+            12,
+            1,
+            2,
+            generator=torch.Generator().manual_seed(9100),
+        )
+        return build_video_nystrom_reference(
+            reference,
+            num_landmarks=5,
+            metadata={
+                "feature": {
+                    "name": "tiny-partial-video",
+                    "matrix_shape": [1, 2],
+                },
+                "reference": {"dataset": "partial-unit-test"},
+            },
+            bandwidth=1.6,
+            kmeans_iterations=4,
+            ridge=1.0e-5,
+            chunk_size=3,
+        )
+
+    def _matched_objective(self, *, chunk_size: int = 37):
+        return VideoNystromMMD(
+            self._pack(),
+            kernel_chunk_size=chunk_size,
+            attraction_weight=2.0,
+            current_teacher_attraction_enabled=True,
+            current_teacher_attraction_weight=1.0,
+        )
+
+    def test_128_by_256_formula_denominators_and_suffix_only_gradient(self):
+        objective = self._matched_objective()
+        generator = torch.Generator().manual_seed(9101)
+        all_students = torch.randn(
+            256, 1, 2, generator=generator, requires_grad=True)
+        gradient_students = all_students[-128:]
+        teacher = torch.randn(
+            128, 1, 2, generator=generator, requires_grad=True)
+
+        loss, metadata = objective(
+            gradient_students=gradient_students,
+            student_reference=all_students,
+            current_teacher=teacher,
+            stage=False,
+        )
+
+        kernel_gs = rbf_frobenius_kernel(
+            gradient_students,
+            all_students.detach(),
+            objective.sigma,
+        )
+        fixed_attraction = (
+            rbf_frobenius_kernel(
+                gradient_students,
+                objective.landmarks,
+                objective.sigma,
+            )
+            @ objective.alpha
+        ).mean()
+        kernel_gt = rbf_frobenius_kernel(
+            gradient_students,
+            teacher.detach(),
+            objective.sigma,
+        )
+        kernel_tt = rbf_frobenius_kernel(
+            teacher.detach(),
+            teacher.detach(),
+            objective.sigma,
+        )
+        expected = (
+            kernel_gs.sum() / float(128 * 256)
+            - 2.0
+            * (
+                (2.0 / 3.0) * fixed_attraction
+                + (1.0 / 3.0) * kernel_gt.sum() / float(128 * 128)
+            )
+            + (2.0 / 3.0) * objective.k_rr
+            + (1.0 / 3.0) * kernel_tt.sum() / float(128 * 128)
+        )
+
+        torch.testing.assert_close(
+            loss, expected, rtol=2.0e-5, atol=2.0e-6)
+        torch.testing.assert_close(
+            metadata["repulsion"],
+            (kernel_gs.sum() / float(128 * 256)).detach(),
+            rtol=2.0e-5,
+            atol=2.0e-6,
+        )
+        self.assertEqual(int(metadata["global_samples"]), 256)
+        self.assertEqual(int(metadata["gradient_student_samples"]), 128)
+        self.assertEqual(int(metadata["student_reference_samples"]), 256)
+        self.assertEqual(int(metadata["current_teacher_samples"]), 128)
+        self.assertEqual(int(metadata["repulsion_denominator"]), 32768)
+        self.assertEqual(
+            int(metadata["student_repulsion_denominator"]), 32768)
+        self.assertEqual(
+            int(metadata["current_teacher_attraction_denominator"]), 16384)
+        self.assertEqual(
+            int(metadata["current_teacher_self_kernel_denominator"]), 16384)
+        self.assertEqual(
+            int(metadata["fixed_teacher_attraction_denominator"]), 128)
+        self.assertEqual(
+            int(metadata["partial_student_reference_enabled"]), 1)
+        self.assertEqual(int(metadata["student_reference_detached"]), 1)
+        self.assertNotIn("full_student_partial_gradient", metadata)
+        self.assertAlmostEqual(
+            float(metadata[
+                "fixed_teacher_attraction_normalized_weight"]),
+            2.0 / 3.0,
+            places=7,
+        )
+        self.assertAlmostEqual(
+            float(metadata[
+                "current_teacher_attraction_normalized_weight"]),
+            1.0 / 3.0,
+            places=7,
+        )
+
+        actual_student_gradient, actual_teacher_gradient = (
+            torch.autograd.grad(
+                loss,
+                (all_students, teacher),
+                allow_unused=True,
+                retain_graph=True,
+            )
+        )
+        expected_student_gradient, = torch.autograd.grad(
+            expected, all_students)
+        torch.testing.assert_close(
+            actual_student_gradient,
+            expected_student_gradient,
+            rtol=3.0e-5,
+            atol=3.0e-6,
+        )
+        self.assertIsNone(actual_teacher_gradient)
+        torch.testing.assert_close(
+            actual_student_gradient[:128],
+            torch.zeros_like(actual_student_gradient[:128]),
+            rtol=0.0,
+            atol=0.0,
+        )
+        self.assertGreater(
+            float(actual_student_gradient[-128:].norm()), 0.0)
+        self.assertTrue(torch.isfinite(actual_student_gradient).all())
+
+    def test_full_student_partial_gradient_formula_and_denominators(self):
+        objective = self._matched_objective()
+        generator = torch.Generator().manual_seed(9105)
+        student_reference = torch.randn(
+            256, 1, 2, generator=generator, requires_grad=True)
+        gradient_student_indices = torch.cat([
+            torch.arange(rank * 64 + 32, rank * 64 + 64)
+            for rank in range(4)
+        ])
+        gradient_students = (
+            student_reference.detach()
+            .index_select(0, gradient_student_indices)
+            .clone()
+            .requires_grad_(True)
+        )
+        teacher = torch.randn(
+            128, 1, 2, generator=generator, requires_grad=True)
+
+        loss, metadata = objective(
+            gradient_students=gradient_students,
+            student_reference=student_reference,
+            current_teacher=teacher,
+            gradient_student_indices=gradient_student_indices,
+            stage=False,
+        )
+
+        full_students = student_reference.detach().index_copy(
+            0, gradient_student_indices, gradient_students)
+        kernel_ss = rbf_frobenius_kernel(
+            full_students, full_students, objective.sigma)
+        fixed_attraction = (
+            rbf_frobenius_kernel(
+                full_students, objective.landmarks, objective.sigma)
+            @ objective.alpha
+        ).mean()
+        kernel_st = rbf_frobenius_kernel(
+            full_students, teacher.detach(), objective.sigma)
+        kernel_tt = rbf_frobenius_kernel(
+            teacher.detach(), teacher.detach(), objective.sigma)
+        expected = (
+            kernel_ss.sum() / float(256 * 256)
+            - 2.0
+            * (
+                (2.0 / 3.0) * fixed_attraction
+                + (1.0 / 3.0) * kernel_st.sum() / float(256 * 128)
+            )
+            + (2.0 / 3.0) * objective.k_rr
+            + (1.0 / 3.0) * kernel_tt.sum() / float(128 * 128)
+        )
+
+        torch.testing.assert_close(
+            loss, expected, rtol=2.0e-5, atol=2.0e-6)
+        self.assertEqual(int(metadata["global_samples"]), 256)
+        self.assertEqual(int(metadata["gradient_student_samples"]), 128)
+        self.assertEqual(int(metadata["student_reference_samples"]), 256)
+        self.assertEqual(int(metadata["current_teacher_samples"]), 128)
+        self.assertEqual(int(metadata["repulsion_denominator"]), 65536)
+        self.assertEqual(
+            int(metadata["student_repulsion_denominator"]), 65536)
+        self.assertEqual(
+            int(metadata["fixed_teacher_attraction_denominator"]), 256)
+        self.assertEqual(
+            int(metadata["current_teacher_attraction_denominator"]),
+            32768,
+        )
+        self.assertEqual(
+            int(metadata["current_teacher_self_kernel_denominator"]),
+            16384,
+        )
+        self.assertEqual(
+            int(metadata["full_student_partial_gradient"]), 1)
+        self.assertEqual(
+            int(metadata["student_repulsion_gradient_sides"]), 2)
+
+        actual_gradient, reference_gradient, teacher_gradient = (
+            torch.autograd.grad(
+                loss,
+                (gradient_students, student_reference, teacher),
+                allow_unused=True,
+                retain_graph=True,
+            )
+        )
+        expected_gradient, = torch.autograd.grad(
+            expected, gradient_students, retain_graph=True)
+        torch.testing.assert_close(
+            actual_gradient,
+            expected_gradient,
+            rtol=3.0e-5,
+            atol=3.0e-6,
+        )
+        self.assertIsNone(reference_gradient)
+        self.assertIsNone(teacher_gradient)
+        self.assertGreater(float(actual_gradient.norm()), 0.0)
+        self.assertTrue(torch.isfinite(actual_gradient).all())
+
+        full_repulsion_gradient, = torch.autograd.grad(
+            kernel_ss.sum() / float(256 * 256),
+            gradient_students,
+            retain_graph=True,
+        )
+        rectangular_repulsion = rbf_frobenius_kernel(
+            gradient_students, student_reference.detach(), objective.sigma
+        ).sum() / float(128 * 256)
+        rectangular_repulsion_gradient, = torch.autograd.grad(
+            rectangular_repulsion, gradient_students)
+        torch.testing.assert_close(
+            full_repulsion_gradient,
+            rectangular_repulsion_gradient,
+            rtol=3.0e-5,
+            atol=3.0e-6,
+        )
+
+    def test_inputs_are_detached_and_converted_to_fp32_for_kernels(self):
+        objective = self._matched_objective(chunk_size=2)
+        generator = torch.Generator().manual_seed(9102)
+        gradient_students = torch.randn(
+            3, 1, 2, generator=generator, requires_grad=True)
+        student_reference = torch.randn(
+            5,
+            1,
+            2,
+            generator=generator,
+            dtype=torch.float64,
+            requires_grad=True,
+        )
+        teacher = torch.randn(
+            3,
+            1,
+            2,
+            generator=generator,
+            dtype=torch.float64,
+            requires_grad=True,
+        )
+
+        loss, _ = objective(
+            gradient_students=gradient_students,
+            student_reference=student_reference,
+            current_teacher=teacher,
+            stage=False,
+        )
+        student_gradient, reference_gradient, teacher_gradient = (
+            torch.autograd.grad(
+                loss,
+                (gradient_students, student_reference, teacher),
+                allow_unused=True,
+            )
+        )
+        self.assertEqual(loss.dtype, torch.float32)
+        self.assertIsNotNone(student_gradient)
+        self.assertIsNone(reference_gradient)
+        self.assertIsNone(teacher_gradient)
+
+    def test_legacy_interface_preserves_full_symmetric_formula_and_schema(self):
+        objective = self._matched_objective(chunk_size=2)
+        generator = torch.Generator().manual_seed(9103)
+        generated = torch.randn(
+            4, 1, 2, generator=generator, requires_grad=True)
+        teacher = torch.randn(4, 1, 2, generator=generator)
+
+        loss, metadata = objective(
+            generated,
+            current_teacher=teacher,
+            stage=False,
+        )
+        kernel_gg = rbf_frobenius_kernel(
+            generated, generated, objective.sigma)
+        fixed_attraction = (
+            rbf_frobenius_kernel(
+                generated, objective.landmarks, objective.sigma)
+            @ objective.alpha
+        ).mean()
+        kernel_gt = rbf_frobenius_kernel(
+            generated, teacher, objective.sigma)
+        kernel_tt = rbf_frobenius_kernel(
+            teacher, teacher, objective.sigma)
+        expected = (
+            kernel_gg.mean()
+            - 2.0
+            * (
+                (2.0 / 3.0) * fixed_attraction
+                + (1.0 / 3.0) * kernel_gt.mean()
+            )
+            + (2.0 / 3.0) * objective.k_rr
+            + (1.0 / 3.0) * kernel_tt.mean()
+        )
+        torch.testing.assert_close(loss, expected)
+        actual_gradient, = torch.autograd.grad(
+            loss, generated, retain_graph=True)
+        expected_gradient, = torch.autograd.grad(expected, generated)
+        torch.testing.assert_close(
+            actual_gradient,
+            expected_gradient,
+            rtol=2.0e-5,
+            atol=2.0e-6,
+        )
+        self.assertEqual(int(metadata["global_samples"]), 4)
+        for partial_key in (
+            "partial_student_reference_enabled",
+            "gradient_student_samples",
+            "student_reference_samples",
+            "student_reference_detached",
+            "student_repulsion_denominator",
+            "fixed_teacher_attraction_denominator",
+            "current_teacher_attraction_denominator",
+            "current_teacher_self_kernel_denominator",
+            "full_student_partial_gradient",
+        ):
+            self.assertNotIn(partial_key, metadata)
+
+
+
+if __name__ == "__main__":
+    unittest.main()

1.3B/code/tools/build_triple_nystrom_references.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""Build three fixed Nyström references sequentially from one triple bank."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import os
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+from losses.video_nystrom_mmd import (
+    build_video_nystrom_reference,
+    compute_reference_pack_hash,
+    load_reference_pack,
+    save_reference_pack,
+)
+from utils.triple_video_feature_bank import (
+    TripleVideoFeatureBank,
+    file_sha256,
+)
+
+
+OUTPUT_NAMES = {
+    "dinov3": "dinov3_vitl16_f20mean5_k1024_m256_matched5k.pt",
+    "vjepa2": "vjepa2_l15_last64_k1024_m1024_matched5k.pt",
+    "videomae_v2": "videomae_v2_l15_fixed5x16_k256_m256_matched5k.pt",
+}
+LANDMARKS = {"dinov3": 256, "vjepa2": 1024, "videomae_v2": 256}
+
+
+def reference_feature_metadata(bank: TripleVideoFeatureBank, branch: str) -> dict:
+    record = bank.manifest["branches"][branch]
+    raw = copy.deepcopy(record["feature_contract"])
+    if branch != "videomae_v2":
+        return raw
+    from video_fd.encoders.videomae_v2_dense_metadata import (
+        canonical_videomae_v2_dense_feature_metadata,
+    )
+    expected = canonical_videomae_v2_dense_feature_metadata(
+        model_id=record["model_identity"]["model_id"],
+        layer_index=int(raw.get("layer_index", 15)),
+        num_tokens=int(raw.get("num_tokens_per_segment", 256)),
+        token_indices=raw.get("relative_token_indices"),
+        vae_id=str(raw.get("vae_id", "Wan2.1-T2V-1.3B/Wan2.1_VAE.pth")),
+        num_segments=int(raw.get("num_segments", 5)),
+        frames_per_segment=int(raw.get("frames_per_segment", 16)),
+        raw_tokens_per_segment=int(raw.get("raw_tokens_per_segment", 1568)),
+        feature_dim=int(raw.get("feature_dim", 1024)),
+        resize_short_side=int(raw.get("resize_short_side", 224)),
+        crop_size=int(raw.get("crop_size", 224)),
+        skip_first_frame=bool(raw.get("skip_first_frame", True)),
+    )
+    ignored = {"model_id", "feature_identity_sha256"}
+    if (
+        {key: value for key, value in raw.items() if key not in ignored}
+        != {key: value for key, value in expected.items() if key not in ignored}
+    ):
+        raise ValueError("VideoMAE bank and runtime protocols differ beyond model ID")
+    return expected
+
+
+def atomic_torch_save(value, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    torch.save(value, temporary)
+    os.replace(temporary, path)
+
+
+def load_branch(bank: TripleVideoFeatureBank, branch: str) -> torch.Tensor:
+    parts = []
+    for shard in bank.shards:
+        with safe_open(
+            shard["resolved_path"], framework="pt", device="cpu"
+        ) as handle:
+            parts.append(handle.get_tensor(branch))
+    value = torch.cat(parts, dim=0)
+    expected = (len(bank), *bank.manifest["branches"][branch]["sample_shape"])
+    if value.dtype != torch.float16 or tuple(value.shape) != expected:
+        raise ValueError(f"bank branch {branch} geometry/dtype mismatch")
+    return value
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--bank-manifest", required=True)
+    parser.add_argument("--bank-manifest-sha256", required=True)
+    parser.add_argument("--bank-content-hash", required=True)
+    parser.add_argument("--source-manifest-sha256", required=True)
+    parser.add_argument("--videomae-revision", default=(
+        "9981a9c8f77118c421e5228e1b219468a4b0238d"))
+    parser.add_argument(
+        "--videomae-download-manifest-content-hash", required=True)
+    parser.add_argument("--vjepa-model-sha256", required=True)
+    parser.add_argument("--dinov3-checkpoint-sha256", required=True)
+    parser.add_argument("--dinov3-source-commit", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--kmeans-iterations", type=int, default=20)
+    parser.add_argument("--chunk-size", type=int, default=32)
+    parser.add_argument("--ridge", type=float, default=1.0e-6)
+    parser.add_argument("--expected-sample-count", type=int, default=5000)
+    parser.add_argument("--expected-shard-size", type=int, default=32)
+    parser.add_argument("--dataset-name", default="fdwan13b_vidprom5k_matched")
+    parser.add_argument("--skip-bank-shard-hashes", action="store_true")
+    args = parser.parse_args()
+    bank = TripleVideoFeatureBank(
+        args.bank_manifest,
+        expected_manifest_sha256=args.bank_manifest_sha256,
+        expected_content_hash=args.bank_content_hash,
+        expected_source_manifest_sha256=args.source_manifest_sha256,
+        expected_sample_count=args.expected_sample_count,
+        expected_videomae_revision=args.videomae_revision,
+        expected_videomae_download_manifest_content_hash=(
+            args.videomae_download_manifest_content_hash),
+        expected_vjepa_model_sha256=args.vjepa_model_sha256,
+        expected_dinov3_checkpoint_sha256=(
+            args.dinov3_checkpoint_sha256),
+        expected_dinov3_source_commit=args.dinov3_source_commit,
+        expected_shard_size=args.expected_shard_size,
+        validate_all_shards=not args.skip_bank_shard_hashes,
+    )
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    build_request = {
+        "chunk_size": int(args.chunk_size),
+        "kmeans_iterations": int(args.kmeans_iterations),
+        "ridge": float(args.ridge),
+        "seed": 0,
+    }
+    summaries = {}
+    for branch in ("dinov3", "vjepa2", "videomae_v2"):
+        output = output_dir / OUTPUT_NAMES[branch]
+        branch_manifest = bank.manifest["branches"][branch]
+        feature_metadata = reference_feature_metadata(bank, branch)
+        migrated_metadata = False
+        if output.is_file():
+            existing = load_reference_pack(output, map_location="cpu")
+            reference_metadata = existing["metadata"].get("reference", {})
+            source = reference_metadata.get("source", {})
+            nystrom = existing["metadata"].get("nystrom", {})
+            if (
+                source.get("triple_teacher_bank_content_hash")
+                != bank.content_hash
+                or int(nystrom.get("num_landmarks", -1)) != LANDMARKS[branch]
+                or int(nystrom.get("kmeans_iterations", -1))
+                != int(args.kmeans_iterations)
+                or float(nystrom.get("ridge", float("nan")))
+                != float(args.ridge)
+                or int(nystrom.get("seed", -1)) != 0
+                or reference_metadata.get("build_request") != build_request
+            ):
+                raise ValueError(
+                    f"existing reference has another identity: {output}")
+            if existing["metadata"].get("feature") != feature_metadata:
+                old_feature = existing["metadata"].get("feature", {})
+                ignored = {"model_id", "feature_identity_sha256"}
+                if (
+                    branch != "videomae_v2"
+                    or {key: value for key, value in old_feature.items()
+                        if key not in ignored}
+                    != {key: value for key, value in feature_metadata.items()
+                        if key not in ignored}
+                ):
+                    raise ValueError(
+                        f"existing {branch} reference feature protocol differs")
+                source_content_hash = existing["content_hash"]
+                migrated = copy.deepcopy(existing)
+                migrated["metadata"]["feature"] = feature_metadata
+                migrated["metadata"]["reference"][
+                    "metadata_only_migration"] = {
+                        "source_reference_content_hash": source_content_hash,
+                        "numerical_tensors_preserved": True,
+                    }
+                migrated["content_hash"] = compute_reference_pack_hash(migrated)
+                save_reference_pack(migrated, output, atomic=True)
+                existing = load_reference_pack(output, map_location="cpu")
+                migrated_metadata = True
+            summaries[branch] = {
+                "path": str(output),
+                "sha256": file_sha256(output),
+                "content_hash": existing["content_hash"],
+                "resumed": True,
+                "metadata_only_migration": migrated_metadata,
+            }
+            continue
+        cpu_reference = load_branch(bank, branch)
+        reference = cpu_reference.to(
+            device=torch.device(args.device), dtype=torch.float32)
+        del cpu_reference
+        metadata = {
+            "feature": feature_metadata,
+            "reference": {
+                "dataset": args.dataset_name,
+                "build_request": build_request,
+                "source": {
+                    "triple_teacher_bank_content_hash": bank.content_hash,
+                    "branch_name": branch,
+                    "triple_teacher_bank_manifest_sha256": (
+                        bank.manifest_sha256),
+                    "source_manifest_sha256": (
+                        args.source_manifest_sha256),
+                    "selected_samples": len(bank),
+                    "sample_order": "frozen_manifest_order",
+                    "feature_bank_fp16_sha256": branch_manifest[
+                        "feature_bank_fp16_sha256"],
+                    "model_identity": branch_manifest["model_identity"],
+                },
+            },
+        }
+        pack = build_video_nystrom_reference(
+            reference,
+            num_landmarks=LANDMARKS[branch],
+            metadata=metadata,
+            kmeans_iterations=args.kmeans_iterations,
+            ridge=args.ridge,
+            chunk_size=args.chunk_size,
+            seed=0,
+        )
+        del reference
+        cpu_pack = {
+            key: (value.detach().cpu() if isinstance(value, torch.Tensor) else value)
+            for key, value in pack.items()
+        }
+        if compute_reference_pack_hash(cpu_pack) != cpu_pack["content_hash"]:
+            raise RuntimeError("reference hash changed during CPU transfer")
+        atomic_torch_save(cpu_pack, output)
+        validated = load_reference_pack(output, map_location="cpu")
+        summaries[branch] = {
+            "path": str(output),
+            "sha256": file_sha256(output),
+            "content_hash": validated["content_hash"],
+            "resumed": False,
+        }
+        del pack, cpu_pack, validated
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    manifest = {
+        "schema_version": 1,
+        "bank_manifest": str(Path(args.bank_manifest).resolve()),
+        "bank_manifest_sha256": args.bank_manifest_sha256,
+        "bank_content_hash": bank.content_hash,
+        "build_request": build_request,
+        "references": summaries,
+    }
+    path = output_dir / "triple_reference_manifest.json"
+    temporary = path.with_name(f".{path.name}.tmp")
+    temporary.write_text(
+        json.dumps(manifest, sort_keys=True, indent=2) + "\n",
+        encoding="utf-8",
+    )
+    os.replace(temporary, path)
+    print(json.dumps(manifest, sort_keys=True), flush=True)
+
+
+if __name__ == "__main__":
+    main()

1.3B/code/tools/download_videomae_v2_large.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Stage an existing local VideoMAE-v2 Large snapshot; no network access."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import tempfile
+import shutil
+
+
+MODEL_ID = "VideoMAEv2-Large"
+REVISION = "9981a9c8f77118c421e5228e1b219468a4b0238d"
+
+
+def file_sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while block := handle.read(8 * 1024 * 1024):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def canonical_sha256(value) -> str:
+    payload = json.dumps(
+        value, sort_keys=True, separators=(",", ":"), allow_nan=False
+    ).encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()
+
+
+def atomic_json(path: Path, value) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, name = tempfile.mkstemp(
+        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+            json.dump(value, handle, sort_keys=True, indent=2)
+            handle.write("\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(name, path)
+    finally:
+        Path(name).unlink(missing_ok=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--output-dir",
+        default=(
+            "assets/models/"
+            "VideoMAEv2-Large"),
+    )
+    parser.add_argument("--source-dir", required=True, type=Path,
+                        help="Existing local snapshot; not a hosted repository ID")
+    args = parser.parse_args()
+    output = Path(args.output_dir).expanduser().resolve()
+    source = args.source_dir.expanduser().resolve()
+    if not source.is_dir():
+        parser.error("--source-dir must be an existing local directory")
+    if source == output or source in output.parents or output in source.parents:
+        parser.error("Source and output directories must be separate and non-nested")
+    if output.exists() and (not output.is_dir() or any(output.iterdir())):
+        parser.error("--output-dir must be absent or empty")
+    if not (source / "config.json").is_file():
+        parser.error("The local snapshot must include config.json")
+    if not any(source.glob("*.safetensors")) and not any(source.glob("*.bin")):
+        parser.error("The local snapshot must include model weights")
+    shutil.copytree(source, output, dirs_exist_ok=True,
+                    ignore=shutil.ignore_patterns(".git", ".cache", "__pycache__", ".DS_Store"))
+    # Existing frozen banks depend on the original manifest digest. Copy that
+    # manifest unchanged when supplied, rather than rewriting its identity.
+    existing_manifest = output / "download_manifest.json"
+    if existing_manifest.is_file():
+        manifest = json.loads(existing_manifest.read_text())
+        payload = dict(manifest)
+        saved = payload.pop("content_hash", None)
+        if manifest.get("revision") != REVISION or canonical_sha256(payload) != saved:
+            raise ValueError("Invalid pinned local snapshot manifest")
+        for item in manifest["files"]:
+            path = (output / item["path"]).resolve()
+            if (output not in path.parents or path.stat().st_size != item["size_bytes"]
+                    or file_sha256(path) != item["sha256"]):
+                raise ValueError(f"Local snapshot file identity mismatch: {path}")
+        print(json.dumps({"output_dir": str(output), "file_count": len(manifest["files"]),
+                          "content_hash": saved, "preserved_manifest": True}, sort_keys=True))
+        return
+    files = []
+    for path in sorted(output.rglob("*")):
+        if (
+            not path.is_file()
+            or path.name == "download_manifest.json"
+            or ".cache" in path.parts
+        ):
+            continue
+        files.append({
+            "path": path.relative_to(output).as_posix(),
+            "size_bytes": path.stat().st_size,
+            "sha256": file_sha256(path),
+        })
+    payload = {
+        "schema_version": 1,
+        "model_id": MODEL_ID,
+        "revision": REVISION,
+        "files": files,
+    }
+    payload["content_hash"] = canonical_sha256(payload)
+    atomic_json(output / "download_manifest.json", payload)
+    print(json.dumps({
+        "output_dir": str(output),
+        "file_count": len(files),
+        "content_hash": payload["content_hash"],
+    }, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

1.3B/code/tools/extract_triple_matched_teacher_bank.py

@@ -0,0 +1,472 @@
+#!/usr/bin/env python3
+"""Decode each frozen 81-frame MP4 once and write a resumable triple bank."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import tempfile
+import time
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from utils.triple_video_feature_bank import (
+    BANK_KIND,
+    BANK_SCHEMA_VERSION,
+    canonical_sha256,
+    file_sha256,
+)
+from utils.matched_video_manifest import (
+    load_matched_video_manifest,
+    sha256_file,
+)
+from video_fd.encoders.dinov3_temporal_pool_adapter import (
+    DINOv3TemporalPoolAdapter,
+)
+from video_fd.encoders.segmented_videomae_v2 import (
+    SegmentedVideoMAEv2Ensemble,
+)
+from video_fd.encoders.videomae_v2_dense_metadata import (
+    canonical_videomae_v2_dense_feature_metadata,
+    uniform_segment_token_indices,
+)
+from video_fd.encoders.vjepa2_adapter import VJEPA2Adapter
+from video_fd.encoders.vjepa2_dense_metadata import (
+    canonical_vjepa2_dense_feature_metadata,
+    uniform_dense_token_indices,
+)
+
+
+VIDEOMAE_REVISION = "9981a9c8f77118c421e5228e1b219468a4b0238d"
+
+
+def atomic_json(path: Path, value) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, name = tempfile.mkstemp(
+        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+            json.dump(value, handle, sort_keys=True, indent=2, ensure_ascii=False)
+            handle.write("\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(name, path)
+    finally:
+        Path(name).unlink(missing_ok=True)
+
+
+def tensor_sha256(tensor: torch.Tensor) -> str:
+    value = tensor.detach().cpu().contiguous()
+    return hashlib.sha256(value.numpy().tobytes()).hexdigest()
+
+
+def decode_81(path: Path, *, size: int, mtime_ns: int) -> torch.Tensor:
+    stat = path.stat()
+    if stat.st_size != size or stat.st_mtime_ns != mtime_ns:
+        raise ValueError(f"frozen video stat changed: {path}")
+    import av
+
+    frames = []
+    with av.open(str(path), mode="r") as container:
+        if len(container.streams.video) != 1:
+            raise ValueError(f"expected one video stream: {path}")
+        stream = container.streams.video[0]
+        for frame in container.decode(stream):
+            array = frame.to_ndarray(format="rgb24")
+            if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3:
+                raise ValueError(f"invalid RGB frame in {path}")
+            frames.append(torch.from_numpy(
+                np.ascontiguousarray(array)).permute(2, 0, 1))
+    if len(frames) != 81:
+        raise ValueError(f"{path.name} has {len(frames)} frames, expected 81")
+    shape = tuple(frames[0].shape)
+    if any(tuple(frame.shape) != shape for frame in frames):
+        raise ValueError(f"resolution changes within {path}")
+    return torch.stack(frames).float().div_(255.0).unsqueeze(0)
+
+
+def load_download_manifest(model_dir: Path) -> dict:
+    path = model_dir / "download_manifest.json"
+    with path.open("r", encoding="utf-8") as handle:
+        manifest = json.load(handle)
+    payload = dict(manifest)
+    saved = payload.pop("content_hash")
+    if (
+        manifest.get("revision") != VIDEOMAE_REVISION
+        or canonical_sha256(payload) != saved
+    ):
+        raise ValueError("invalid pinned VideoMAE download manifest")
+    for item in manifest["files"]:
+        path = (model_dir / item["path"]).resolve()
+        if (
+            model_dir not in path.parents
+            or path.stat().st_size != int(item["size_bytes"])
+            or file_sha256(path) != item["sha256"]
+        ):
+            raise ValueError(f"VideoMAE file identity mismatch: {path}")
+    return manifest
+
+
+def build_encoders(args, device):
+    dino = DINOv3TemporalPoolAdapter({
+        "repo_path": str(args.dinov3_repo),
+        "weight_path": str(args.dinov3_weights),
+        "source_commit": str(args.dinov3_source_commit),
+        "checkpoint_sha256": str(args.dinov3_checkpoint_sha256),
+        "source_num_frames": 81,
+        "frame_indices": (
+            0, 4, 8, 13, 17, 21, 25, 29, 34, 38,
+            42, 46, 51, 55, 59, 63, 67, 72, 76, 80),
+        "input_size": 256,
+        "patch_size": 16,
+        "spatial_axis_indices": tuple(range(16)),
+        "forward_chunk_frames": 4,
+        "geometry": "temporal_mean5_full256",
+        "temporal_groups": 4,
+        "temporal_pool_frames": 5,
+    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
+    vj = VJEPA2Adapter({
+        "name": "vjepa2_dense",
+        "model_id": str(args.vjepa_model),
+        "input_frames": 64,
+        "dense_num_frames": 64,
+        "input_size": 256,
+        "resize_short_side": 292,
+        "crop_mode": "center",
+        "dense_only": True,
+        "dense_layer_index": 15,
+    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
+    vm = SegmentedVideoMAEv2Ensemble({
+        "name": "videomae_v2_five_segment_dense",
+        "model_id": str(args.videomae_model),
+        "revision": VIDEOMAE_REVISION,
+        "local_files_only": True,
+        "encoder_count": 5,
+        "skip_leading_frames": 1,
+        "segment_num_frames": 16,
+        "source_num_frames": 81,
+        "input_frames": 16,
+        "dense_num_frames": 16,
+        "input_size": 224,
+        "resize_short_side": 224,
+        "crop_mode": "center",
+        "dense_only": True,
+        "dense_layer_index": 15,
+    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
+    if (dino.dense_feature_dim, vj.dense_feature_dim, vm.dense_feature_dim) != (1024, 1024, 1024):
+        raise RuntimeError("encoder feature dimensions differ from contract")
+    return dino, vj, vm
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", required=True)
+    parser.add_argument("--manifest-sha256", required=True)
+    parser.add_argument("--source-csv-sha256", required=True)
+    parser.add_argument("--video-root", required=True)
+    parser.add_argument("--dinov3-repo", required=True)
+    parser.add_argument("--dinov3-weights", required=True)
+    parser.add_argument("--dinov3-source-commit", required=True)
+    parser.add_argument("--dinov3-checkpoint-sha256", required=True)
+    parser.add_argument("--vjepa-model", required=True)
+    parser.add_argument("--vjepa-reference-model-id", required=True)
+    parser.add_argument("--vjepa-model-sha256", required=True)
+    parser.add_argument("--videomae-model", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--shard-rows", type=int, default=32)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--expected-count", type=int, default=5000)
+    args = parser.parse_args()
+    if args.shard_rows != 32:
+        raise ValueError("triple teacher bank requires shard-rows=32")
+    manifest_path = Path(args.manifest).expanduser().resolve()
+    if sha256_file(manifest_path) != args.manifest_sha256:
+        raise ValueError("source manifest SHA-256 mismatch")
+    entries, source_metadata = load_matched_video_manifest(
+        manifest_path,
+        expected_source_csv_sha256=args.source_csv_sha256,
+        expected_count=args.expected_count,
+    )
+    video_root = Path(args.video_root).expanduser().resolve()
+    output = Path(args.output_dir).expanduser().resolve()
+    shard_dir = output / "shards"
+    shard_dir.mkdir(parents=True, exist_ok=True)
+    progress_path = output / "progress.json"
+    final_path = output / "manifest.json"
+    vm_download = load_download_manifest(
+        Path(args.videomae_model).expanduser().resolve())
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    distributed = world_size > 1
+    if distributed:
+        local_rank = int(os.environ["LOCAL_RANK"])
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(backend="nccl")
+        rank = int(dist.get_rank())
+        world_size = int(dist.get_world_size())
+        device = torch.device("cuda", local_rank)
+    else:
+        rank = 0
+        device = torch.device(args.device)
+    dino, vj, vm = build_encoders(args, device)
+    vj_indices = uniform_dense_token_indices(
+        vj.dense_raw_tokens, 1024, device=device)
+    vm_indices = uniform_segment_token_indices(
+        vm.dense_raw_tokens, 256, device=device)
+    vj_contract = canonical_vjepa2_dense_feature_metadata(
+        model_id=args.vjepa_reference_model_id,
+        layer_index=15,
+        num_frames=64,
+        resize_short_side=292,
+        crop_size=256,
+        raw_tokens=vj.dense_raw_tokens,
+        num_tokens=1024,
+        feature_dim=1024,
+        token_indices=vj_indices.cpu(),
+        frame_sampling="last_contiguous",
+        source_num_frames=81,
+    )
+    vm_contract = canonical_videomae_v2_dense_feature_metadata(
+        model_id="VideoMAEv2-Large",
+        layer_index=15,
+        num_tokens=256,
+        token_indices=vm_indices.cpu(),
+    )
+    fixed = {
+        "schema_version": BANK_SCHEMA_VERSION,
+        "kind": BANK_KIND,
+        "sample_count": len(entries),
+        "source_manifest": {
+            "path": str(manifest_path),
+            "sha256": args.manifest_sha256,
+            "source_csv_sha256": args.source_csv_sha256,
+            "sample_identity_sha256": source_metadata.get(
+                "sample_identity_sha256"),
+        },
+        "branches": {
+            "dinov3": {
+                "sample_shape": [1024, 1024],
+                "dtype": "float16",
+                "feature_contract": dino.feature_metadata(),
+                "model_identity": {
+                    "model_id": "dinov3_vitl16",
+                    "repo_path": str(Path(args.dinov3_repo).resolve()),
+                    "weight_path": str(Path(args.dinov3_weights).resolve()),
+                    "source_commit": args.dinov3_source_commit,
+                    "checkpoint_sha256": args.dinov3_checkpoint_sha256,
+                },
+            },
+            "vjepa2": {
+                "sample_shape": [1024, 1024],
+                "dtype": "float16",
+                "feature_contract": vj_contract,
+                "model_identity": {
+                    "model_id": args.vjepa_reference_model_id,
+                    "load_path": str(Path(args.vjepa_model).resolve()),
+                    "sha256": args.vjepa_model_sha256,
+                },
+            },
+            "videomae_v2": {
+                "sample_shape": [1280, 1024],
+                "dtype": "float16",
+                "feature_contract": vm_contract,
+                "model_identity": {
+                    "model_id": "VideoMAEv2-Large",
+                    "revision": VIDEOMAE_REVISION,
+                    "download_manifest_content_hash": vm_download[
+                        "content_hash"],
+                },
+            },
+        },
+        "samples": [{
+            "index": entry.manifest_index,
+            "sample_id": entry.sample_id,
+            "prompt": entry.caption,
+        } for entry in entries],
+    }
+    progress_identity = canonical_sha256(fixed)
+    receipt_dir = output / "receipts"
+    receipt_dir.mkdir(parents=True, exist_ok=True)
+    shard_count = (len(entries) + args.shard_rows - 1) // args.shard_rows
+    started = time.time()
+    for shard_index in range(rank, shard_count, world_size):
+        start = shard_index * args.shard_rows
+        end = min(start + args.shard_rows, len(entries))
+        relative = f"shards/shard-{shard_index:05d}.safetensors"
+        path = output / relative
+        receipt_path = receipt_dir / f"shard-{shard_index:05d}.json"
+        if receipt_path.is_file():
+            receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
+            if (
+                receipt.get("progress_identity") != progress_identity
+                or int(receipt.get("index", -1)) != shard_index
+                or int(receipt.get("sample_start", -1)) != start
+                or int(receipt.get("sample_end", -1)) != end
+                or not path.is_file()
+                or file_sha256(path) != receipt.get("sha256")
+            ):
+                raise ValueError(
+                    f"resumable triple shard identity mismatch: {path}")
+            print(json.dumps({
+                "rank": rank,
+                "resumed_shard": shard_index,
+                "completed": end,
+            }), flush=True)
+            continue
+
+        branch_rows = {
+            "dinov3": [],
+            "vjepa2": [],
+            "videomae_v2": [],
+        }
+        for entry in entries[start:end]:
+            video = decode_81(
+                video_root / entry.video_name,
+                size=entry.video_size_bytes,
+                mtime_ns=entry.video_mtime_ns,
+            ).to(device=device, dtype=torch.float32)
+            temporal = torch.arange(
+                17, 81, device=device, dtype=torch.long).unsqueeze(0)
+            with torch.inference_mode():
+                dino_feature = dino.forward_dense_no_grad(
+                    video, gradient_checkpointing=False)
+                vj_feature = vj.forward_dense(
+                    video,
+                    layer_index=15,
+                    token_indices=vj_indices,
+                    gradient_checkpointing=False,
+                    temporal_window_indices=temporal,
+                )
+                vm_feature = vm.forward_dense(
+                    video,
+                    layer_index=15,
+                    token_indices=vm_indices,
+                    gradient_checkpointing=False,
+                )
+            observed = {
+                "dinov3": tuple(dino_feature.shape),
+                "vjepa2": tuple(vj_feature.shape),
+                "videomae_v2": tuple(vm_feature.shape),
+            }
+            expected = {
+                "dinov3": (1, 1024, 1024),
+                "vjepa2": (1, 1024, 1024),
+                "videomae_v2": (1, 1280, 1024),
+            }
+            if observed != expected:
+                raise RuntimeError(
+                    f"triple bank row shape mismatch: {observed}")
+            branch_rows["dinov3"].append(
+                dino_feature.squeeze(0).half().cpu().contiguous())
+            branch_rows["vjepa2"].append(
+                vj_feature.squeeze(0).half().cpu().contiguous())
+            branch_rows["videomae_v2"].append(
+                vm_feature.squeeze(0).half().cpu().contiguous())
+            del video, dino_feature, vj_feature, vm_feature
+
+        tensors = {
+            name: torch.stack(rows, dim=0)
+            for name, rows in branch_rows.items()
+        }
+        ids = [entry.sample_id for entry in entries[start:end]]
+        shard_payload = {
+            "index": shard_index,
+            "sample_start": start,
+            "sample_end": end,
+            "sample_ids_sha256": canonical_sha256(ids),
+            "tensor_sha256": {
+                name: tensor_sha256(value)
+                for name, value in tensors.items()
+            },
+            "branches": fixed["branches"],
+            "source_manifest_sha256": args.manifest_sha256,
+            "progress_identity": progress_identity,
+        }
+        shard_content_hash = canonical_sha256(shard_payload)
+        temporary = path.with_name(f".{path.name}.rank{rank}.tmp")
+        save_file(tensors, str(temporary), metadata={
+            "content_hash": shard_content_hash,
+            "sample_start": str(start),
+            "sample_end": str(end),
+            "source_manifest_sha256": args.manifest_sha256,
+            "progress_identity": progress_identity,
+        })
+        os.replace(temporary, path)
+        shard = {
+            "index": shard_index,
+            "path": relative,
+            "sample_start": start,
+            "sample_end": end,
+            "sample_ids_sha256": canonical_sha256(ids),
+            "sha256": file_sha256(path),
+            "content_hash": shard_content_hash,
+        }
+        atomic_json(receipt_path, {
+            **shard,
+            "progress_identity": progress_identity,
+            "rank": rank,
+        })
+        print(json.dumps({
+            "rank": rank,
+            "completed": end,
+            "total": len(entries),
+            "elapsed_seconds": time.time() - started,
+            "shard": relative,
+        }), flush=True)
+
+    if distributed:
+        dist.barrier()
+    if rank == 0:
+        shards = []
+        for shard_index in range(shard_count):
+            receipt_path = receipt_dir / f"shard-{shard_index:05d}.json"
+            if not receipt_path.is_file():
+                raise FileNotFoundError(
+                    f"missing completed shard receipt: {receipt_path}")
+            receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
+            if receipt.pop("progress_identity", None) != progress_identity:
+                raise ValueError("triple shard progress identity mismatch")
+            receipt.pop("rank", None)
+            path = output / receipt["path"]
+            if file_sha256(path) != receipt["sha256"]:
+                raise ValueError(f"triple shard SHA-256 mismatch: {path}")
+            shards.append(receipt)
+
+        branch_hashers = {
+            name: hashlib.sha256() for name in fixed["branches"]
+        }
+        for shard in shards:
+            with safe_open(
+                str(output / shard["path"]),
+                framework="pt",
+                device="cpu",
+            ) as handle:
+                for name in branch_hashers:
+                    value = handle.get_tensor(name).contiguous()
+                    branch_hashers[name].update(value.numpy().tobytes())
+        for name, hasher in branch_hashers.items():
+            fixed["branches"][name]["feature_bank_fp16_sha256"] = (
+                hasher.hexdigest())
+        final_payload = {**fixed, "shards": shards}
+        final_payload["content_hash"] = canonical_sha256(final_payload)
+        atomic_json(final_path, final_payload)
+        print(json.dumps({
+            "manifest": str(final_path),
+            "manifest_sha256": file_sha256(final_path),
+            "content_hash": final_payload["content_hash"],
+            "samples": len(entries),
+            "world_size": world_size,
+        }, sort_keys=True), flush=True)
+    if distributed:
+        dist.barrier()
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

1.3B/code/tools/precompute_dinov3_dense_nystrom_reference.py

@@ -0,0 +1,1160 @@
+#!/usr/bin/env python3
+"""Build the fixed DINOv3 ViT-L/16 whole-video Nyström reference.
+
+The teacher and objective are unchanged from the paired V-JEPA2 experiment:
+the final latent from every row of the 1,537-video CausVid ODE LMDB is decoded
+with the frozen Wan VAE, then represented by a frozen DINOv3 image encoder.
+The DINOv3 feature map is fixed to 16 uniformly distributed source frames and
+an 8x8 subset of the 16x16 patch grid per frame.  Frame-major concatenation
+therefore produces one aligned ``[16 * 64, 1024]`` matrix per video.
+
+Every torchrun rank atomically writes one indexed FP16 feature shard.  Rank 0
+validates and merges all shards in LMDB-index order, constructs the existing
+video-matrix Nyström objective (m=256, seed=0), atomically publishes the pack,
+and reloads the *final path* for validation.  Temporary shards are removed only
+after that final validation succeeds; failed runs retain them for diagnosis.
+"""
+
+from __future__ import annotations
+
+import argparse
+from datetime import timedelta
+import gc
+import hashlib
+import json
+import math
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple
+
+import torch
+import torch.distributed as dist
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+
+TOOL_FORMAT_VERSION = 1
+MIB = 1024 ** 2
+EXPECTED_TEACHER_COUNT = 1537
+EXPECTED_SOURCE_COMMIT = "6876159a11b4df116f30f667f8c9888617df0751"
+SOURCE_NUM_FRAMES = 81
+FRAME_INDICES = (0, 5, 11, 16, 21, 27, 32, 37, 43, 48, 53, 59, 64, 69, 75, 80)
+INPUT_SIZE = 256
+PATCH_SIZE = 16
+PATCH_GRID_SIZE = 16
+SPATIAL_AXIS_INDICES = (0, 2, 4, 6, 9, 11, 13, 15)
+SPATIAL_TOKEN_INDICES = tuple(
+    row * PATCH_GRID_SIZE + column
+    for row in SPATIAL_AXIS_INDICES
+    for column in SPATIAL_AXIS_INDICES
+)
+PER_FRAME_TOKENS = 64
+TOTAL_TOKENS = len(FRAME_INDICES) * PER_FRAME_TOKENS
+FEATURE_DIM = 1024
+FORWARD_CHUNK_FRAMES = 4
+NUM_LANDMARKS = 256
+NYSTROM_SEED = 0
+MEDIAN_MAX_SAMPLES = EXPECTED_TEACHER_COUNT
+KMEANS_ITERATIONS = 20
+NYSTROM_RIDGE = 1.0e-6
+KERNEL_CHUNK_SIZE = 32
+WAN_VAE_MODEL_NAME = "Wan2.1-T2V-1.3B"
+WAN_VAE_MEAN = (
+    -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+    0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921,
+)
+WAN_VAE_STD = (
+    2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+    3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160,
+)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Build the fixed full-1537 DINOv3 ViT-L/16 video-matrix "
+            "Nyström teacher reference."
+        )
+    )
+    parser.add_argument("--teacher-lmdb", type=Path, required=True)
+    parser.add_argument("--repo-path", type=Path, required=True)
+    parser.add_argument("--weight-path", type=Path, required=True)
+    parser.add_argument(
+        "--source-commit", type=str, default=EXPECTED_SOURCE_COMMIT)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument(
+        "--expected-teacher-count", type=int, default=EXPECTED_TEACHER_COUNT)
+    parser.add_argument(
+        "--forward-chunk-frames", type=int, default=FORWARD_CHUNK_FRAMES)
+    parser.add_argument("--num-landmarks", type=int, default=NUM_LANDMARKS)
+    parser.add_argument("--seed", type=int, default=NYSTROM_SEED)
+    parser.add_argument(
+        "--median-max-samples", type=int, default=MEDIAN_MAX_SAMPLES)
+    parser.add_argument("--bandwidth", type=float, default=None)
+    parser.add_argument(
+        "--kmeans-iterations", type=int, default=KMEANS_ITERATIONS)
+    parser.add_argument("--ridge", type=float, default=NYSTROM_RIDGE)
+    parser.add_argument("--chunk-size", type=int, default=KERNEL_CHUNK_SIZE)
+    parser.add_argument("--tmp-dir", type=Path, default=None)
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--log-every", type=int, default=25)
+    parser.add_argument("--distributed-timeout-minutes", type=int, default=360)
+    parser.add_argument("--dry-run", action="store_true")
+    return parser
+
+
+def _validate_args(args: argparse.Namespace) -> None:
+    positive = {
+        "batch_size": args.batch_size,
+        "expected_teacher_count": args.expected_teacher_count,
+        "forward_chunk_frames": args.forward_chunk_frames,
+        "num_landmarks": args.num_landmarks,
+        "median_max_samples": args.median_max_samples,
+        "kmeans_iterations": args.kmeans_iterations,
+        "chunk_size": args.chunk_size,
+        "distributed_timeout_minutes": args.distributed_timeout_minutes,
+    }
+    invalid = {name: value for name, value in positive.items() if int(value) <= 0}
+    if invalid:
+        raise ValueError(f"arguments must be positive: {invalid}")
+    if int(args.expected_teacher_count) != EXPECTED_TEACHER_COUNT:
+        raise ValueError(
+            "the canonical reference requires exactly "
+            f"{EXPECTED_TEACHER_COUNT} teachers")
+    if int(args.forward_chunk_frames) != FORWARD_CHUNK_FRAMES:
+        raise ValueError(
+            f"--forward-chunk-frames must be {FORWARD_CHUNK_FRAMES}")
+    if int(args.num_landmarks) != NUM_LANDMARKS:
+        raise ValueError(f"--num-landmarks must be {NUM_LANDMARKS}")
+    if int(args.seed) != NYSTROM_SEED:
+        raise ValueError(f"--seed must be {NYSTROM_SEED}")
+    if int(args.median_max_samples) != MEDIAN_MAX_SAMPLES:
+        raise ValueError(
+            f"--median-max-samples must be {MEDIAN_MAX_SAMPLES}")
+    if args.bandwidth is not None:
+        raise ValueError(
+            "the canonical reference requires median-estimated bandwidth")
+    if int(args.kmeans_iterations) != KMEANS_ITERATIONS:
+        raise ValueError(
+            f"--kmeans-iterations must be {KMEANS_ITERATIONS}")
+    if not math.isclose(
+        float(args.ridge), NYSTROM_RIDGE, rel_tol=0.0, abs_tol=0.0
+    ):
+        raise ValueError(f"--ridge must be {NYSTROM_RIDGE}")
+    if int(args.chunk_size) != KERNEL_CHUNK_SIZE:
+        raise ValueError(f"--chunk-size must be {KERNEL_CHUNK_SIZE}")
+    if str(args.source_commit) != EXPECTED_SOURCE_COMMIT:
+        raise ValueError(
+            "--source-commit must pin the approved DINOv3 source commit "
+            f"{EXPECTED_SOURCE_COMMIT}")
+    if float(args.ridge) <= 0.0 or not math.isfinite(float(args.ridge)):
+        raise ValueError("--ridge must be finite and positive")
+    if args.bandwidth is not None and (
+        float(args.bandwidth) <= 0.0
+        or not math.isfinite(float(args.bandwidth))
+    ):
+        raise ValueError("--bandwidth must be finite and positive")
+    if int(args.log_every) < 0:
+        raise ValueError("--log-every must be non-negative")
+
+
+def _git_head(repo_path: Path) -> str:
+    try:
+        result = subprocess.run(
+            ["git", "-C", str(repo_path), "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (OSError, subprocess.CalledProcessError) as exc:
+        raise RuntimeError(
+            f"failed to inspect DINOv3 source checkout {repo_path}") from exc
+    return result.stdout.strip()
+
+
+def _git_source_status(repo_path: Path) -> str:
+    try:
+        result = subprocess.run(
+            [
+                "git",
+                "-C",
+                str(repo_path),
+                "status",
+                "--porcelain",
+                "--untracked-files=all",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (OSError, subprocess.CalledProcessError) as exc:
+        raise RuntimeError(
+            f"failed to inspect DINOv3 source worktree {repo_path}") from exc
+    return result.stdout
+
+
+def _inspect_assets(args: argparse.Namespace) -> Dict[str, Any]:
+    repo_path = args.repo_path.expanduser().resolve()
+    weight_path = args.weight_path.expanduser().resolve()
+    if not repo_path.is_dir():
+        raise FileNotFoundError(f"DINOv3 source directory not found: {repo_path}")
+    if not weight_path.is_file():
+        raise FileNotFoundError(f"DINOv3 weight file not found: {weight_path}")
+    source_commit = _git_head(repo_path)
+    if source_commit != str(args.source_commit):
+        raise ValueError(
+            "DINOv3 source checkout commit mismatch: "
+            f"{source_commit} != {args.source_commit}")
+    source_status = _git_source_status(repo_path)
+    if source_status.strip():
+        raise ValueError(
+            "DINOv3 source checkout has worktree changes; refusing "
+            "to build a reference from unpinned source")
+    stat = weight_path.stat()
+    return {
+        "repo_path": str(repo_path),
+        "weight_path": str(weight_path),
+        "source_commit": source_commit,
+        "weight_size_bytes": int(stat.st_size),
+        "weight_mtime_ns": int(stat.st_mtime_ns),
+    }
+
+
+def _inspect_lmdb(path: Path) -> Dict[str, Any]:
+    try:
+        import lmdb
+    except ImportError as exc:
+        raise ImportError(
+            "python-lmdb is required; use the same Python environment as training"
+        ) from exc
+    from utils.lmdb import get_array_shape_from_lmdb
+
+    resolved = path.expanduser().resolve()
+    if not resolved.is_dir():
+        raise FileNotFoundError(f"teacher LMDB directory not found: {resolved}")
+    env = lmdb.open(
+        str(resolved), readonly=True, lock=False, readahead=False, meminit=False)
+    try:
+        latents_shape = tuple(
+            int(value) for value in get_array_shape_from_lmdb(env, "latents"))
+        prompts_shape = tuple(
+            int(value) for value in get_array_shape_from_lmdb(env, "prompts"))
+    finally:
+        env.close()
+    if len(latents_shape) != 6:
+        raise ValueError(
+            "teacher LMDB latents must be [N,trajectory,T,C,H,W], got "
+            f"{latents_shape}")
+    if prompts_shape != (latents_shape[0],):
+        raise ValueError(
+            f"teacher prompts/latents count mismatch: {prompts_shape} vs "
+            f"{latents_shape}")
+    data_path = resolved / "data.mdb"
+    stat = data_path.stat()
+    return {
+        "path": str(resolved),
+        "latents_shape": latents_shape,
+        "prompts_shape": prompts_shape,
+        "data_mdb_size_bytes": int(stat.st_size),
+        "data_mdb_mtime_ns": int(stat.st_mtime_ns),
+    }
+
+
+def _rank_indices(count: int, rank: int, world_size: int) -> range:
+    return range(int(rank), int(count), int(world_size))
+
+
+def _required_temp_bytes(count: int) -> int:
+    feature_bytes = int(count) * TOTAL_TOKENS * FEATURE_DIM * 2
+    return int(math.ceil(feature_bytes * 1.15)) + 64 * MIB
+
+
+def _required_output_bytes() -> int:
+    landmark_bytes = NUM_LANDMARKS * TOTAL_TOKENS * FEATURE_DIM * 4
+    return int(math.ceil(landmark_bytes * 1.25)) + 32 * MIB
+
+
+def _choose_temp_parent(
+    requested: Optional[Path], required_bytes: int
+) -> Tuple[Path, int]:
+    candidates: Sequence[Path] = (
+        (requested.expanduser(),)
+        if requested is not None
+        else (Path("/dev/shm"), Path(tempfile.gettempdir()))
+    )
+    errors = []
+    for candidate in candidates:
+        try:
+            candidate = candidate.resolve()
+            if not candidate.is_dir():
+                errors.append(f"{candidate}: not a directory")
+                continue
+            if not os.access(candidate, os.W_OK | os.X_OK):
+                errors.append(f"{candidate}: not writable")
+                continue
+            free = int(shutil.disk_usage(candidate).free)
+            if free < required_bytes:
+                errors.append(
+                    f"{candidate}: {free / MIB:.1f} MiB free, "
+                    f"{required_bytes / MIB:.1f} MiB required")
+                continue
+            return candidate, free
+        except OSError as exc:
+            errors.append(f"{candidate}: {exc}")
+    raise RuntimeError(
+        "no temporary shard directory has sufficient capacity: "
+        + "; ".join(errors))
+
+
+def _check_output(
+    output: Path, *, required_bytes: int, overwrite: bool
+) -> Dict[str, Any]:
+    output = output.expanduser().resolve()
+    if output.exists() and not overwrite:
+        raise FileExistsError(
+            f"output already exists: {output}; pass --overwrite to replace it")
+    existing = output.parent
+    while not existing.exists() and existing != existing.parent:
+        existing = existing.parent
+    if not existing.is_dir():
+        raise FileNotFoundError(
+            f"no existing output parent for capacity check: {output.parent}")
+    free = int(shutil.disk_usage(existing).free)
+    if free < required_bytes:
+        raise RuntimeError(
+            f"output filesystem has only {free / MIB:.1f} MiB free; "
+            f"{required_bytes / MIB:.1f} MiB required")
+    return {
+        "path": str(output),
+        "capacity_parent": str(existing),
+        "free_bytes": free,
+        "required_bytes": required_bytes,
+    }
+
+
+def _plan(args: argparse.Namespace, world_size: int) -> Dict[str, Any]:
+    lmdb = _inspect_lmdb(args.teacher_lmdb)
+    count = int(lmdb["latents_shape"][0])
+    if count != EXPECTED_TEACHER_COUNT:
+        raise ValueError(
+            "canonical DINOv3 reference requires the complete 1,537-row LMDB; "
+            f"found {count} rows")
+    if int(args.num_landmarks) > count:
+        raise ValueError("number of landmarks exceeds teacher population")
+    assets = _inspect_assets(args)
+    temp_required = _required_temp_bytes(count)
+    temp_parent, temp_free = _choose_temp_parent(args.tmp_dir, temp_required)
+    output = _check_output(
+        args.output,
+        required_bytes=_required_output_bytes(),
+        overwrite=bool(args.overwrite),
+    )
+    return {
+        "tool_format_version": TOOL_FORMAT_VERSION,
+        "teacher_lmdb": lmdb,
+        "assets": assets,
+        "selected_samples": count,
+        "world_size": int(world_size),
+        "samples_per_rank": [
+            len(_rank_indices(count, rank, world_size))
+            for rank in range(world_size)
+        ],
+        "feature_geometry": {
+            "source_num_frames": SOURCE_NUM_FRAMES,
+            "frame_indices": list(FRAME_INDICES),
+            "input_size": INPUT_SIZE,
+            "patch_size": PATCH_SIZE,
+            "patch_grid_size": PATCH_GRID_SIZE,
+            "spatial_axis_indices": list(SPATIAL_AXIS_INDICES),
+            "spatial_token_indices": list(SPATIAL_TOKEN_INDICES),
+            "raw_tokens_per_frame": PATCH_GRID_SIZE ** 2,
+            "per_frame_tokens": PER_FRAME_TOKENS,
+            "num_tokens": TOTAL_TOKENS,
+            "feature_dim": FEATURE_DIM,
+            "sample_shape": [TOTAL_TOKENS, FEATURE_DIM],
+            "concat_order": "frame_major_then_spatial_row_major",
+        },
+        "num_landmarks": NUM_LANDMARKS,
+        "temporary": {
+            "parent": str(temp_parent),
+            "required_bytes": temp_required,
+            "free_bytes": temp_free,
+        },
+        "output": output,
+    }
+
+
+def _setup_distributed(timeout_minutes: int) -> Tuple[torch.device, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "reference extraction requires CUDA; use --dry-run on CPU")
+    torch.cuda.set_device(local_rank)
+    if world_size > 1:
+        dist.init_process_group(
+            backend="nccl",
+            timeout=timedelta(minutes=int(timeout_minutes)),
+        )
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    return torch.device("cuda", local_rank), rank, world_size
+
+
+def _barrier() -> None:
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+
+def _broadcast_rank0_object(value: Any, rank: int) -> Any:
+    if not (dist.is_available() and dist.is_initialized()):
+        return value
+    payload = [value if rank == 0 else None]
+    dist.broadcast_object_list(payload, src=0)
+    return payload[0]
+
+
+def _atomic_torch_save(value: Any, path: Path) -> None:
+    temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}")
+    try:
+        torch.save(value, temporary)
+        os.replace(temporary, path)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def _canonical_sha256(value: Dict[str, Any]) -> str:
+    payload = json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+        allow_nan=False,
+    ).encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()
+
+
+def _file_identity(path: Path) -> Dict[str, Any]:
+    resolved = path.expanduser().resolve(strict=True)
+    if not resolved.is_file():
+        raise ValueError(f"identity target is not a regular file: {resolved}")
+    before = resolved.stat()
+    digest = hashlib.sha256()
+    with resolved.open("rb") as handle:
+        while True:
+            block = handle.read(8 << 20)
+            if not block:
+                break
+            digest.update(block)
+    after = resolved.stat()
+    if (
+        int(before.st_size) != int(after.st_size)
+        or int(before.st_mtime_ns) != int(after.st_mtime_ns)
+    ):
+        raise RuntimeError(f"identity target changed while hashing: {resolved}")
+    return {
+        "path": str(resolved),
+        "size_bytes": int(after.st_size),
+        "mtime_ns": int(after.st_mtime_ns),
+        "sha256": digest.hexdigest(),
+    }
+
+
+def _vae_identity() -> Dict[str, Any]:
+    """Resolve the exact Wan decoder used before DINOv3 extraction."""
+    from utils.paths import model_path
+
+    weights = model_path(WAN_VAE_MODEL_NAME) / "Wan2.1_VAE.pth"
+    identity = {
+        "wrapper": "utils.wan_wrapper.WanVAEWrapper",
+        "wrapper_source": _file_identity(REPO_ROOT / "utils" / "wan_wrapper.py"),
+        "model_name": WAN_VAE_MODEL_NAME,
+        "weights": _file_identity(weights),
+        "latent_channels": 16,
+        "latent_mean": list(WAN_VAE_MEAN),
+        "latent_std": list(WAN_VAE_STD),
+        "mean_std_contract": "WanVAEWrapper.inline_constants.v1",
+        "compute_dtype": "bfloat16",
+        "decode_use_cache": False,
+        "decode_output_range": "[-1,1]",
+    }
+    identity["identity_sha256"] = _canonical_sha256(identity)
+    return identity
+
+
+def _validate_reference_current_assets(
+    pack: Mapping[str, Any],
+    *,
+    repo_path: Path,
+    weight_path: Path,
+    teacher_lmdb: Path,
+) -> Dict[str, Any]:
+    """Require a validated reference to match every current production asset.
+
+    Callers must first run ``load_reference_pack`` (which verifies schema and
+    content hash). This second gate prevents a valid but stale pack from being
+    accepted after the DINO checkpoint, Wan VAE, or teacher LMDB changes.
+    """
+    if not isinstance(pack, Mapping):
+        raise TypeError("reference pack must be a mapping")
+    metadata = pack.get("metadata")
+    if not isinstance(metadata, Mapping):
+        raise ValueError("reference pack metadata is missing")
+    feature = metadata.get("feature")
+    reference = metadata.get("reference")
+    if not isinstance(feature, Mapping) or not isinstance(reference, Mapping):
+        raise ValueError("reference feature/source metadata is missing")
+    source = reference.get("source")
+    saved_vae = reference.get("vae")
+    if not isinstance(source, Mapping) or not isinstance(saved_vae, Mapping):
+        raise ValueError("reference source or Wan VAE identity is missing")
+
+    from video_fd.encoders.dinov3_dense_metadata import file_sha256
+
+    resolved_repo = repo_path.expanduser().resolve(strict=True)
+    resolved_weight = weight_path.expanduser().resolve(strict=True)
+    current_weight_sha256 = file_sha256(resolved_weight)
+    expected_feature_identity = {
+        "repo_path": str(resolved_repo),
+        "weight_path": str(resolved_weight),
+        "weight_sha256": current_weight_sha256,
+    }
+    feature_mismatches = {
+        key: {"saved": feature.get(key), "current": value}
+        for key, value in expected_feature_identity.items()
+        if feature.get(key) != value
+    }
+    if feature_mismatches:
+        raise ValueError(
+            "reference DINOv3 asset identity does not match current assets: "
+            f"{feature_mismatches}")
+
+    current_vae = _vae_identity()
+    if dict(saved_vae) != current_vae:
+        raise ValueError(
+            "reference Wan VAE identity does not match current assets")
+
+    current_lmdb = _inspect_lmdb(teacher_lmdb)
+    expected_source_identity = {
+        "lmdb_path": current_lmdb["path"],
+        "latents_shape": list(current_lmdb["latents_shape"]),
+        "prompts_shape": list(current_lmdb["prompts_shape"]),
+        "selected_samples": int(current_lmdb["latents_shape"][0]),
+        "data_mdb_size_bytes": int(current_lmdb["data_mdb_size_bytes"]),
+        "data_mdb_mtime_ns": int(current_lmdb["data_mdb_mtime_ns"]),
+    }
+    source_mismatches = {
+        key: {"saved": source.get(key), "current": value}
+        for key, value in expected_source_identity.items()
+        if source.get(key) != value
+    }
+    if source_mismatches:
+        raise ValueError(
+            "reference teacher LMDB identity does not match current assets: "
+            f"{source_mismatches}")
+    return {
+        "dinov3_weight_sha256": current_weight_sha256,
+        "vae_identity_sha256": current_vae["identity_sha256"],
+        "teacher_lmdb": current_lmdb,
+    }
+
+
+def _prompt_entry_digest(index: int, prompt: str) -> bytes:
+    digest = hashlib.sha256()
+    digest.update(int(index).to_bytes(8, byteorder="little", signed=False))
+    digest.update(b"\0")
+    digest.update(str(prompt).encode("utf-8"))
+    return digest.digest()
+
+
+def _load_dense_encoder(args: argparse.Namespace, device: torch.device):
+    """The sole runtime boundary with the separately-owned DINOv3 adapter."""
+    from video_fd.encoders.dinov3_adapter import DINOv3DenseAdapter
+
+    encoder = DINOv3DenseAdapter({
+        "name": "dinov3_dense",
+        "repo_path": str(args.repo_path.expanduser().resolve()),
+        "weight_path": str(args.weight_path.expanduser().resolve()),
+        "source_commit": str(args.source_commit),
+        "source_num_frames": SOURCE_NUM_FRAMES,
+        "frame_indices": list(FRAME_INDICES),
+        "input_size": INPUT_SIZE,
+        "patch_size": PATCH_SIZE,
+        "patch_grid_size": PATCH_GRID_SIZE,
+        "spatial_axis_indices": list(SPATIAL_AXIS_INDICES),
+        "forward_chunk_frames": int(args.forward_chunk_frames),
+        "compute_dtype": "float32",
+    })
+    return encoder.to(device=device, dtype=torch.float32).eval().requires_grad_(False)
+
+
+@torch.no_grad()
+def _extract_rank_shard(
+    args: argparse.Namespace,
+    plan: Dict[str, Any],
+    *,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+    shard_path: Path,
+    vae_identity: Dict[str, Any],
+) -> None:
+    from utils.dataset import ODERegressionLMDBDataset
+    from utils.wan_wrapper import WanVAEWrapper
+
+    selected_count = int(plan["selected_samples"])
+    assigned = list(_rank_indices(selected_count, rank, world_size))
+    dataset = ODERegressionLMDBDataset(
+        str(args.teacher_lmdb.expanduser().resolve()), max_pair=selected_count)
+    if len(dataset) != selected_count:
+        raise RuntimeError(
+            f"LMDB length changed after planning: {len(dataset)} != "
+            f"{selected_count}")
+    vae = (
+        WanVAEWrapper()
+        .to(device=device, dtype=torch.bfloat16)
+        .eval()
+        .requires_grad_(False)
+    )
+    if not torch.equal(vae.mean, torch.tensor(WAN_VAE_MEAN)):
+        raise RuntimeError("WanVAEWrapper latent mean changed from the reference contract")
+    if not torch.equal(vae.std, torch.tensor(WAN_VAE_STD)):
+        raise RuntimeError("WanVAEWrapper latent std changed from the reference contract")
+    encoder = _load_dense_encoder(args, device)
+    feature_metadata = encoder.feature_metadata()
+    if not isinstance(feature_metadata, dict):
+        raise TypeError("DINOv3 adapter feature_metadata() must return a dict")
+    spatial_indices = torch.tensor(
+        SPATIAL_TOKEN_INDICES, device=device, dtype=torch.long)
+
+    feature_chunks = []
+    index_chunks = []
+    prompt_digests = []
+    expected_tail = tuple(plan["teacher_lmdb"]["latents_shape"][1:])
+    started = time.monotonic()
+    for offset in range(0, len(assigned), int(args.batch_size)):
+        batch_indices = assigned[offset:offset + int(args.batch_size)]
+        items = [dataset[index] for index in batch_indices]
+        trajectories = [item["ode_latent"] for item in items]
+        for index, trajectory in zip(batch_indices, trajectories):
+            if tuple(trajectory.shape) != expected_tail:
+                raise ValueError(
+                    f"LMDB row {index} shape {tuple(trajectory.shape)} != "
+                    f"{expected_tail}")
+            if not bool(torch.isfinite(trajectory).all().item()):
+                raise FloatingPointError(
+                    f"LMDB row {index} contains non-finite latents")
+        latent = torch.stack(
+            [trajectory[-1] for trajectory in trajectories], dim=0
+        ).to(device=device, dtype=torch.bfloat16)
+        pixels = vae.decode_to_pixel(latent.detach())
+        video01 = ((pixels + 1.0) * 0.5).clamp(0.0, 1.0)
+        if (
+            video01.ndim != 5
+            or int(video01.shape[1]) != SOURCE_NUM_FRAMES
+            or int(video01.shape[2]) != 3
+        ):
+            raise ValueError(
+                "Wan VAE must decode [B,81,3,H,W], got "
+                f"{tuple(video01.shape)}")
+        features = encoder.forward_dense(
+            video01,
+            token_indices=spatial_indices,
+            gradient_checkpointing=False,
+        ).float()
+        expected_shape = (len(batch_indices), TOTAL_TOKENS, FEATURE_DIM)
+        if tuple(features.shape) != expected_shape:
+            raise RuntimeError(
+                f"DINOv3 features {tuple(features.shape)} != {expected_shape}")
+        if not bool(torch.isfinite(features).all().item()):
+            raise FloatingPointError(
+                f"non-finite DINOv3 features for indices {batch_indices}")
+        features_fp16 = features.to(dtype=torch.float16, device="cpu")
+        if not bool(torch.isfinite(features_fp16).all().item()):
+            raise FloatingPointError(
+                f"FP16 conversion overflowed for indices {batch_indices}")
+        feature_chunks.append(features_fp16)
+        index_chunks.append(torch.tensor(batch_indices, dtype=torch.long))
+        prompt_digests.extend(
+            _prompt_entry_digest(index, item["prompts"])
+            for index, item in zip(batch_indices, items)
+        )
+        del latent, pixels, video01, features, features_fp16, items, trajectories
+        completed = min(offset + len(batch_indices), len(assigned))
+        if int(args.log_every) > 0 and (
+            completed == len(assigned) or completed % int(args.log_every) == 0
+        ):
+            print(
+                f"[extract rank {rank}] {completed}/{len(assigned)} "
+                f"({time.monotonic() - started:.1f}s)",
+                flush=True,
+            )
+
+    if feature_chunks:
+        shard_features = torch.cat(feature_chunks, dim=0).contiguous()
+        shard_indices = torch.cat(index_chunks, dim=0).contiguous()
+        digest_tensor = torch.tensor(
+            [list(value) for value in prompt_digests], dtype=torch.uint8)
+    else:
+        shard_features = torch.empty(
+            0, TOTAL_TOKENS, FEATURE_DIM, dtype=torch.float16)
+        shard_indices = torch.empty(0, dtype=torch.long)
+        digest_tensor = torch.empty(0, 32, dtype=torch.uint8)
+    payload = {
+        "shard_format_version": TOOL_FORMAT_VERSION,
+        "rank": int(rank),
+        "world_size": int(world_size),
+        "indices": shard_indices,
+        "features": shard_features,
+        "prompt_entry_sha256": digest_tensor,
+        "sample_shape": [TOTAL_TOKENS, FEATURE_DIM],
+        "feature_dtype": "float16",
+        "feature_metadata": feature_metadata,
+        "vae_identity": vae_identity,
+    }
+    _atomic_torch_save(payload, shard_path)
+    print(
+        f"[extract rank {rank}] wrote {shard_path} "
+        f"{tuple(shard_features.shape)}",
+        flush=True,
+    )
+    del encoder, vae, dataset, payload, feature_chunks, shard_features
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def _load_and_merge_shards(
+    shard_dir: Path,
+    *,
+    world_size: int,
+    selected_count: int,
+    sample_shape: Sequence[int],
+) -> Tuple[torch.Tensor, str, str, Dict[str, Any], Dict[str, Any]]:
+    feature_parts = []
+    index_parts = []
+    digest_parts = []
+    feature_metadata = None
+    vae_identity = None
+    expected_tail = tuple(int(value) for value in sample_shape)
+    for rank in range(int(world_size)):
+        path = shard_dir / f"rank_{rank:05d}.pt"
+        if not path.is_file():
+            raise FileNotFoundError(f"missing rank shard: {path}")
+        payload = torch.load(path, map_location="cpu", weights_only=False)
+        if int(payload.get("shard_format_version", -1)) != TOOL_FORMAT_VERSION:
+            raise ValueError(f"incompatible shard format in {path}")
+        if int(payload.get("rank", -1)) != rank:
+            raise ValueError(f"rank identity mismatch in {path}")
+        if int(payload.get("world_size", -1)) != int(world_size):
+            raise ValueError(f"world-size mismatch in {path}")
+        features = payload.get("features")
+        indices = payload.get("indices")
+        digests = payload.get("prompt_entry_sha256")
+        shard_feature_metadata = payload.get("feature_metadata")
+        shard_vae_identity = payload.get("vae_identity")
+        if (
+            not isinstance(features, torch.Tensor)
+            or features.dtype != torch.float16
+            or features.ndim != 3
+            or tuple(features.shape[1:]) != expected_tail
+            or not bool(torch.isfinite(features).all().item())
+        ):
+            raise ValueError(f"invalid FP16 feature shard {path}")
+        if (
+            not isinstance(indices, torch.Tensor)
+            or indices.dtype != torch.long
+            or tuple(indices.shape) != (features.shape[0],)
+        ):
+            raise ValueError(f"invalid indices in {path}")
+        if (
+            not isinstance(digests, torch.Tensor)
+            or digests.dtype != torch.uint8
+            or tuple(digests.shape) != (features.shape[0], 32)
+        ):
+            raise ValueError(f"invalid prompt hashes in {path}")
+        if not isinstance(shard_feature_metadata, dict):
+            raise ValueError(f"missing DINOv3 feature metadata in {path}")
+        if feature_metadata is None:
+            feature_metadata = shard_feature_metadata
+        elif shard_feature_metadata != feature_metadata:
+            raise ValueError(
+                f"DINOv3 feature metadata differs across rank shards: {path}")
+        if not isinstance(shard_vae_identity, dict):
+            raise ValueError(f"missing Wan VAE identity in {path}")
+        if vae_identity is None:
+            vae_identity = shard_vae_identity
+        elif shard_vae_identity != vae_identity:
+            raise ValueError(
+                f"Wan VAE identity differs across rank shards: {path}")
+        feature_parts.append(features)
+        index_parts.append(indices)
+        digest_parts.append(digests)
+    features = torch.cat(feature_parts, dim=0)
+    indices = torch.cat(index_parts, dim=0)
+    digests = torch.cat(digest_parts, dim=0)
+    if int(features.shape[0]) != int(selected_count):
+        raise ValueError(
+            f"rank shards contain {features.shape[0]} samples, expected "
+            f"{selected_count}")
+    order = torch.argsort(indices, stable=True)
+    indices = indices.index_select(0, order)
+    if not torch.equal(indices, torch.arange(selected_count, dtype=torch.long)):
+        raise ValueError(
+            "rank shards do not cover every selected LMDB index exactly once")
+    features = features.index_select(0, order).contiguous()
+    digests = digests.index_select(0, order).contiguous()
+    prompt_hash = hashlib.sha256(digests.numpy().tobytes()).hexdigest()
+    feature_hash = hashlib.sha256(
+        features.view(torch.uint8).numpy().tobytes()).hexdigest()
+    if feature_metadata is None:
+        raise ValueError("rank shards contained no DINOv3 feature metadata")
+    if vae_identity is None:
+        raise ValueError("rank shards contained no Wan VAE identity")
+    return features, prompt_hash, feature_hash, feature_metadata, vae_identity
+
+
+def _validate_feature_metadata(
+    args: argparse.Namespace,
+    feature_metadata: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Rebuild the canonical schema and reject even consistent bad shards."""
+    from video_fd.encoders.dinov3_dense_metadata import (
+        canonical_dinov3_dense_feature_metadata,
+        file_sha256,
+    )
+
+    repo_path = args.repo_path.expanduser().resolve()
+    if _git_head(repo_path) != str(args.source_commit):
+        raise RuntimeError("DINOv3 source commit changed during extraction")
+    if _git_source_status(repo_path).strip():
+        raise RuntimeError("DINOv3 source worktree changed during extraction")
+    current_weight_sha256 = file_sha256(args.weight_path)
+    expected = canonical_dinov3_dense_feature_metadata(
+        repo_path=str(repo_path),
+        weight_path=str(args.weight_path.expanduser().resolve()),
+        source_commit=str(args.source_commit),
+        weight_sha256=current_weight_sha256,
+        source_num_frames=SOURCE_NUM_FRAMES,
+        frame_indices=list(FRAME_INDICES),
+        input_size=INPUT_SIZE,
+        patch_size=PATCH_SIZE,
+        patch_grid_size=PATCH_GRID_SIZE,
+        spatial_axis_indices=list(SPATIAL_AXIS_INDICES),
+        spatial_token_indices=list(SPATIAL_TOKEN_INDICES),
+        per_frame_tokens=PER_FRAME_TOKENS,
+        num_tokens=TOTAL_TOKENS,
+        forward_chunk_frames=FORWARD_CHUNK_FRAMES,
+        feature_dim=FEATURE_DIM,
+    )
+    if feature_metadata != expected:
+        raise ValueError(
+            "rank-shard DINOv3 metadata does not equal the canonical schema")
+    return expected
+
+
+def _reference_metadata(
+    args: argparse.Namespace,
+    plan: Dict[str, Any],
+    *,
+    feature_metadata: Dict[str, Any],
+    vae_identity: Dict[str, Any],
+    prompt_index_hash: str,
+    feature_bank_hash: str,
+) -> Dict[str, Any]:
+    feature_metadata = _validate_feature_metadata(args, feature_metadata)
+    if vae_identity != _vae_identity():
+        raise RuntimeError(
+            "Wan VAE asset or wrapper changed during teacher extraction")
+    data_mdb = Path(plan["teacher_lmdb"]["path"]) / "data.mdb"
+    data_stat = data_mdb.stat()
+    if (
+        int(data_stat.st_size)
+        != int(plan["teacher_lmdb"]["data_mdb_size_bytes"])
+        or int(data_stat.st_mtime_ns)
+        != int(plan["teacher_lmdb"]["data_mdb_mtime_ns"])
+    ):
+        raise RuntimeError("teacher LMDB changed during DINOv3 extraction")
+    source = {
+        "kind": "causvid_ode_teacher_lmdb_final_latent",
+        "lmdb_path": plan["teacher_lmdb"]["path"],
+        "latents_shape": list(plan["teacher_lmdb"]["latents_shape"]),
+        "prompts_shape": list(plan["teacher_lmdb"]["prompts_shape"]),
+        "selected_samples": EXPECTED_TEACHER_COUNT,
+        "selected_index_start": 0,
+        "selected_index_stop_exclusive": EXPECTED_TEACHER_COUNT,
+        "trajectory_index": -1,
+        "lmdb_storage_dtype": "float16",
+        "data_mdb_size_bytes": int(
+            plan["teacher_lmdb"]["data_mdb_size_bytes"]),
+        "data_mdb_mtime_ns": int(
+            plan["teacher_lmdb"]["data_mdb_mtime_ns"]),
+        "prompt_index_sha256": prompt_index_hash,
+        "feature_bank_fp16_sha256": feature_bank_hash,
+    }
+    source["identity_sha256"] = _canonical_sha256(source)
+    return {
+        "feature": feature_metadata,
+        "reference": {
+            "source": source,
+            "vae": vae_identity,
+            "extraction": {
+                "tool": "tools/precompute_dinov3_dense_nystrom_reference.py",
+                "tool_format_version": TOOL_FORMAT_VERSION,
+                "world_size": int(plan["world_size"]),
+                "vae_compute_dtype": "bfloat16",
+                "vae_decode_use_cache": False,
+                "decoded_pixel_transform": "((pixels+1)*0.5).clamp(0,1)",
+                "dinov3_compute_dtype": "float32",
+                "forward_chunk_frames": FORWARD_CHUNK_FRAMES,
+                "rank_shard_dtype": "float16",
+            },
+            "build_request": {
+                "num_landmarks": NUM_LANDMARKS,
+                "seed": NYSTROM_SEED,
+                "bandwidth": (
+                    None if args.bandwidth is None else float(args.bandwidth)),
+                "median_max_samples": int(args.median_max_samples),
+                "kmeans_iterations": int(args.kmeans_iterations),
+                "ridge": float(args.ridge),
+                "chunk_size": int(args.chunk_size),
+            },
+        },
+    }
+
+
+def _build_and_save_reference(
+    args: argparse.Namespace,
+    plan: Dict[str, Any],
+    shard_dir: Path,
+    *,
+    device: torch.device,
+) -> Dict[str, Any]:
+    from losses.video_nystrom_mmd import (
+        build_video_nystrom_reference,
+        load_reference_pack,
+        save_reference_pack,
+        validate_reference_pack,
+    )
+
+    features, prompt_hash, feature_hash, feature_metadata, vae_identity = (
+        _load_and_merge_shards(
+            shard_dir,
+            world_size=int(plan["world_size"]),
+            selected_count=EXPECTED_TEACHER_COUNT,
+            sample_shape=(TOTAL_TOKENS, FEATURE_DIM),
+        )
+    )
+    metadata = _reference_metadata(
+        args,
+        plan,
+        feature_metadata=feature_metadata,
+        vae_identity=vae_identity,
+        prompt_index_hash=prompt_hash,
+        feature_bank_hash=feature_hash,
+    )
+    print(
+        f"[merge] building m={NUM_LANDMARKS} reference from "
+        f"{tuple(features.shape)} FP16 matrices on {device}",
+        flush=True,
+    )
+    reference = features.to(device=device, non_blocking=False)
+    pack = build_video_nystrom_reference(
+        reference,
+        num_landmarks=NUM_LANDMARKS,
+        metadata=metadata,
+        bandwidth=args.bandwidth,
+        median_max_samples=int(args.median_max_samples),
+        kmeans_iterations=int(args.kmeans_iterations),
+        ridge=float(args.ridge),
+        chunk_size=int(args.chunk_size),
+        seed=NYSTROM_SEED,
+    )
+    validate_reference_pack(pack)
+    del reference
+    torch.cuda.empty_cache()
+
+    output = args.output.expanduser().resolve()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_name(f".{output.name}.tmp-{os.getpid()}")
+    try:
+        save_reference_pack(pack, temporary)
+        staged = load_reference_pack(temporary, map_location="cpu")
+        validate_reference_pack(staged)
+        if output.exists() and not bool(args.overwrite):
+            raise FileExistsError(
+                f"output appeared during construction: {output}")
+        os.replace(temporary, output)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+    # This final-path reload is the gate for temporary-shard deletion.
+    published = load_reference_pack(output, map_location="cpu")
+    validate_reference_pack(published)
+    _validate_published_pack_identity(pack, published)
+    return {
+        "output": str(output),
+        "samples": int(features.shape[0]),
+        "sample_shape": list(features.shape[1:]),
+        "num_landmarks": int(pack["landmarks"].shape[0]),
+        "sigma": float(torch.as_tensor(pack["sigma"]).item()),
+        "content_hash": str(pack["content_hash"]),
+        "prompt_index_sha256": prompt_hash,
+        "feature_bank_fp16_sha256": feature_hash,
+        "final_path_validated": True,
+    }
+
+
+def _validate_published_pack_identity(
+    pack: Mapping[str, Any],
+    published: Mapping[str, Any],
+) -> None:
+    """Require the final-path reload to preserve the built pack identity."""
+
+    if published.get("content_hash") != pack.get("content_hash"):
+        raise ValueError("published reference content hash changed after rename")
+    # The builder canonicalizes user metadata and adds its reserved ``nystrom``
+    # block. Compare against that exact pack payload, not the pre-build input.
+    if published.get("metadata") != pack.get("metadata"):
+        raise ValueError("published reference metadata changed after rename")
+
+
+def _remove_validated_shards(
+    shard_dir: Path,
+    build_summary: Dict[str, Any],
+) -> None:
+    """Delete shards only with the explicit final-path validation receipt."""
+    if build_summary.get("final_path_validated") is not True:
+        raise RuntimeError(
+            "refusing to delete shards before final reference validation")
+    output = Path(str(build_summary.get("output", ""))).expanduser().resolve()
+    if not output.is_file():
+        raise FileNotFoundError(
+            f"refusing to delete shards; published reference is missing: {output}")
+    shutil.rmtree(shard_dir)
+
+
+def _run(args: argparse.Namespace) -> None:
+    _validate_args(args)
+    env_world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    env_rank = int(os.environ.get("RANK", "0"))
+    plan = _plan(args, env_world_size)
+    if bool(args.dry_run):
+        if env_rank == 0:
+            print(json.dumps(plan, indent=2, sort_keys=True), flush=True)
+        return
+
+    device, rank, world_size = _setup_distributed(
+        int(args.distributed_timeout_minutes))
+    if world_size != int(plan["world_size"]):
+        raise RuntimeError(
+            f"world size changed after planning: {world_size} != "
+            f"{plan['world_size']}")
+    shard_dir = None
+    success = False
+    validated_summary = None
+    try:
+        vae_result = None
+        if rank == 0:
+            try:
+                vae_result = {"identity": _vae_identity(), "error": None}
+            except Exception as exc:
+                vae_result = {
+                    "identity": None,
+                    "error": f"{type(exc).__name__}: {exc}",
+                }
+        vae_result = _broadcast_rank0_object(vae_result, rank)
+        if vae_result["error"] is not None:
+            raise RuntimeError(
+                "rank 0 could not resolve the Wan VAE identity: "
+                f"{vae_result['error']}")
+        vae_identity = vae_result["identity"]
+
+        temp_result = None
+        if rank == 0:
+            try:
+                created = Path(tempfile.mkdtemp(
+                    prefix="dinov3_dense_nystrom_",
+                    dir=Path(plan["temporary"]["parent"]),
+                ))
+                temp_result = {"path": str(created), "error": None}
+            except Exception as exc:
+                temp_result = {
+                    "path": None,
+                    "error": f"{type(exc).__name__}: {exc}",
+                }
+        temp_result = _broadcast_rank0_object(temp_result, rank)
+        if temp_result["error"] is not None:
+            raise RuntimeError(
+                "rank 0 could not create shard directory: "
+                f"{temp_result['error']}")
+        shard_dir = Path(temp_result["path"])
+
+        _extract_rank_shard(
+            args,
+            plan,
+            device=device,
+            rank=rank,
+            world_size=world_size,
+            shard_path=shard_dir / f"rank_{rank:05d}.pt",
+            vae_identity=vae_identity,
+        )
+        _barrier()
+        summary = None
+        build_error = None
+        if rank == 0:
+            try:
+                summary = _build_and_save_reference(
+                    args, plan, shard_dir, device=device)
+            except Exception as exc:
+                import traceback
+                traceback.print_exc()
+                build_error = f"{type(exc).__name__}: {exc}"
+        outcome = _broadcast_rank0_object(
+            {"summary": summary, "error": build_error}, rank)
+        if outcome["error"] is not None:
+            raise RuntimeError(
+                f"rank 0 reference construction failed: {outcome['error']}")
+        _barrier()
+        validated_summary = outcome["summary"]
+        success = True
+        if rank == 0:
+            print(json.dumps(
+                {"status": "ok", **outcome["summary"]},
+                indent=2,
+                sort_keys=True,
+            ), flush=True)
+    finally:
+        if rank == 0 and shard_dir is not None:
+            if success:
+                _remove_validated_shards(shard_dir, validated_summary)
+                print(
+                    f"[cleanup] removed validated temporary shards {shard_dir}",
+                    flush=True,
+                )
+            else:
+                print(
+                    f"[cleanup] preserving failed-run shards at {shard_dir}",
+                    flush=True,
+                )
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    _run(build_parser().parse_args(argv))
+
+
+if __name__ == "__main__":
+    main()

1.3B/code/train.py

@@ -0,0 +1,90 @@
+import argparse
+import os
+import torch.distributed as dist
+from omegaconf import OmegaConf
+from utils.config import load_config
+from utils.optional_wandb import wandb
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True)
+    parser.add_argument("--no_save", action="store_true")
+    parser.add_argument("--no_visualize", action="store_true")
+    parser.add_argument(
+        "--max_steps", "--training-updates",
+        type=int,
+        default=None,
+        help="Optional runtime override for training_updates (useful for smoke tests)",
+    )
+    parser.add_argument(
+        "--fresh-video-local-samples", "--distribution-samples-per-rank",
+        type=int,
+        default=None,
+        help=(
+            "Optional smoke-test override for the number of fresh videos "
+            "collected per rank and optimizer update."
+        ),
+    )
+    parser.add_argument(
+        "--total-batch-size", "--distribution-batch-size",
+        type=int,
+        default=None,
+        help="Optional matching logical global-batch override.",
+    )
+    parser.add_argument(
+        "--eval-interval",
+        type=int,
+        default=None,
+        help="Optional runtime eval interval override; 0 disables eval.",
+    )
+    parser.add_argument("--logdir", type=str, default="", help="Path to the directory to save logs")
+    parser.add_argument("--wandb-save-dir", type=str, default="", help="Path to the directory to save wandb logs")
+    parser.add_argument("--disable-wandb", action="store_true")
+
+    args = parser.parse_args()
+
+    config = load_config(args.config_path, "configs/default_config.yaml")
+    config.no_save = args.no_save
+    config.no_visualize = args.no_visualize
+    if args.max_steps is not None:
+        if args.max_steps <= 0:
+            parser.error("--max_steps must be positive")
+        config.max_steps = args.max_steps
+    if args.fresh_video_local_samples is not None:
+        if args.fresh_video_local_samples <= 0:
+            parser.error("--fresh-video-local-samples must be positive")
+        config.fresh_video_local_samples = args.fresh_video_local_samples
+    if args.total_batch_size is not None:
+        if args.total_batch_size <= 0:
+            parser.error("--total-batch-size must be positive")
+        config.total_batch_size = args.total_batch_size
+    if args.eval_interval is not None:
+        if args.eval_interval < 0:
+            parser.error("--eval-interval must be non-negative")
+        config.eval_interval = args.eval_interval
+
+    # get the filename of config_path
+    config_name = os.path.basename(args.config_path).split(".")[0]
+    config.config_name = config_name
+    # Preserve the exact YAML identity for strict offline-reference checks.
+    # Runtime-only smoke overrides remain separate and do not change the
+    # generator trajectory contract encoded by this file.
+    config.config_path = os.path.realpath(args.config_path)
+    config.logdir = args.logdir
+    config.wandb_save_dir = args.wandb_save_dir
+    config.disable_wandb = args.disable_wandb
+
+    if config.trainer != "score_distillation":
+        raise ValueError("This release supports only three-encoder distribution matching")
+    from trainer.distribution_matching import DistributionMatchingTrainer
+    trainer = DistributionMatchingTrainer(config)
+    try:
+        trainer.train()
+    finally:
+        wandb.finish()
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

1.3B/code/trainer/__init__.py

@@ -0,0 +1,31 @@
+"""Lazy public trainer exports.
+
+Loading one trainer must not import optional dependencies belonging only to a
+different experiment.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+
+
+_EXPORTS = {
+    "DiffusionTrainer": (".diffusion", "Trainer"),
+    "GANTrainer": (".gan", "Trainer"),
+    "ODETrainer": (".ode", "Trainer"),
+    "ScoreDistillationTrainer": (".distillation", "Trainer"),
+    "ARLStage1Trainer": (".arl", "ARLStage1Trainer"),
+    "ARLStage2Trainer": (".arl", "ARLStage2Trainer"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name):
+    try:
+        module_name, attribute = _EXPORTS[name]
+    except KeyError as exc:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+    value = getattr(import_module(module_name, __name__), attribute)
+    globals()[name] = value
+    return value

1.3B/code/trainer/distillation.py

@@ -0,0 +1,6 @@
+"""Compatibility import; read trainer/distribution_matching.py for the implementation."""
+import sys
+from . import distribution_matching as _implementation
+
+# Share the module object so existing monkeypatches and private imports still work.
+sys.modules[__name__] = _implementation

1.3B/code/trainer/feature_gradient_replay.py

@@ -0,0 +1,414 @@
+"""Separate distribution estimation from generator backpropagation (Sec. 3.3).
+
+The distribution batch B supplies all cross-video interactions. Only the smaller
+configured differentiation batch b needs feature gradients and rollout replay.
+This module caches features and returns their cotangents (dL/d feature); the
+trainer later replays the encoders, VAE, and generator to obtain dL/d parameters.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Mapping, Sequence
+
+import torch
+import torch.distributed as dist
+
+from losses.multi_encoder_mmd import (
+    MultiEncoderMMD,
+)
+
+
+def _detach(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach()
+    if isinstance(value, Mapping):
+        return {key: _detach(item) for key, item in value.items()}
+    if isinstance(value, tuple):
+        return tuple(_detach(item) for item in value)
+    if isinstance(value, list):
+        return [_detach(item) for item in value]
+    return value
+
+
+@dataclass(frozen=True)
+class FeatureGradientReplayResult:
+    """Detached loss, per-microbatch feature gradients, and replay bookkeeping.
+
+    ``local_cotangents`` are the derivatives of the global distribution loss.
+    ``fsdp_cotangents`` additionally compensate for FSDP's gradient averaging.
+    Field names stay compatible with existing training logs and checkpoints.
+    """
+
+    loss: torch.Tensor
+    optimized_loss: torch.Tensor
+    metadata: Mapping[str, Any]
+    local_cotangents: tuple[Mapping[str, torch.Tensor], ...]
+    fsdp_cotangents: tuple[Mapping[str, torch.Tensor], ...]
+    replay_batch_indices: tuple[int, ...]
+    replay_sample_ids: tuple[tuple[str, ...], ...]
+    local_samples: int
+    global_samples: int
+    gradient_local_samples: int
+    gradient_global_samples: int
+    current_teacher_local_samples: int
+    current_teacher_global_samples: int
+    world_size: int
+
+
+class FeatureGradientReplay:
+    """Cache B generated feature sets and differentiate the selected b samples.
+
+    B and b denote global sample counts; constructor counts are per rank. Each
+    encoder keeps its own [tokens, channels] layout while sharing sample order.
+    This implementation selects a per-rank suffix. All B samples still define
+    the distribution loss, including interactions with unselected samples.
+    """
+
+    def __init__(
+        self,
+        loss_module: MultiEncoderMMD,
+        *,
+        target_local_samples: int,
+        gradient_local_samples: int,
+        loss_scale: float = 1.0,
+        require_sample_id_alignment: bool = True,
+        full_student_partial_gradient: bool = True,
+    ) -> None:
+        if not isinstance(loss_module, MultiEncoderMMD):
+            raise TypeError(
+                "loss_module must be MultiRepresentationVideoNystromMMD")
+        if not 0 < int(gradient_local_samples) < int(target_local_samples):
+            raise ValueError(
+                "dual GradCache requires a strict positive gradient suffix")
+        if not full_student_partial_gradient:
+            raise ValueError(
+                "dual GradCache requires full_student_partial_gradient=true")
+        self.loss_module = loss_module
+        self.target_local_samples = int(target_local_samples)
+        self.gradient_local_samples = int(gradient_local_samples)
+        self.loss_scale = float(loss_scale)
+        self.require_sample_id_alignment = bool(require_sample_id_alignment)
+        self.full_student_partial_gradient = True
+        self.partial_gradient_enabled = True
+        self.current_teacher_attraction_enabled = True
+        self._feature_batches: list[dict[str, torch.Tensor]] = []
+        self._teacher_batches: list[dict[str, torch.Tensor]] = []
+        self._batch_sizes: list[int] = []
+        self._gradient_batch_indices: list[int] = []
+        self._gradient_batch_sizes: list[int] = []
+        self._gradient_sample_ids: list[tuple[str, ...]] = []
+        self._selection_started = False
+        self._device: torch.device | None = None
+        self._result: FeatureGradientReplayResult | None = None
+
+    @staticmethod
+    def _world_size() -> int:
+        return dist.get_world_size() if dist.is_initialized() else 1
+
+    @staticmethod
+    def _rank() -> int:
+        return dist.get_rank() if dist.is_initialized() else 0
+
+    @property
+    def local_samples(self) -> int:
+        return sum(self._batch_sizes)
+
+    @property
+    def selected_local_samples(self) -> int:
+        return sum(self._gradient_batch_sizes)
+
+    @property
+    def finalized(self) -> bool:
+        return self._result is not None
+
+    def _validate_mapping(
+        self,
+        value: Mapping[str, torch.Tensor],
+        label: str,
+        *,
+        expected_batch: int | None = None,
+    ) -> tuple[int, torch.device]:
+        if not isinstance(value, Mapping):
+            raise TypeError(f"{label} must be a branch tensor mapping")
+        if set(value) != set(self.loss_module.branch_names):
+            raise ValueError(
+                f"{label} keys must be {self.loss_module.branch_names}")
+        batch_size = None
+        device = None
+        for name in self.loss_module.branch_names:
+            tensor = value[name]
+            if (
+                not isinstance(tensor, torch.Tensor)
+                or tensor.ndim != 3
+                or not tensor.is_floating_point()
+                or min(tensor.shape) <= 0
+            ):
+                raise ValueError(f"{label}[{name!r}] must be floating [B,N,D]")
+            if tuple(tensor.shape[1:]) != self.loss_module.sample_shapes[name]:
+                raise ValueError(
+                    f"{label}[{name!r}] geometry {tuple(tensor.shape[1:])} "
+                    f"!= {self.loss_module.sample_shapes[name]}")
+            if not bool(torch.isfinite(tensor).all().item()):
+                raise FloatingPointError(f"{label}[{name!r}] is non-finite")
+            if batch_size is None:
+                batch_size = int(tensor.shape[0])
+                device = tensor.device
+            elif int(tensor.shape[0]) != batch_size or tensor.device != device:
+                raise ValueError(
+                    f"all {label} branches must share batch size and device")
+        if expected_batch is not None and batch_size != expected_batch:
+            raise ValueError(
+                f"{label} batch size {batch_size} != {expected_batch}")
+        return int(batch_size), device
+
+    def add(
+        self,
+        generated_features: Mapping[str, torch.Tensor],
+        current_teacher_features: Mapping[str, torch.Tensor] | None = None,
+        *,
+        gradient_selected: bool = True,
+        sample_ids: Sequence[str] | None = None,
+        current_teacher_sample_ids: Sequence[str] | None = None,
+    ) -> int:
+        """Store one rollout's features without retaining its generator graph."""
+        if self.finalized:
+            raise RuntimeError("cannot add after finalize")
+        batch_size, device = self._validate_mapping(
+            generated_features, "generated_features")
+        if self._device is None:
+            self._device = device
+        elif device != self._device:
+            raise ValueError("all microbatches must share a device")
+        if self.local_samples + batch_size > self.target_local_samples:
+            raise ValueError("dual GradCache exceeds target_local_samples")
+        if gradient_selected:
+            self._selection_started = True
+            if self.selected_local_samples + batch_size > self.gradient_local_samples:
+                raise ValueError("dual GradCache exceeds gradient suffix")
+        elif self._selection_started:
+            raise ValueError("gradient-selected microbatches must form a suffix")
+
+        normalized_ids = None
+        if sample_ids is not None:
+            normalized_ids = tuple(str(value) for value in sample_ids)
+            if len(normalized_ids) != batch_size or any(not value for value in normalized_ids):
+                raise ValueError("sample_ids must contain one non-empty ID per row")
+        if gradient_selected and self.require_sample_id_alignment:
+            if normalized_ids is None:
+                raise ValueError("selected dual samples require sample_ids")
+            teacher_ids = (
+                tuple(str(value) for value in current_teacher_sample_ids)
+                if current_teacher_sample_ids is not None else None
+            )
+            if teacher_ids != normalized_ids:
+                raise ValueError("student and teacher sample IDs/order differ")
+
+        if gradient_selected:
+            if current_teacher_features is None:
+                raise ValueError("selected dual samples require teacher features")
+            self._validate_mapping(
+                current_teacher_features,
+                "current_teacher_features",
+                expected_batch=batch_size,
+            )
+        elif current_teacher_features is not None:
+            raise ValueError("prefix microbatches must not carry teacher features")
+
+        # Collection stage (Algorithm 1): keep only feature values and sample
+        # order. Replay state is stored by the trainer, outside this cache.
+        staged = {
+            name: generated_features[name].detach().float().clone()
+            for name in self.loss_module.branch_names
+        }
+        self._feature_batches.append(staged)
+        index = len(self._feature_batches) - 1
+        if gradient_selected:
+            self._gradient_batch_indices.append(index)
+            self._gradient_batch_sizes.append(batch_size)
+            self._gradient_sample_ids.append(normalized_ids)
+            self._teacher_batches.append({
+                name: current_teacher_features[name].detach().float().clone()
+                for name in self.loss_module.branch_names
+            })
+        self._batch_sizes.append(batch_size)
+        return index
+
+    def _gather(self, value: torch.Tensor) -> list[torch.Tensor]:
+        if self._world_size() == 1:
+            return [value.detach()]
+        gathered = [torch.empty_like(value) for _ in range(self._world_size())]
+        dist.all_gather(gathered, value.detach().contiguous())
+        return gathered
+
+    def _validate_ready(self) -> None:
+        return None
+
+    def finalize(self) -> FeatureGradientReplayResult:
+        """Evaluate the global objective and cache dL/d feature for replay."""
+        if self._result is not None:
+            return self._result
+        self._validate_ready()
+        world_size = self._world_size()
+        global_distribution_features = {}
+        global_differentiation_features = {}
+        global_reference_features = {}
+        local_leaves = {}
+        # Map each rank's selected suffix into the concatenated B-sample batch.
+        # These indices preserve full-batch interactions while only b rows have
+        # an autograd path. This function does not randomly resample the subset.
+        suffix_start = self.target_local_samples - self.gradient_local_samples
+        gradient_indices = torch.cat([
+            torch.arange(
+                rank * self.target_local_samples + suffix_start,
+                (rank + 1) * self.target_local_samples,
+                device=self._device,
+                dtype=torch.long,
+            )
+            for rank in range(world_size)
+        ])
+        for name in self.loss_module.branch_names:
+            local_distribution_features = torch.cat([
+                batch[name] for batch in self._feature_batches
+            ], dim=0)
+            local_differentiation_features = torch.cat([
+                self._feature_batches[index][name]
+                for index in self._gradient_batch_indices
+            ], dim=0)
+            local_reference_features = torch.cat([
+                batch[name] for batch in self._teacher_batches
+            ], dim=0)
+            distribution_parts = self._gather(local_distribution_features)
+            differentiation_parts = self._gather(local_differentiation_features)
+            reference_parts = self._gather(local_reference_features)
+            # Only this rank's selected features are differentiable leaves.
+            # Remote features contribute values; their owning rank gets their
+            # derivatives. No generator or encoder graph is retained here.
+            leaf = local_differentiation_features.detach().requires_grad_(True)
+            differentiation_parts[self._rank()] = leaf
+            local_leaves[name] = leaf
+            global_distribution_features[name] = torch.cat(distribution_parts, dim=0).detach()
+            global_differentiation_features[name] = torch.cat(differentiation_parts, dim=0)
+            global_reference_features[name] = torch.cat(reference_parts, dim=0).detach()
+
+        # Resolve all cross-video interactions once. The loss combines the
+        # persistent Nystrom summary with current reference features (Eq. 9),
+        # then aggregates the encoder objectives (Eqs. 11 and 13).
+        raw_loss, metadata = self.loss_module(
+            gradient_students=global_differentiation_features,
+            student_reference=global_distribution_features,
+            current_teacher=global_reference_features,
+            gradient_student_indices=gradient_indices,
+        )
+        optimized = raw_loss * self.loss_scale
+        gradients = torch.autograd.grad(
+            optimized,
+            tuple(local_leaves[name] for name in self.loss_module.branch_names),
+            retain_graph=False,
+            create_graph=False,
+        )
+        # Restore the collection order so each feature gradient is replayed
+        # through the exact sample and encoder that produced its feature value.
+        split_by_branch = {
+            name: tuple(part.contiguous() for part in gradient.detach().split(
+                self._gradient_batch_sizes, dim=0))
+            for name, gradient in zip(self.loss_module.branch_names, gradients)
+        }
+        local_parts = tuple({
+            name: split_by_branch[name][index]
+            for name in self.loss_module.branch_names
+        } for index in range(len(self._gradient_batch_sizes)))
+        # FSDP averages parameter gradients across ranks, so cancel that factor
+        # here. This world_size multiplier is separate from the B/b selective
+        # differentiation correction (Eq. 12), applied later by the trainer.
+        fsdp_parts = tuple({
+            name: (value * float(world_size)).contiguous()
+            for name, value in part.items()
+        } for part in local_parts)
+        global_samples = self.target_local_samples * world_size
+        gradient_global = self.gradient_local_samples * world_size
+        kss_denominator = global_samples * global_samples
+        # Keep established metadata names for checkpoint/log compatibility.
+        detached_metadata = dict(_detach(metadata))
+        detached_metadata.update({
+            "raw_loss": raw_loss.detach(),
+            "optimized_loss": optimized.detach(),
+            "loss_scale": raw_loss.new_tensor(self.loss_scale),
+            "gradcache_local_samples": raw_loss.new_tensor(
+                self.target_local_samples, dtype=torch.long),
+            "gradcache_global_samples": raw_loss.new_tensor(
+                global_samples, dtype=torch.long),
+            "gradcache_forward_local_samples": raw_loss.new_tensor(
+                self.target_local_samples, dtype=torch.long),
+            "gradcache_forward_global_samples": raw_loss.new_tensor(
+                global_samples, dtype=torch.long),
+            "gradcache_gradient_local_samples": raw_loss.new_tensor(
+                self.gradient_local_samples, dtype=torch.long),
+            "gradcache_gradient_global_samples": raw_loss.new_tensor(
+                gradient_global, dtype=torch.long),
+            "gradcache_current_teacher_local_samples": raw_loss.new_tensor(
+                self.gradient_local_samples, dtype=torch.long),
+            "gradcache_current_teacher_global_samples": raw_loss.new_tensor(
+                gradient_global, dtype=torch.long),
+            "gradcache_kss_denominator": raw_loss.new_tensor(
+                kss_denominator, dtype=torch.long),
+            "gradcache_partial_gradient_enabled": raw_loss.new_ones(
+                (), dtype=torch.long),
+            "gradcache_full_student_partial_gradient": raw_loss.new_ones(
+                (), dtype=torch.long),
+            "gradcache_world_size": raw_loss.new_tensor(
+                world_size, dtype=torch.long),
+            "gradcache_fsdp_cotangent_scale": raw_loss.new_tensor(
+                float(world_size)),
+            "gradcache_branch_names": self.loss_module.branch_names,
+        })
+        self._result = FeatureGradientReplayResult(
+            loss=raw_loss.detach(),
+            optimized_loss=optimized.detach(),
+            metadata=detached_metadata,
+            local_cotangents=local_parts,
+            fsdp_cotangents=fsdp_parts,
+            replay_batch_indices=tuple(self._gradient_batch_indices),
+            replay_sample_ids=tuple(self._gradient_sample_ids),
+            local_samples=self.target_local_samples,
+            global_samples=global_samples,
+            gradient_local_samples=self.gradient_local_samples,
+            gradient_global_samples=gradient_global,
+            current_teacher_local_samples=self.gradient_local_samples,
+            current_teacher_global_samples=gradient_global,
+            world_size=world_size,
+        )
+        return self._result
+
+    def _clear(self) -> None:
+        self._feature_batches.clear()
+        self._teacher_batches.clear()
+        self._batch_sizes.clear()
+        self._gradient_batch_indices.clear()
+        self._gradient_batch_sizes.clear()
+        self._gradient_sample_ids.clear()
+        self._selection_started = False
+        self._device = None
+        self._result = None
+
+    def commit(self) -> None:
+        """Release cached values after the corresponding optimizer update."""
+        if self._result is None:
+            raise RuntimeError("cannot commit before finalize")
+        self._clear()
+
+    def discard(self) -> None:
+        """Release an abandoned batch without marking an optimizer update."""
+        self._clear()
+
+
+# Preserve old imports and serialized class references during the naming cleanup.
+MultiRepresentationFreshVideoGradCache = FeatureGradientReplay
+MultiRepresentationFreshVideoGradCacheResult = FeatureGradientReplayResult
+
+__all__ = [
+    "FeatureGradientReplay",
+    "FeatureGradientReplayResult",
+    "MultiRepresentationFreshVideoGradCache",
+    "MultiRepresentationFreshVideoGradCacheResult",
+]

1.3B/code/trainer/fresh_video_gradcache.py

@@ -0,0 +1,634 @@
+"""Fresh-update gradient cache for video-level distribution losses.
+
+The cache deliberately has no state across optimizer updates.  It collects a
+fixed number of detached ``[B, N, D]`` feature matrices on every rank, builds
+the exact global fresh batch in deterministic rank order, and differentiates
+the frozen reference loss only with respect to the local rank's feature leaf.
+
+The two returned cotangents have distinct meanings:
+
+``local_cotangents``
+    The mathematical derivatives ``d L_global / d feature_local``.
+
+``fsdp_cotangents``
+    ``world_size * local_cotangents``.  These are the cotangents that replay
+    should use when FSDP/DDP will average parameter gradients across ranks:
+
+    ``mean_r(world_size * J_r.T @ dL/dF_r) = sum_r J_r.T @ dL/dF_r``.
+
+Collection is transactional.  ``finalize`` stages a result but retains the
+fresh features until the caller reports optimizer success with ``commit`` or
+failure with ``discard``.  Both terminal operations clear all update-local
+state, so no generated sample can leak into the next optimizer update.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+from typing import Any, Mapping, Sequence
+
+import torch
+import torch.distributed as dist
+
+from losses.video_nystrom_mmd import VideoNystromMMD
+from losses.multilayer_video_nystrom_mmd import (
+    MultiLayerVideoNystromMMD,
+)
+
+
+def _detach_metadata(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach()
+    if isinstance(value, Mapping):
+        return {
+            key: _detach_metadata(item)
+            for key, item in value.items()
+        }
+    if isinstance(value, tuple):
+        return tuple(_detach_metadata(item) for item in value)
+    if isinstance(value, list):
+        return [_detach_metadata(item) for item in value]
+    return value
+
+
+@dataclass(frozen=True)
+class FreshVideoGradCacheResult:
+    """Finalized loss and replay cotangents for one fresh optimizer update."""
+
+    loss: torch.Tensor
+    optimized_loss: torch.Tensor
+    metadata: Mapping[str, Any]
+    local_cotangents: tuple[torch.Tensor, ...]
+    fsdp_cotangents: tuple[torch.Tensor, ...]
+    replay_batch_indices: tuple[int, ...]
+    replay_sample_ids: tuple[tuple[str, ...], ...]
+    local_samples: int
+    global_samples: int
+    gradient_local_samples: int
+    gradient_global_samples: int
+    current_teacher_local_samples: int
+    current_teacher_global_samples: int
+    world_size: int
+
+
+class FreshVideoGradCache:
+    """Collect and differentiate one globally fresh video feature batch."""
+
+    def __init__(
+        self,
+        loss_module: VideoNystromMMD | MultiLayerVideoNystromMMD,
+        *,
+        target_local_samples: int = 16,
+        gradient_local_samples: int | None = None,
+        loss_scale: float = 1.0,
+        require_sample_id_alignment: bool = False,
+        full_student_partial_gradient: bool = False,
+    ) -> None:
+        if not isinstance(
+            loss_module,
+            (VideoNystromMMD, MultiLayerVideoNystromMMD),
+        ):
+            raise TypeError(
+                "loss_module must be a VideoNystromMMD or "
+                "MultiLayerVideoNystromMMD instance")
+        if int(target_local_samples) <= 0:
+            raise ValueError("target_local_samples must be positive")
+        if gradient_local_samples is None:
+            gradient_local_samples = int(target_local_samples)
+        if not 0 < int(gradient_local_samples) <= int(target_local_samples):
+            raise ValueError(
+                "gradient_local_samples must be in [1, target_local_samples]")
+        if (
+            not math.isfinite(float(loss_scale))
+            or float(loss_scale) < 0.0
+        ):
+            raise ValueError("loss_scale must be finite and non-negative")
+        if any(parameter.requires_grad for parameter in loss_module.parameters()):
+            raise ValueError(
+                "VideoNystromMMD reference parameters must be frozen")
+
+        self.loss_module = loss_module
+        self.current_teacher_attraction_enabled = bool(getattr(
+            loss_module, "current_teacher_attraction_enabled", False))
+        self.target_local_samples = int(target_local_samples)
+        self.gradient_local_samples = int(gradient_local_samples)
+        self.partial_gradient_enabled = bool(
+            self.gradient_local_samples < self.target_local_samples)
+        self.full_student_partial_gradient = bool(
+            full_student_partial_gradient)
+        if (
+            self.full_student_partial_gradient
+            and not self.partial_gradient_enabled
+        ):
+            raise ValueError(
+                "full_student_partial_gradient requires "
+                "gradient_local_samples < target_local_samples")
+        if (
+            self.partial_gradient_enabled
+            and isinstance(loss_module, MultiLayerVideoNystromMMD)
+        ):
+            raise TypeError(
+                "partial-gradient GradCache currently requires a single "
+                "VideoNystromMMD layer")
+        self.loss_scale = float(loss_scale)
+        self.require_sample_id_alignment = bool(require_sample_id_alignment)
+        self._feature_batches: list[torch.Tensor] = []
+        self._teacher_feature_batches: list[torch.Tensor] = []
+        self._batch_sizes: list[int] = []
+        self._gradient_batch_indices: list[int] = []
+        self._gradient_batch_sizes: list[int] = []
+        self._gradient_sample_ids: list[tuple[str, ...]] = []
+        self._selection_started = False
+        self._sample_shape: tuple[int, int] | None = None
+        self._device: torch.device | None = None
+        self._result: FreshVideoGradCacheResult | None = None
+
+    @staticmethod
+    def _world_size() -> int:
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_world_size()
+        return 1
+
+    @staticmethod
+    def _rank() -> int:
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_rank()
+        return 0
+
+    @staticmethod
+    def _all_ranks_true(local_value: bool, device: torch.device) -> bool:
+        flag = torch.tensor(
+            int(bool(local_value)), device=device, dtype=torch.int32)
+        if (
+            dist.is_available()
+            and dist.is_initialized()
+            and dist.get_world_size() > 1
+        ):
+            dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+        return bool(flag.item())
+
+    @property
+    def local_samples(self) -> int:
+        return sum(self._batch_sizes)
+
+    @property
+    def num_microbatches(self) -> int:
+        return len(self._feature_batches)
+
+    @property
+    def selected_local_samples(self) -> int:
+        return sum(self._gradient_batch_sizes)
+
+    @property
+    def ready(self) -> bool:
+        return (
+            self.local_samples == self.target_local_samples
+            and self.selected_local_samples == self.gradient_local_samples
+        )
+
+    @property
+    def finalized(self) -> bool:
+        return self._result is not None
+
+    def add(
+        self,
+        generated_features: torch.Tensor,
+        current_teacher_features: torch.Tensor | None = None,
+        *,
+        gradient_selected: bool | None = None,
+        sample_ids: Sequence[str] | None = None,
+        current_teacher_sample_ids: Sequence[str] | None = None,
+    ) -> int:
+        """Add one detached microbatch and return its zero-based cache index."""
+        if self.finalized:
+            raise RuntimeError(
+                "cannot add features after finalize; commit or discard first")
+        if not isinstance(generated_features, torch.Tensor):
+            raise TypeError("generated_features must be a torch.Tensor")
+        if (
+            generated_features.ndim != 3
+            or min(generated_features.shape) <= 0
+            or not generated_features.is_floating_point()
+        ):
+            raise ValueError(
+                "generated_features must be a non-empty floating [B, N, D] "
+                "tensor")
+        if not bool(torch.isfinite(generated_features).all().item()):
+            raise FloatingPointError(
+                "generated_features contains non-finite values")
+
+        if gradient_selected is None:
+            if self.partial_gradient_enabled:
+                raise TypeError(
+                    "gradient_selected must be explicit when partial-gradient "
+                    "GradCache is enabled")
+            gradient_selected = True
+        gradient_selected = bool(gradient_selected)
+        if self._selection_started and not gradient_selected:
+            raise ValueError(
+                "gradient-selected microbatches must form a contiguous suffix")
+        if gradient_selected:
+            self._selection_started = True
+
+        batch_size = int(generated_features.shape[0])
+        normalized_sample_ids: tuple[str, ...] | None = None
+        normalized_teacher_ids: tuple[str, ...] | None = None
+        if sample_ids is not None:
+            normalized_sample_ids = tuple(str(value) for value in sample_ids)
+            if len(normalized_sample_ids) != batch_size or any(
+                not value for value in normalized_sample_ids
+            ):
+                raise ValueError(
+                    "sample_ids must contain one non-empty ID per student")
+        if current_teacher_sample_ids is not None:
+            normalized_teacher_ids = tuple(
+                str(value) for value in current_teacher_sample_ids)
+            if len(normalized_teacher_ids) != batch_size or any(
+                not value for value in normalized_teacher_ids
+            ):
+                raise ValueError(
+                    "current_teacher_sample_ids must contain one non-empty ID "
+                    "per selected student")
+        if self.require_sample_id_alignment and normalized_sample_ids is None:
+            raise ValueError(
+                "sample_ids are required by this GradCache experiment")
+
+        if self.current_teacher_attraction_enabled and gradient_selected:
+            if not isinstance(current_teacher_features, torch.Tensor):
+                raise TypeError(
+                    "current_teacher_features must be a torch.Tensor for each "
+                    "selected student microbatch")
+            if (
+                current_teacher_features.ndim != 3
+                or min(current_teacher_features.shape) <= 0
+                or not current_teacher_features.is_floating_point()
+            ):
+                raise ValueError(
+                    "current_teacher_features must be a non-empty floating "
+                    "[B, N, D] tensor")
+            if tuple(current_teacher_features.shape) != tuple(
+                generated_features.shape
+            ):
+                raise ValueError(
+                    "current teacher features must exactly match selected "
+                    "student [B,N,D]: "
+                    f"{tuple(current_teacher_features.shape)} != "
+                    f"{tuple(generated_features.shape)}")
+            if current_teacher_features.device != generated_features.device:
+                raise ValueError(
+                    "generated and current teacher features must share a device")
+            if not bool(torch.isfinite(
+                current_teacher_features.detach()
+            ).all().item()):
+                raise FloatingPointError(
+                    "current_teacher_features contains non-finite values")
+            if self.require_sample_id_alignment and normalized_teacher_ids is None:
+                raise ValueError(
+                    "current_teacher_sample_ids are required for selected "
+                    "students")
+            if (
+                normalized_teacher_ids is not None
+                and normalized_sample_ids != normalized_teacher_ids
+            ):
+                raise ValueError(
+                    "matched teacher sample IDs must exactly equal selected "
+                    "student sample IDs in order")
+        elif current_teacher_features is not None:
+            raise ValueError(
+                "current_teacher_features may only be provided for a selected "
+                "microbatch while exact current-teacher attraction is enabled")
+        elif current_teacher_sample_ids is not None:
+            raise ValueError(
+                "current_teacher_sample_ids may only be provided for a "
+                "selected current-teacher microbatch")
+
+        if self.local_samples + batch_size > self.target_local_samples:
+            raise ValueError(
+                "fresh feature cache would exceed target_local_samples: "
+                f"{self.local_samples} + {batch_size} > "
+                f"{self.target_local_samples}")
+        if (
+            gradient_selected
+            and self.selected_local_samples + batch_size
+            > self.gradient_local_samples
+        ):
+            raise ValueError(
+                "fresh feature cache would exceed gradient_local_samples: "
+                f"{self.selected_local_samples} + {batch_size} > "
+                f"{self.gradient_local_samples}")
+
+        sample_shape = (
+            int(generated_features.shape[1]),
+            int(generated_features.shape[2]),
+        )
+        if self._sample_shape is None:
+            self._sample_shape = sample_shape
+            self._device = generated_features.device
+        elif (
+            sample_shape != self._sample_shape
+            or generated_features.device != self._device
+        ):
+            raise ValueError(
+                "all fresh feature batches on a rank must share [N, D] and "
+                "device; expected "
+                f"{self._sample_shape} on {self._device}, got "
+                f"{sample_shape} on {generated_features.device}")
+
+        staged = generated_features.detach().to(dtype=torch.float32).clone()
+        self._feature_batches.append(staged)
+        batch_index = len(self._feature_batches) - 1
+        if gradient_selected:
+            self._gradient_batch_indices.append(batch_index)
+            self._gradient_batch_sizes.append(batch_size)
+            self._gradient_sample_ids.append(
+                normalized_sample_ids or tuple("" for _ in range(batch_size)))
+        if self.current_teacher_attraction_enabled and gradient_selected:
+            staged_teacher = current_teacher_features.detach().to(
+                dtype=torch.float32).clone()
+            self._teacher_feature_batches.append(staged_teacher)
+        self._batch_sizes.append(batch_size)
+        return batch_index
+
+    def _validate_distributed_geometry(
+        self,
+        local_features: torch.Tensor,
+        *,
+        expected_samples: int,
+        label: str,
+    ) -> None:
+        return None
+
+    def _validate_ready_consensus(self) -> None:
+        return None
+
+    def _gather_detached(
+        self,
+        local_features: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        if self._world_size() == 1:
+            return [local_features.detach()]
+        gathered = [
+            torch.empty_like(local_features)
+            for _ in range(self._world_size())
+        ]
+        dist.all_gather(gathered, local_features.detach().contiguous())
+        return gathered
+
+    def finalize(self) -> FreshVideoGradCacheResult:
+        """Compute the global loss and stage suffix-only replay cotangents."""
+        if self._result is not None:
+            return self._result
+        self._validate_ready_consensus()
+
+        local_reference = torch.cat(self._feature_batches, dim=0)
+        self._validate_distributed_geometry(
+            local_reference,
+            expected_samples=self.target_local_samples,
+            label="student reference",
+        )
+        gathered_reference = self._gather_detached(local_reference)
+
+        if self.partial_gradient_enabled:
+            selected_batches = [
+                self._feature_batches[index]
+                for index in self._gradient_batch_indices
+            ]
+            local_gradient = torch.cat(selected_batches, dim=0)
+            self._validate_distributed_geometry(
+                local_gradient,
+                expected_samples=self.gradient_local_samples,
+                label="gradient students",
+            )
+            gathered_gradient = self._gather_detached(local_gradient)
+            local_leaf = local_gradient.detach().requires_grad_(True)
+            gathered_gradient[self._rank()] = local_leaf
+            global_gradient = torch.cat(gathered_gradient, dim=0)
+            global_reference = torch.cat(
+                gathered_reference, dim=0).detach()
+            gradient_student_indices = None
+            if self.full_student_partial_gradient:
+                # ``global_gradient`` is rank-major, with each rank
+                # contributing its contiguous selected suffix. Build the
+                # matching positions in the rank-major full student batch so
+                # the objective can inject every selected value back into its
+                # original location without reordering either population.
+                suffix_start = (
+                    self.target_local_samples
+                    - self.gradient_local_samples
+                )
+                gradient_student_indices = torch.cat([
+                    torch.arange(
+                        rank * self.target_local_samples + suffix_start,
+                        (rank + 1) * self.target_local_samples,
+                        device=global_reference.device,
+                        dtype=torch.long,
+                    )
+                    for rank in range(self._world_size())
+                ])
+        else:
+            local_leaf = local_reference.detach().requires_grad_(True)
+            gathered_reference[self._rank()] = local_leaf
+            global_gradient = torch.cat(gathered_reference, dim=0)
+            global_reference = None
+
+        global_teacher = None
+        if self.current_teacher_attraction_enabled:
+            if len(self._teacher_feature_batches) != len(
+                self._gradient_batch_indices
+            ):
+                raise RuntimeError(
+                    "matched teacher and selected student GradCache batch "
+                    "counts differ")
+            local_teacher = torch.cat(
+                self._teacher_feature_batches, dim=0)
+            if tuple(local_teacher.shape) != tuple(
+                local_gradient.shape
+                if self.partial_gradient_enabled else local_reference.shape
+            ):
+                raise RuntimeError(
+                    "matched teacher and gradient student shapes differ")
+            self._validate_distributed_geometry(
+                local_teacher,
+                expected_samples=self.gradient_local_samples,
+                label="current teacher",
+            )
+            gathered_teacher = self._gather_detached(local_teacher)
+            global_teacher = torch.cat(
+                gathered_teacher, dim=0).detach().to(dtype=torch.float32)
+
+        if self.partial_gradient_enabled:
+            if isinstance(self.loss_module, MultiLayerVideoNystromMMD):
+                raise TypeError(
+                    "partial-gradient GradCache currently requires a single "
+                    "VideoNystromMMD layer")
+            partial_loss_kwargs = {
+                "gradient_students": global_gradient,
+                "student_reference": global_reference,
+                "current_teacher": global_teacher,
+            }
+            if self.full_student_partial_gradient:
+                partial_loss_kwargs["gradient_student_indices"] = (
+                    gradient_student_indices)
+            raw_loss, metadata = self.loss_module(**partial_loss_kwargs)
+        elif self.current_teacher_attraction_enabled:
+            raw_loss, metadata = self.loss_module(
+                global_gradient, current_teacher=global_teacher)
+        else:
+            raw_loss, metadata = self.loss_module(global_gradient)
+        if not isinstance(raw_loss, torch.Tensor) or raw_loss.ndim != 0:
+            raise TypeError(
+                "VideoNystromMMD must return a scalar tensor loss")
+        if not self._all_ranks_true(
+            bool(torch.isfinite(raw_loss.detach()).item()),
+            raw_loss.device,
+        ):
+            raise FloatingPointError(
+                "VideoNystromMMD returned a non-finite loss on at least one "
+                "rank")
+        optimized_loss = raw_loss * self.loss_scale
+
+        (local_cotangent,) = torch.autograd.grad(
+            optimized_loss,
+            local_leaf,
+            retain_graph=False,
+            create_graph=False,
+        )
+        local_cotangent = local_cotangent.detach()
+        if not self._all_ranks_true(
+            bool(torch.isfinite(local_cotangent).all().item()),
+            local_cotangent.device,
+        ):
+            raise FloatingPointError(
+                "VideoNystromMMD returned a non-finite feature cotangent on "
+                "at least one rank")
+
+        world_size = self._world_size()
+        local_parts = tuple(
+            value.contiguous()
+            for value in local_cotangent.split(
+                self._gradient_batch_sizes, dim=0)
+        )
+        fsdp_parts = tuple(
+            (value * float(world_size)).contiguous()
+            for value in local_parts
+        )
+        global_samples = self.target_local_samples * world_size
+        gradient_global_samples = self.gradient_local_samples * world_size
+        teacher_local_samples = (
+            self.gradient_local_samples
+            if self.current_teacher_attraction_enabled else 0
+        )
+        teacher_global_samples = teacher_local_samples * world_size
+        kss_denominator = global_samples * (
+            global_samples
+            if self.full_student_partial_gradient
+            else gradient_global_samples
+        )
+        detached_metadata = dict(_detach_metadata(metadata))
+        detached_metadata.update({
+            "raw_loss": raw_loss.detach(),
+            "optimized_loss": optimized_loss.detach(),
+            "loss_scale": raw_loss.new_tensor(self.loss_scale),
+            "gradcache_local_samples": raw_loss.new_tensor(
+                self.target_local_samples, dtype=torch.long),
+            "gradcache_global_samples": raw_loss.new_tensor(
+                global_samples, dtype=torch.long),
+            "gradcache_forward_local_samples": raw_loss.new_tensor(
+                self.target_local_samples, dtype=torch.long),
+            "gradcache_forward_global_samples": raw_loss.new_tensor(
+                global_samples, dtype=torch.long),
+            "gradcache_gradient_local_samples": raw_loss.new_tensor(
+                self.gradient_local_samples, dtype=torch.long),
+            "gradcache_gradient_global_samples": raw_loss.new_tensor(
+                gradient_global_samples, dtype=torch.long),
+            "gradcache_kss_denominator": raw_loss.new_tensor(
+                kss_denominator, dtype=torch.long),
+            "gradcache_partial_gradient_enabled": raw_loss.new_tensor(
+                int(self.partial_gradient_enabled), dtype=torch.long),
+            "gradcache_full_student_partial_gradient": raw_loss.new_tensor(
+                int(self.full_student_partial_gradient), dtype=torch.long),
+            "gradcache_world_size": raw_loss.new_tensor(
+                world_size, dtype=torch.long),
+            "gradcache_fsdp_cotangent_scale": raw_loss.new_tensor(
+                float(world_size)),
+        })
+        detached_metadata.update({
+            "gradcache_current_teacher_local_samples": raw_loss.new_tensor(
+                teacher_local_samples, dtype=torch.long),
+            "gradcache_current_teacher_global_samples": raw_loss.new_tensor(
+                teacher_global_samples, dtype=torch.long),
+        })
+        core_global_samples = detached_metadata.get("global_samples")
+        if (
+            core_global_samples is not None
+            and int(torch.as_tensor(core_global_samples).item())
+            != global_samples
+        ):
+            raise ValueError(
+                "VideoNystromMMD global_samples metadata disagrees with the "
+                f"fresh GradCache: {int(torch.as_tensor(core_global_samples))} "
+                f"!= {global_samples}")
+        core_gradient_samples = detached_metadata.get(
+            "gradient_student_samples")
+        if (
+            self.partial_gradient_enabled
+            and core_gradient_samples is not None
+            and int(torch.as_tensor(core_gradient_samples).item())
+            != gradient_global_samples
+        ):
+            raise ValueError(
+                "VideoNystromMMD gradient_student_samples metadata disagrees "
+                "with the fresh GradCache")
+        core_repulsion_denominator = detached_metadata.get(
+            "repulsion_denominator")
+        if (
+            self.partial_gradient_enabled
+            and core_repulsion_denominator is not None
+            and int(torch.as_tensor(core_repulsion_denominator).item())
+            != kss_denominator
+        ):
+            raise ValueError(
+                "VideoNystromMMD Kss denominator disagrees with the fresh "
+                f"GradCache: expected {kss_denominator}")
+
+        self._result = FreshVideoGradCacheResult(
+            loss=raw_loss.detach(),
+            optimized_loss=optimized_loss.detach(),
+            metadata=detached_metadata,
+            local_cotangents=local_parts,
+            fsdp_cotangents=fsdp_parts,
+            replay_batch_indices=tuple(self._gradient_batch_indices),
+            replay_sample_ids=tuple(self._gradient_sample_ids),
+            local_samples=self.target_local_samples,
+            global_samples=global_samples,
+            gradient_local_samples=self.gradient_local_samples,
+            gradient_global_samples=gradient_global_samples,
+            current_teacher_local_samples=teacher_local_samples,
+            current_teacher_global_samples=teacher_global_samples,
+            world_size=world_size,
+        )
+        return self._result
+
+    def _clear(self) -> None:
+        self._feature_batches.clear()
+        self._teacher_feature_batches.clear()
+        self._batch_sizes.clear()
+        self._gradient_batch_indices.clear()
+        self._gradient_batch_sizes.clear()
+        self._gradient_sample_ids.clear()
+        self._selection_started = False
+        self._sample_shape = None
+        self._device = None
+        self._result = None
+
+    def commit(self) -> None:
+        """Clear a finalized cache after a successful optimizer update."""
+        if self._result is None:
+            raise RuntimeError(
+                "cannot commit a fresh feature cache before finalize")
+        self._clear()
+
+    def discard(self) -> None:
+        """Clear collected/finalized data after a skipped or failed update."""
+        self._clear()

1.3B/code/trainer/multi_representation_fresh_video_gradcache.py

@@ -0,0 +1,7 @@
+"""Compatibility module; read ``trainer.feature_gradient_replay`` instead."""
+
+import sys
+from trainer import feature_gradient_replay as _implementation
+
+# Preserve private helpers, monkeypatch targets, and legacy import paths.
+sys.modules[__name__] = _implementation

1.3B/code/utils/config.py

@@ -0,0 +1,83 @@
+"""Readable paper terminology at the config boundary, stable runtime keys inside.
+
+The public YAML uses distribution/differentiation/reference batches. Translate
+once before loading defaults, constructing models, or relocating assets. The
+resolved config retains the released keys and values so checkpoint/reference
+identities and existing scripts remain compatible. Legacy YAML is still valid.
+"""
+from collections.abc import Mapping
+from copy import deepcopy
+
+# New public spelling -> existing runtime/serialized spelling.
+CONFIG_ALIASES = {
+    'initial_generator_checkpoint': 'generator_ckpt',
+    'initial_generator_checkpoint_sha256': 'generator_ckpt_sha256',
+    'training_objective': 'distribution_loss',
+    'training_loop': 'trainer',
+    'feature_space': 'wan_feature_type',
+    'mmd_estimator': 'wan_mmd_estimator',
+    'objective_weight': 'wan_kv_loss_weight',
+    'normalize_encoder_losses': 'wan_kv_self_normalize',
+    'encoder_loss_normalization_epsilon': 'wan_kv_norm_eps',
+    'reference_feature_source': 'current_teacher_source',
+    'reference_feature_manifest': 'triple_teacher_bank_manifest',
+    'reference_feature_manifest_sha256': 'triple_teacher_bank_manifest_sha256',
+    'reference_feature_content_hash': 'triple_teacher_bank_content_hash',
+    'reference_video_manifest': 'matched_manifest_path',
+    'reference_video_count': 'matched_manifest_samples',
+    'nystrom_reference_weight': 'triple_nystrom_attraction_weight',
+    'monte_carlo_reference_weight': 'triple_nystrom_current_teacher_attraction_weight',
+    'dino_loss_weight': 'triple_dinov3_weight',
+    'vjepa_loss_weight': 'triple_vjepa2_weight',
+    'videomae_loss_weight': 'triple_videomae_v2_weight',
+    'encoder_loss_divisor': 'triple_aggregation_divisor',
+    'active_feature_encoders': 'triple_active_branches',
+    'per_device_microbatch_size': 'batch_size',
+    'distribution_samples_per_rank': 'fresh_video_local_samples',
+    'differentiation_samples_per_rank': 'fresh_video_generator_grad_local_samples',
+    'differentiate_full_distribution_loss': 'fresh_video_full_student_partial_gradient',
+    'differentiation_gradient_scale': 'fresh_video_generator_grad_sample_scale',
+    'distribution_batch_size': 'total_batch_size',
+    'reference_minibatch_size': 'expected_current_teacher_samples',
+    'latent_video_shape': 'image_or_video_shape',
+    'training_updates': 'max_steps',
+    'generator_learning_rate': 'lr',
+}
+
+VALUE_ALIASES = {
+    'training_objective': {'distribution_matching': 'wan_kv_mmd'},
+    'training_loop': {'distribution_matching': 'score_distillation'},
+    'feature_space': {'frozen_video_encoders': 'triple_dinov3_vjepa2_videomae_v2_dense_latent'},
+    'mmd_estimator': {'hybrid_nystrom_monte_carlo': 'nystrom_exact_repulsion'},
+    'reference_feature_source': {'matched_encoder_features': 'matched_triple_feature_bank'},
+}
+
+
+def resolve_config_aliases(config: Mapping) -> dict:
+    """Return the released runtime schema without mutating the caller's data.
+
+    An explicitly supplied readable key takes precedence over its legacy alias.
+    This also permits overriding a legacy base config with readable settings.
+    """
+    result = {}
+    overridden = {legacy for readable, legacy in CONFIG_ALIASES.items() if readable in config}
+    for key, value in config.items():
+        if key in CONFIG_ALIASES:
+            legacy = CONFIG_ALIASES[key]
+            replacements = VALUE_ALIASES.get(key, {})
+            result[legacy] = replacements.get(value, value) if isinstance(value, str) else deepcopy(value)
+        elif key not in overridden:
+            result[key] = deepcopy(value)
+    return result
+
+
+def load_config(path, default_path=None):
+    """Resolve public names before merging defaults; preserve OmegaConf behavior."""
+    from omegaconf import OmegaConf
+
+    def read(filename):
+        config = OmegaConf.to_container(OmegaConf.load(filename), resolve=False)
+        return OmegaConf.create(resolve_config_aliases(config))
+
+    config = read(path)
+    return OmegaConf.merge(read(default_path), config) if default_path else config

1.3B/code/utils/dataset.py

@@ -0,0 +1,283 @@
+from utils.lmdb import get_array_shape_from_lmdb, retrieve_row_from_lmdb
+from torch.utils.data import Dataset
+import numpy as np
+import torch
+import lmdb
+import json
+from pathlib import Path
+from PIL import Image
+import os
+
+
+class TextDataset(Dataset):
+    def __init__(self, prompt_path, extended_prompt_path=None):
+        with open(prompt_path, encoding="utf-8") as f:
+            self.prompt_list = [line.rstrip() for line in f]
+
+        if extended_prompt_path is not None:
+            with open(extended_prompt_path, encoding="utf-8") as f:
+                self.extended_prompt_list = [line.rstrip() for line in f]
+            assert len(self.extended_prompt_list) == len(self.prompt_list)
+        else:
+            self.extended_prompt_list = None
+
+    def __len__(self):
+        return len(self.prompt_list)
+
+    def __getitem__(self, idx):
+        batch = {
+            "prompts": self.prompt_list[idx],
+            "idx": idx,
+        }
+        if self.extended_prompt_list is not None:
+            batch["extended_prompts"] = self.extended_prompt_list[idx]
+        return batch
+
+
+class ODERegressionLMDBDataset(Dataset):
+    def __init__(
+        self,
+        data_path: str,
+        max_pair: int = int(1e8),
+        *,
+        load_latents: bool = True,
+        load_terminal_latent_only: bool = False,
+    ):
+        self.env = lmdb.open(data_path, readonly=True,
+                             lock=False, readahead=False, meminit=False)
+
+        self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')
+        self.max_pair = max_pair
+        self.load_latents = bool(load_latents)
+        self.load_terminal_latent_only = bool(
+            load_terminal_latent_only)
+        if self.load_terminal_latent_only and not self.load_latents:
+            raise ValueError(
+                "load_terminal_latent_only requires load_latents=True")
+
+    def __len__(self):
+        return min(self.latents_shape[0], self.max_pair)
+
+    def __getitem__(self, idx):
+        """
+        Outputs:
+            - prompts: List of Strings
+            - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
+        """
+        prompts = retrieve_row_from_lmdb(
+            self.env,
+            "prompts", str, idx
+        )
+        batch = {
+            "prompts": prompts,
+            "idx": idx,
+        }
+        if self.load_latents:
+            latents = retrieve_row_from_lmdb(
+                self.env,
+                "latents", np.float16, idx, shape=self.latents_shape[1:]
+            )
+            if len(latents.shape) == 4:
+                if not self.load_terminal_latent_only:
+                    latents = latents[None, ...]
+            elif self.load_terminal_latent_only:
+                # Keep only the clean terminal trajectory row in host memory.
+                latents = latents[-1]
+            batch["ode_latent"] = torch.tensor(
+                latents, dtype=torch.float32)
+        return batch
+
+
+class ShardingLMDBDataset(Dataset):
+    def __init__(self, data_path: str, max_pair: int = int(1e8)):
+        self.envs = []
+        self.index = []
+
+        for fname in sorted(os.listdir(data_path)):
+            path = os.path.join(data_path, fname)
+            env = lmdb.open(path,
+                            readonly=True,
+                            lock=False,
+                            readahead=False,
+                            meminit=False)
+            self.envs.append(env)
+
+        self.latents_shape = [None] * len(self.envs)
+        for shard_id, env in enumerate(self.envs):
+            self.latents_shape[shard_id] = get_array_shape_from_lmdb(env, 'latents')
+            for local_i in range(self.latents_shape[shard_id][0]):
+                self.index.append((shard_id, local_i))
+
+            # print("shard_id ", shard_id, " local_i ", local_i)
+
+        self.max_pair = max_pair
+
+    def __len__(self):
+        return len(self.index)
+
+    def __getitem__(self, idx):
+        """
+            Outputs:
+                - prompts: List of Strings
+                - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
+        """
+        shard_id, local_idx = self.index[idx]
+
+        latents = retrieve_row_from_lmdb(
+            self.envs[shard_id],
+            "latents", np.float16, local_idx,
+            shape=self.latents_shape[shard_id][1:]
+        )
+
+        if len(latents.shape) == 4:
+            latents = latents[None, ...]
+
+        prompts = retrieve_row_from_lmdb(
+            self.envs[shard_id],
+            "prompts", str, local_idx
+        )
+
+        return {
+            "prompts": prompts,
+            "ode_latent": torch.tensor(latents, dtype=torch.float32)
+        }
+
+
+class TextImagePairDataset(Dataset):
+    def __init__(
+        self,
+        data_dir,
+        transform=None,
+        eval_first_n=-1,
+        pad_to_multiple_of=None
+    ):
+        """
+        Args:
+            data_dir (str): Path to the directory containing:
+                - target_crop_info_*.json (metadata file)
+                - */ (subdirectory containing images with matching aspect ratio)
+            transform (callable, optional): Optional transform to be applied on the image
+        """
+        self.transform = transform
+        data_dir = Path(data_dir)
+
+        # Find the metadata JSON file
+        metadata_files = list(data_dir.glob('target_crop_info_*.json'))
+        if not metadata_files:
+            raise FileNotFoundError(f"No metadata file found in {data_dir}")
+        if len(metadata_files) > 1:
+            raise ValueError(f"Multiple metadata files found in {data_dir}")
+
+        metadata_path = metadata_files[0]
+        # Extract aspect ratio from metadata filename (e.g. target_crop_info_26-15.json -> 26-15)
+        aspect_ratio = metadata_path.stem.split('_')[-1]
+
+        # Use aspect ratio subfolder for images
+        self.image_dir = data_dir / aspect_ratio
+        if not self.image_dir.exists():
+            raise FileNotFoundError(f"Image directory not found: {self.image_dir}")
+
+        # Load metadata
+        with open(metadata_path, 'r') as f:
+            self.metadata = json.load(f)
+
+        eval_first_n = eval_first_n if eval_first_n != -1 else len(self.metadata)
+        self.metadata = self.metadata[:eval_first_n]
+
+        # Verify all images exist
+        for item in self.metadata:
+            image_path = self.image_dir / item['file_name']
+            if not image_path.exists():
+                raise FileNotFoundError(f"Image not found: {image_path}")
+
+        self.dummy_prompt = "DUMMY PROMPT"
+        self.pre_pad_len = len(self.metadata)
+        if pad_to_multiple_of is not None and len(self.metadata) % pad_to_multiple_of != 0:
+            # Duplicate the last entry
+            self.metadata += [self.metadata[-1]] * (
+                pad_to_multiple_of - len(self.metadata) % pad_to_multiple_of
+            )
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        """
+        Returns:
+            dict: A dictionary containing:
+                - image: PIL Image
+                - caption: str
+                - target_bbox: list of int [x1, y1, x2, y2]
+                - target_ratio: str
+                - type: str
+                - origin_size: tuple of int (width, height)
+        """
+        item = self.metadata[idx]
+
+        # Load image
+        image_path = self.image_dir / item['file_name']
+        image = Image.open(image_path).convert('RGB')
+
+        # Apply transform if specified
+        if self.transform:
+            image = self.transform(image)
+
+        return {
+            'image': image,
+            'prompts': item['caption'],
+            'target_bbox': item['target_crop']['target_bbox'],
+            'target_ratio': item['target_crop']['target_ratio'],
+            'type': item['type'],
+            'origin_size': (item['origin_width'], item['origin_height']),
+            'idx': idx
+        }
+
+
+class _CyclingDataLoader:
+    """Infinite DataLoader iterator with epoch-aware DistributedSampler state."""
+
+    def __init__(self, dataloader):
+        self.dataloader = dataloader
+        self.epoch = 0
+        self._iterator = None
+
+    def __iter__(self):
+        return self
+
+    def _new_iterator(self):
+        if hasattr(self.dataloader.sampler, "set_epoch"):
+            self.dataloader.sampler.set_epoch(self.epoch)
+        self._iterator = iter(self.dataloader)
+
+    def __next__(self):
+        if self._iterator is None:
+            self._new_iterator()
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self.epoch += 1
+            self._new_iterator()
+            return next(self._iterator)
+
+    def seek(self, consumed_batches):
+        """Position the next item at an absolute batch counter."""
+        length = len(self.dataloader)
+        if length <= 0:
+            raise ValueError("cannot cycle an empty DataLoader")
+        self.epoch, offset = divmod(int(consumed_batches), length)
+        self._new_iterator()
+        for _ in range(offset):
+            next(self._iterator)
+        return offset
+
+
+def cycle(dl):
+    return _CyclingDataLoader(dl)
+
+
+# Compatibility re-exports keep legacy callers on ``utils.dataset`` while the
+# prompt/video implementation remains isolated from the LMDB dataset module.
+from utils.openvid_dataset import (  # noqa: E402
+    OpenVidPromptVideoDataset,
+    collate_openvid_prompt_video,
+)

1.3B/code/utils/dinov3_vjepa2_feature_bank.py

@@ -0,0 +1,301 @@
+"""Strict mmap-backed teacher bank for DINOv3 + V-JEPA2 features."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any, Mapping, Sequence
+
+import torch
+from safetensors import safe_open
+from torch.utils.data import Dataset
+
+
+BANK_KIND = "dual_dinov3_vjepa2_matched_teacher_bank"
+BANK_SCHEMA_VERSION = 1
+BRANCH_SHAPES = {
+    "dinov3": (1024, 1024),
+    "vjepa2": (1024, 1024),
+}
+
+
+def _canonical_bytes(value: Any) -> bytes:
+    return json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+        allow_nan=False,
+    ).encode("utf-8")
+
+
+def canonical_sha256(value: Any) -> str:
+    return hashlib.sha256(_canonical_bytes(value)).hexdigest()
+
+
+def file_sha256(path: str | os.PathLike) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while True:
+            block = handle.read(8 * 1024 * 1024)
+            if not block:
+                break
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _strict_sha256(value: object, label: str) -> str:
+    normalized = str(value)
+    if (
+        len(normalized) != 64
+        or any(character not in "0123456789abcdef" for character in normalized)
+    ):
+        raise ValueError(f"{label} must be a lowercase SHA-256")
+    return normalized
+
+
+class DINOv3VJEPA2FeatureBank(Dataset):
+    """Metadata dataset plus random-access frozen FP16 feature rows."""
+
+    def __init__(
+        self,
+        manifest_path: str | os.PathLike,
+        *,
+        expected_manifest_sha256: str,
+        expected_content_hash: str,
+        expected_source_manifest_sha256: str,
+        expected_source_csv_sha256: str,
+        expected_sample_identity_sha256: str,
+        expected_sample_count: int = 5000,
+        expected_vjepa_model_sha256: str | None = None,
+        expected_dinov3_checkpoint_sha256: str | None = None,
+        expected_dinov3_source_commit: str | None = None,
+        expected_shard_size: int = 32,
+        validate_all_shards: bool = True,
+    ) -> None:
+        self.manifest_path = Path(manifest_path).expanduser().resolve()
+        if not self.manifest_path.is_file():
+            raise FileNotFoundError(self.manifest_path)
+        actual_manifest_sha = file_sha256(self.manifest_path)
+        if actual_manifest_sha != _strict_sha256(
+            expected_manifest_sha256, "expected_manifest_sha256"
+        ):
+            raise ValueError("DINOv3+V-JEPA2 teacher bank manifest SHA-256 mismatch")
+        with self.manifest_path.open("r", encoding="utf-8") as handle:
+            manifest = json.load(handle)
+        if not isinstance(manifest, dict):
+            raise TypeError("DINOv3+V-JEPA2 teacher bank manifest must be an object")
+        required = {
+            "schema_version",
+            "kind",
+            "sample_count",
+            "source_manifest",
+            "branches",
+            "samples",
+            "shards",
+            "content_hash",
+        }
+        if set(manifest) != required:
+            raise ValueError(
+                "DINOv3+V-JEPA2 teacher bank manifest keys mismatch: "
+                f"{sorted(manifest)} != {sorted(required)}")
+        if (
+            manifest["schema_version"] != BANK_SCHEMA_VERSION
+            or manifest["kind"] != BANK_KIND
+            or int(manifest["sample_count"]) != int(expected_sample_count)
+        ):
+            raise ValueError("DINOv3+V-JEPA2 teacher bank kind/schema/sample count mismatch")
+        payload = dict(manifest)
+        content_hash = payload.pop("content_hash")
+        if canonical_sha256(payload) != content_hash:
+            raise ValueError("DINOv3+V-JEPA2 teacher bank content_hash is invalid")
+        if content_hash != _strict_sha256(
+            expected_content_hash, "expected_content_hash"
+        ):
+            raise ValueError("DINOv3+V-JEPA2 teacher bank content identity mismatch")
+        source = manifest["source_manifest"]
+        if source.get("sha256") != _strict_sha256(
+            expected_source_manifest_sha256,
+            "expected_source_manifest_sha256",
+        ):
+            raise ValueError("DINOv3+V-JEPA2 bank source manifest SHA-256 mismatch")
+        if source.get("source_csv_sha256") != _strict_sha256(
+            expected_source_csv_sha256,
+            "expected_source_csv_sha256",
+        ):
+            raise ValueError("DINOv3+V-JEPA2 bank source CSV SHA-256 mismatch")
+        if source.get("sample_identity_sha256") != _strict_sha256(
+            expected_sample_identity_sha256,
+            "expected_sample_identity_sha256",
+        ):
+            raise ValueError("DINOv3+V-JEPA2 bank sample identity mismatch")
+        branches = manifest["branches"]
+        if set(branches) != set(BRANCH_SHAPES):
+            raise ValueError("DINOv3+V-JEPA2 teacher bank branch set mismatch")
+        for name, expected_shape in BRANCH_SHAPES.items():
+            branch = branches[name]
+            if (
+                tuple(branch.get("sample_shape", ())) != expected_shape
+                or branch.get("dtype") != "float16"
+                or not isinstance(branch.get("feature_contract"), Mapping)
+                or not isinstance(branch.get("model_identity"), Mapping)
+            ):
+                raise ValueError(f"DINOv3+V-JEPA2 teacher branch {name!r} contract mismatch")
+        if expected_vjepa_model_sha256 is not None and (
+            branches["vjepa2"]["model_identity"].get("sha256")
+            != _strict_sha256(
+                expected_vjepa_model_sha256,
+                "expected_vjepa_model_sha256",
+            )
+        ):
+            raise ValueError("V-JEPA2 bank model SHA-256 mismatch")
+        if expected_dinov3_checkpoint_sha256 is not None and (
+            branches["dinov3"]["model_identity"].get("checkpoint_sha256")
+            != _strict_sha256(
+                expected_dinov3_checkpoint_sha256,
+                "expected_dinov3_checkpoint_sha256",
+            )
+        ):
+            raise ValueError("DINOv3 bank checkpoint SHA-256 mismatch")
+        if expected_dinov3_source_commit is not None and (
+            branches["dinov3"]["model_identity"].get("source_commit")
+            != str(expected_dinov3_source_commit)
+        ):
+            raise ValueError("DINOv3 bank source commit mismatch")
+        if type(expected_shard_size) is not int or expected_shard_size <= 0:
+            raise ValueError("expected_shard_size must be positive")
+
+        samples = manifest["samples"]
+        if len(samples) != int(expected_sample_count):
+            raise ValueError("DINOv3+V-JEPA2 bank sample table length mismatch")
+        self.samples = []
+        self._sample_index = {}
+        for expected_index, item in enumerate(samples):
+            if (
+                not isinstance(item, dict)
+                or set(item) != {"index", "sample_id", "prompt"}
+                or int(item["index"]) != expected_index
+                or not isinstance(item["sample_id"], str)
+                or not item["sample_id"]
+                or not isinstance(item["prompt"], str)
+            ):
+                raise ValueError(
+                    f"invalid DINOv3+V-JEPA2 bank sample record at {expected_index}")
+            if item["sample_id"] in self._sample_index:
+                raise ValueError("DINOv3+V-JEPA2 bank sample IDs must be unique")
+            self._sample_index[item["sample_id"]] = expected_index
+            self.samples.append(dict(item))
+
+        self.root = self.manifest_path.parent
+        self.shards = []
+        self._row_to_shard: list[tuple[int, int]] = [None] * len(samples)
+        expected_start = 0
+        for shard_index, item in enumerate(manifest["shards"]):
+            required_shard = {
+                "index", "path", "sample_start", "sample_end",
+                "sample_ids_sha256", "sha256", "content_hash",
+            }
+            if not isinstance(item, dict) or set(item) != required_shard:
+                raise ValueError(f"invalid shard record {shard_index}")
+            start = int(item["sample_start"])
+            end = int(item["sample_end"])
+            if (
+                int(item["index"]) != shard_index
+                or start != expected_start
+                or end <= start
+                or end > len(samples)
+                or end - start != min(expected_shard_size, len(samples) - start)
+            ):
+                raise ValueError(
+                    "DINOv3+V-JEPA2 bank shard ranges must be contiguous fixed-size "
+                    f"chunks of {expected_shard_size}")
+            relative = Path(str(item["path"]))
+            if relative.is_absolute() or ".." in relative.parts:
+                raise ValueError("DINOv3+V-JEPA2 bank shard path must stay under bank root")
+            path = (self.root / relative).resolve()
+            if self.root not in path.parents or not path.is_file():
+                raise FileNotFoundError(path)
+            ids = [samples[row]["sample_id"] for row in range(start, end)]
+            if canonical_sha256(ids) != item["sample_ids_sha256"]:
+                raise ValueError("DINOv3+V-JEPA2 bank shard sample ID hash mismatch")
+            if validate_all_shards and file_sha256(path) != item["sha256"]:
+                raise ValueError(f"DINOv3+V-JEPA2 bank shard SHA-256 mismatch: {path}")
+            with safe_open(str(path), framework="pt", device="cpu") as handle:
+                metadata = handle.metadata() or {}
+                if metadata.get("content_hash") != item["content_hash"]:
+                    raise ValueError("DINOv3+V-JEPA2 bank shard embedded content hash mismatch")
+                for name, shape in BRANCH_SHAPES.items():
+                    tensor_shape = tuple(handle.get_slice(name).get_shape())
+                    if tensor_shape != (end - start, *shape):
+                        raise ValueError(
+                            f"DINOv3+V-JEPA2 bank shard tensor {name!r} shape mismatch")
+            for row in range(start, end):
+                self._row_to_shard[row] = (shard_index, row - start)
+            self.shards.append({**item, "resolved_path": str(path)})
+            expected_start = end
+        if expected_start != len(samples):
+            raise ValueError("DINOv3+V-JEPA2 bank shards do not cover all samples")
+        self.manifest = manifest
+        self.content_hash = content_hash
+        self.manifest_sha256 = actual_manifest_sha
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> dict[str, object]:
+        item = self.samples[int(index)]
+        return {
+            "prompts": item["prompt"],
+            "sample_ids": item["sample_id"],
+            "teacher_bank_rows": int(item["index"]),
+        }
+
+    @staticmethod
+    def collate_fn(items: Sequence[Mapping[str, object]]) -> dict[str, object]:
+        return {
+            "prompts": [str(item["prompts"]) for item in items],
+            "sample_ids": [str(item["sample_ids"]) for item in items],
+            "teacher_bank_rows": torch.tensor([
+                int(item["teacher_bank_rows"]) for item in items
+            ], dtype=torch.long),
+        }
+
+    def load_rows(
+        self,
+        sample_ids: Sequence[str],
+    ) -> dict[str, torch.Tensor]:
+        """Read only requested rows, preserving order and duplicate IDs."""
+        ids = tuple(str(value) for value in sample_ids)
+        if not ids or any(value not in self._sample_index for value in ids):
+            raise KeyError("DINOv3+V-JEPA2 bank request contains an unknown sample ID")
+        output = {name: [None] * len(ids) for name in BRANCH_SHAPES}
+        by_shard: dict[int, list[tuple[int, int]]] = {}
+        for output_index, sample_id in enumerate(ids):
+            row = self._sample_index[sample_id]
+            shard_index, offset = self._row_to_shard[row]
+            by_shard.setdefault(shard_index, []).append((output_index, offset))
+        for shard_index, requests in by_shard.items():
+            path = self.shards[shard_index]["resolved_path"]
+            with safe_open(path, framework="pt", device="cpu") as handle:
+                slices = {name: handle.get_slice(name) for name in BRANCH_SHAPES}
+                for output_index, offset in requests:
+                    for name in BRANCH_SHAPES:
+                        output[name][output_index] = slices[name][
+                            offset:offset + 1
+                        ].squeeze(0)
+        return {
+            name: torch.stack(rows, dim=0)
+            for name, rows in output.items()
+        }
+
+
+__all__ = [
+    "BANK_KIND",
+    "BANK_SCHEMA_VERSION",
+    "BRANCH_SHAPES",
+    "DINOv3VJEPA2FeatureBank",
+    "canonical_sha256",
+    "file_sha256",
+]

1.3B/code/utils/distributed.py

@@ -0,0 +1,188 @@
+from datetime import timedelta
+from functools import partial
+import inspect
+import os
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import FullStateDictConfig, FullOptimStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
+from torch.distributed.fsdp.api import CPUOffload
+from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+
+
+def fsdp_state_dict(model):
+    fsdp_fullstate_save_policy = FullStateDictConfig(
+        offload_to_cpu=True, rank0_only=True
+    )
+    with FSDP.state_dict_type(
+        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
+    ):
+        checkpoint = model.state_dict()
+
+    return checkpoint
+
+
+def fsdp_optim_state_dict(model, optimizer):
+    """Gather the full (unsharded) optimizer state on rank 0 (CPU).
+
+    Collective: must be called by every rank.  Returns the full optim state
+    dict on rank 0 and an empty/partial dict on other ranks (rank0_only=True).
+    """
+    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    optim_policy = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(
+        model,
+        StateDictType.FULL_STATE_DICT,
+        state_dict_config=save_policy,
+        optim_state_dict_config=optim_policy,
+    ):
+        osd = FSDP.optim_state_dict(model, optimizer)
+    return osd
+
+
+def fsdp_load_optim_state_dict(model, optimizer, full_osd):
+    """Load a full optimizer state dict (present on rank 0) back into ``optimizer``.
+
+    Collective: every rank must call this.  ``full_osd`` need only be non-None on
+    rank 0 (FSDP broadcasts/shards it to the other ranks), which keeps host memory
+    bounded to a single copy of the optimizer state.
+    """
+    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    optim_policy = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(
+        model,
+        StateDictType.FULL_STATE_DICT,
+        state_dict_config=save_policy,
+        optim_state_dict_config=optim_policy,
+    ):
+        sharded = FSDP.optim_state_dict_to_load(
+            model=model, optim=optimizer, optim_state_dict=full_osd
+        )
+    optimizer.load_state_dict(sharded)
+
+
+def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
+    if mixed_precision:
+        mixed_precision_policy = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.float32,
+            cast_forward_inputs=False
+        )
+    else:
+        mixed_precision_policy = None
+
+    if wrap_strategy == "transformer":
+        if transformer_module is None:
+            raise ValueError("transformer wrap strategy requires transformer_module")
+        transformer_classes = (
+            {transformer_module}
+            if isinstance(transformer_module, type)
+            else set(transformer_module)
+        )
+        auto_wrap_policy = partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls=transformer_classes
+        )
+    elif wrap_strategy == "size":
+        auto_wrap_policy = partial(
+            size_based_auto_wrap_policy,
+            min_num_params=min_num_params
+        )
+    else:
+        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")
+
+    os.environ["NCCL_CROSS_NIC"] = "1"
+
+    sharding_strategy = {
+        "full": ShardingStrategy.FULL_SHARD,
+        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
+        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
+        "no_shard": ShardingStrategy.NO_SHARD,
+    }[sharding_strategy]
+
+    module = FSDP(
+        module,
+        auto_wrap_policy=auto_wrap_policy,
+        sharding_strategy=sharding_strategy,
+        mixed_precision=mixed_precision_policy,
+        device_id=torch.cuda.current_device(),
+        limit_all_gathers=True,
+        use_orig_params=True,
+        cpu_offload=CPUOffload(offload_params=cpu_offload),
+        # Every rank constructs and loads the same pretrained/checkpoint state
+        # before wrapping; no rank-0 broadcast is implied here.
+        sync_module_states=False
+    )
+    return module
+
+
+def barrier():
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def launch_distributed_job(backend: str = "nccl"):
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    host = os.environ["MASTER_ADDR"]
+    port = int(os.environ["MASTER_PORT"])
+    # Bind CUDA before NCCL creates communicators.
+    torch.cuda.set_device(local_rank)
+
+    if ":" in host:  # IPv6
+        init_method = f"tcp://[{host}]:{port}"
+    else:  # IPv4
+        init_method = f"tcp://{host}:{port}"
+    init_kwargs = {}
+    if (
+        backend == "nccl"
+        and "device_id" in inspect.signature(
+            dist.init_process_group).parameters
+    ):
+        init_kwargs["device_id"] = torch.device("cuda", local_rank)
+    dist.init_process_group(
+        rank=rank,
+        world_size=world_size,
+        backend=backend,
+        init_method=init_method,
+        timeout=timedelta(minutes=30),
+        **init_kwargs,
+    )
+
+
+class EMA_FSDP:
+    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
+        self.decay = decay
+        self.shadow = {}
+        self._init_shadow(fsdp_module)
+
+    @torch.no_grad()
+    def _init_shadow(self, fsdp_module):
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                self.shadow[n] = p.detach().clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, fsdp_module):
+        d = self.decay
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                self.shadow[n].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
+
+    # Optional helpers ---------------------------------------------------
+    def state_dict(self):
+        return self.shadow            # picklable
+
+    def load_state_dict(self, sd):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    def copy_to(self, fsdp_module):
+        # load EMA weights into an (unwrapped) copy of the generator
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=True):
+            for n, p in fsdp_module.module.named_parameters():
+                if n in self.shadow:
+                    p.data.copy_(self.shadow[n].to(p.dtype, device=p.device))

1.3B/code/utils/dual_video_feature_bank.py

@@ -0,0 +1,286 @@
+"""Strict mmap-backed teacher bank for V-JEPA2 + VideoMAE-v2 features."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any, Mapping, Sequence
+
+import torch
+from safetensors import safe_open
+from torch.utils.data import Dataset
+
+
+BANK_KIND = "dual_vjepa2_videomae_v2_matched_teacher_bank"
+BANK_SCHEMA_VERSION = 1
+BRANCH_SHAPES = {
+    "vjepa2": (1024, 1024),
+    "videomae_v2": (1280, 1024),
+}
+
+
+def _canonical_bytes(value: Any) -> bytes:
+    return json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+        allow_nan=False,
+    ).encode("utf-8")
+
+
+def canonical_sha256(value: Any) -> str:
+    return hashlib.sha256(_canonical_bytes(value)).hexdigest()
+
+
+def file_sha256(path: str | os.PathLike) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while True:
+            block = handle.read(8 * 1024 * 1024)
+            if not block:
+                break
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _strict_sha256(value: object, label: str) -> str:
+    normalized = str(value)
+    if (
+        len(normalized) != 64
+        or any(character not in "0123456789abcdef" for character in normalized)
+    ):
+        raise ValueError(f"{label} must be a lowercase SHA-256")
+    return normalized
+
+
+class DualVideoFeatureBank(Dataset):
+    """Metadata dataset plus random-access frozen FP16 feature rows."""
+
+    def __init__(
+        self,
+        manifest_path: str | os.PathLike,
+        *,
+        expected_manifest_sha256: str,
+        expected_content_hash: str,
+        expected_source_manifest_sha256: str,
+        expected_sample_count: int = 5000,
+        expected_videomae_revision: str | None = None,
+        expected_videomae_download_manifest_content_hash: str | None = None,
+        expected_vjepa_model_sha256: str | None = None,
+        validate_all_shards: bool = True,
+    ) -> None:
+        self.manifest_path = Path(manifest_path).expanduser().resolve()
+        if not self.manifest_path.is_file():
+            raise FileNotFoundError(self.manifest_path)
+        actual_manifest_sha = file_sha256(self.manifest_path)
+        if actual_manifest_sha != _strict_sha256(
+            expected_manifest_sha256, "expected_manifest_sha256"
+        ):
+            raise ValueError("dual teacher bank manifest SHA-256 mismatch")
+        with self.manifest_path.open("r", encoding="utf-8") as handle:
+            manifest = json.load(handle)
+        if not isinstance(manifest, dict):
+            raise TypeError("dual teacher bank manifest must be an object")
+        required = {
+            "schema_version",
+            "kind",
+            "sample_count",
+            "source_manifest",
+            "branches",
+            "samples",
+            "shards",
+            "content_hash",
+        }
+        if set(manifest) != required:
+            raise ValueError(
+                "dual teacher bank manifest keys mismatch: "
+                f"{sorted(manifest)} != {sorted(required)}")
+        if (
+            manifest["schema_version"] != BANK_SCHEMA_VERSION
+            or manifest["kind"] != BANK_KIND
+            or int(manifest["sample_count"]) != int(expected_sample_count)
+        ):
+            raise ValueError("dual teacher bank kind/schema/sample count mismatch")
+        payload = dict(manifest)
+        content_hash = payload.pop("content_hash")
+        if canonical_sha256(payload) != content_hash:
+            raise ValueError("dual teacher bank content_hash is invalid")
+        if content_hash != _strict_sha256(
+            expected_content_hash, "expected_content_hash"
+        ):
+            raise ValueError("dual teacher bank content identity mismatch")
+        source = manifest["source_manifest"]
+        if source.get("sha256") != _strict_sha256(
+            expected_source_manifest_sha256,
+            "expected_source_manifest_sha256",
+        ):
+            raise ValueError("dual bank source manifest SHA-256 mismatch")
+        branches = manifest["branches"]
+        if set(branches) != set(BRANCH_SHAPES):
+            raise ValueError("dual teacher bank branch set mismatch")
+        for name, expected_shape in BRANCH_SHAPES.items():
+            branch = branches[name]
+            if (
+                tuple(branch.get("sample_shape", ())) != expected_shape
+                or branch.get("dtype") != "float16"
+                or not isinstance(branch.get("feature_contract"), Mapping)
+                or not isinstance(branch.get("model_identity"), Mapping)
+            ):
+                raise ValueError(f"dual teacher branch {name!r} contract mismatch")
+        if expected_videomae_revision is not None and (
+            branches["videomae_v2"]["model_identity"].get("revision")
+            != str(expected_videomae_revision)
+        ):
+            raise ValueError("VideoMAE-v2 bank model revision mismatch")
+        if expected_videomae_download_manifest_content_hash is not None and (
+            branches["videomae_v2"]["model_identity"].get(
+                "download_manifest_content_hash"
+            )
+            != _strict_sha256(
+                expected_videomae_download_manifest_content_hash,
+                "expected_videomae_download_manifest_content_hash",
+            )
+        ):
+            raise ValueError(
+                "VideoMAE-v2 bank download manifest identity mismatch")
+        if expected_vjepa_model_sha256 is not None and (
+            branches["vjepa2"]["model_identity"].get("sha256")
+            != _strict_sha256(
+                expected_vjepa_model_sha256,
+                "expected_vjepa_model_sha256",
+            )
+        ):
+            raise ValueError("V-JEPA2 bank model SHA-256 mismatch")
+
+        samples = manifest["samples"]
+        if len(samples) != int(expected_sample_count):
+            raise ValueError("dual bank sample table length mismatch")
+        self.samples = []
+        self._sample_index = {}
+        for expected_index, item in enumerate(samples):
+            if (
+                not isinstance(item, dict)
+                or set(item) != {"index", "sample_id", "prompt"}
+                or int(item["index"]) != expected_index
+                or not isinstance(item["sample_id"], str)
+                or not item["sample_id"]
+                or not isinstance(item["prompt"], str)
+            ):
+                raise ValueError(
+                    f"invalid dual bank sample record at {expected_index}")
+            if item["sample_id"] in self._sample_index:
+                raise ValueError("dual bank sample IDs must be unique")
+            self._sample_index[item["sample_id"]] = expected_index
+            self.samples.append(dict(item))
+
+        self.root = self.manifest_path.parent
+        self.shards = []
+        self._row_to_shard: list[tuple[int, int]] = [None] * len(samples)
+        expected_start = 0
+        for shard_index, item in enumerate(manifest["shards"]):
+            required_shard = {
+                "index", "path", "sample_start", "sample_end",
+                "sample_ids_sha256", "sha256", "content_hash",
+            }
+            if not isinstance(item, dict) or set(item) != required_shard:
+                raise ValueError(f"invalid shard record {shard_index}")
+            start = int(item["sample_start"])
+            end = int(item["sample_end"])
+            if (
+                int(item["index"]) != shard_index
+                or start != expected_start
+                or end <= start
+                or end > len(samples)
+            ):
+                raise ValueError("dual bank shard ranges are not contiguous")
+            relative = Path(str(item["path"]))
+            if relative.is_absolute() or ".." in relative.parts:
+                raise ValueError("dual bank shard path must stay under bank root")
+            path = (self.root / relative).resolve()
+            if self.root not in path.parents or not path.is_file():
+                raise FileNotFoundError(path)
+            ids = [samples[row]["sample_id"] for row in range(start, end)]
+            if canonical_sha256(ids) != item["sample_ids_sha256"]:
+                raise ValueError("dual bank shard sample ID hash mismatch")
+            if validate_all_shards and file_sha256(path) != item["sha256"]:
+                raise ValueError(f"dual bank shard SHA-256 mismatch: {path}")
+            with safe_open(str(path), framework="pt", device="cpu") as handle:
+                metadata = handle.metadata() or {}
+                if metadata.get("content_hash") != item["content_hash"]:
+                    raise ValueError("dual bank shard embedded content hash mismatch")
+                for name, shape in BRANCH_SHAPES.items():
+                    tensor_shape = tuple(handle.get_slice(name).get_shape())
+                    if tensor_shape != (end - start, *shape):
+                        raise ValueError(
+                            f"dual bank shard tensor {name!r} shape mismatch")
+            for row in range(start, end):
+                self._row_to_shard[row] = (shard_index, row - start)
+            self.shards.append({**item, "resolved_path": str(path)})
+            expected_start = end
+        if expected_start != len(samples):
+            raise ValueError("dual bank shards do not cover all samples")
+        self.manifest = manifest
+        self.content_hash = content_hash
+        self.manifest_sha256 = actual_manifest_sha
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> dict[str, object]:
+        item = self.samples[int(index)]
+        return {
+            "prompts": item["prompt"],
+            "sample_ids": item["sample_id"],
+            "teacher_bank_rows": int(item["index"]),
+        }
+
+    @staticmethod
+    def collate_fn(items: Sequence[Mapping[str, object]]) -> dict[str, object]:
+        return {
+            "prompts": [str(item["prompts"]) for item in items],
+            "sample_ids": [str(item["sample_ids"]) for item in items],
+            "teacher_bank_rows": torch.tensor([
+                int(item["teacher_bank_rows"]) for item in items
+            ], dtype=torch.long),
+        }
+
+    def load_rows(
+        self,
+        sample_ids: Sequence[str],
+    ) -> dict[str, torch.Tensor]:
+        """Read only requested rows, preserving order and duplicate IDs."""
+        ids = tuple(str(value) for value in sample_ids)
+        if not ids or any(value not in self._sample_index for value in ids):
+            raise KeyError("dual bank request contains an unknown sample ID")
+        output = {name: [None] * len(ids) for name in BRANCH_SHAPES}
+        by_shard: dict[int, list[tuple[int, int]]] = {}
+        for output_index, sample_id in enumerate(ids):
+            row = self._sample_index[sample_id]
+            shard_index, offset = self._row_to_shard[row]
+            by_shard.setdefault(shard_index, []).append((output_index, offset))
+        for shard_index, requests in by_shard.items():
+            path = self.shards[shard_index]["resolved_path"]
+            with safe_open(path, framework="pt", device="cpu") as handle:
+                slices = {name: handle.get_slice(name) for name in BRANCH_SHAPES}
+                for output_index, offset in requests:
+                    for name in BRANCH_SHAPES:
+                        output[name][output_index] = slices[name][
+                            offset:offset + 1
+                        ].squeeze(0)
+        return {
+            name: torch.stack(rows, dim=0)
+            for name, rows in output.items()
+        }
+
+
+__all__ = [
+    "BANK_KIND",
+    "BANK_SCHEMA_VERSION",
+    "BRANCH_SHAPES",
+    "DualVideoFeatureBank",
+    "canonical_sha256",
+    "file_sha256",
+]

1.3B/code/utils/lmdb.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+
+def get_array_shape_from_lmdb(env, array_name):
+    with env.begin() as txn:
+        image_shape = txn.get(f"{array_name}_shape".encode()).decode()
+        image_shape = tuple(map(int, image_shape.split()))
+    return image_shape
+
+
+def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
+    """
+    Store rows of multiple numpy arrays in a single LMDB.
+    Each row is stored separately with a naming convention.
+    """
+    with env.begin(write=True) as txn:
+        for array_name, array in arrays_dict.items():
+            for i, row in enumerate(array):
+                # Convert row to bytes
+                if isinstance(row, str):
+                    row_bytes = row.encode()
+                else:
+                    row_bytes = row.tobytes()
+
+                data_key = f'{array_name}_{start_index + i}_data'.encode()
+
+                txn.put(data_key, row_bytes)
+
+
+def process_data_dict(data_dict, seen_prompts):
+    output_dict = {}
+
+    all_videos = []
+    all_prompts = []
+    for prompt, video in data_dict.items():
+        if prompt in seen_prompts:
+            continue
+        else:
+            seen_prompts.add(prompt)
+
+        video = video.half().numpy()
+        all_videos.append(video)
+        all_prompts.append(prompt)
+
+    if len(all_videos) == 0:
+        return {"latents": np.array([]), "prompts": np.array([])}
+
+    all_videos = np.concatenate(all_videos, axis=0)
+
+    output_dict['latents'] = all_videos
+    output_dict['prompts'] = np.array(all_prompts)
+
+    return output_dict
+
+
+def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):
+    """
+    Retrieve a specific row from a specific array in the LMDB.
+    """
+    data_key = f'{array_name}_{row_index}_data'.encode()
+
+    with lmdb_env.begin() as txn:
+        row_bytes = txn.get(data_key)
+
+    if dtype == str:
+        array = row_bytes.decode()
+    else:
+        array = np.frombuffer(row_bytes, dtype=dtype)
+
+    if shape is not None and len(shape) > 0:
+        array = array.reshape(shape)
+    return array

1.3B/code/utils/loss.py

@@ -0,0 +1,81 @@
+from abc import ABC, abstractmethod
+import torch
+
+
+class DenoisingLoss(ABC):
+    @abstractmethod
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Base class for denoising loss.
+        Input:
+            - x: the clean data with shape [B, F, C, H, W]
+            - x_pred: the predicted clean data with shape [B, F, C, H, W]
+            - noise: the noise with shape [B, F, C, H, W]
+            - noise_pred: the predicted noise with shape [B, F, C, H, W]
+            - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]
+            - timestep: the current timestep with shape [B, F]
+        """
+        pass
+
+
+class X0PredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((x - x_pred) ** 2)
+
+
+class VPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))
+        return torch.mean(weights * (x - x_pred) ** 2)
+
+
+class NoisePredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((noise - noise_pred) ** 2)
+
+
+class FlowPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((kwargs["flow_pred"] - (noise - x)) ** 2)
+
+
+NAME_TO_CLASS = {
+    "x0": X0PredLoss,
+    "v": VPredLoss,
+    "noise": NoisePredLoss,
+    "flow": FlowPredLoss
+}
+
+
+def get_denoising_loss(loss_type: str) -> DenoisingLoss:
+    return NAME_TO_CLASS[loss_type]

1.3B/code/utils/matched_video_dataset.py

@@ -0,0 +1,381 @@
+"""Metadata-only matched-video dataset with explicitly lazy teacher I/O."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+from pathlib import Path
+import time
+from typing import Any, Callable, Mapping, Sequence
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from utils.matched_video_feature_cache import (
+    MATCHED_PYAV_LAST64_DECODER_CONTRACT,
+    MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT,
+    MatchedVideoFeatureCache,
+)
+from utils.matched_video_manifest import (
+    MatchedVideoManifestEntry,
+    load_matched_video_manifest,
+)
+
+
+@dataclass(frozen=True)
+class DecodedMatchedVideo:
+    """The final contiguous RGB frame window in V-JEPA2 input range."""
+
+    video01: torch.Tensor
+    frame_indices: torch.Tensor
+    source_frame_count: int
+
+
+def _frames_to_video01(frames: Sequence[Any]) -> torch.Tensor:
+    tensors: list[torch.Tensor] = []
+    shape: tuple[int, int, int] | None = None
+    for frame in frames:
+        array = frame.to_ndarray(format="rgb24")
+        if not isinstance(array, np.ndarray):
+            raise TypeError("PyAV frame.to_ndarray did not return ndarray")
+        if array.ndim != 3 or int(array.shape[2]) != 3:
+            raise ValueError(
+                f"decoded RGB frame must be [H,W,3], got {array.shape}"
+            )
+        if array.dtype != np.uint8:
+            raise ValueError(f"decoded RGB frame must be uint8, got {array.dtype}")
+        current_shape = tuple(int(value) for value in array.shape)
+        if shape is None:
+            shape = current_shape
+        elif current_shape != shape:
+            raise ValueError(
+                "video changes resolution inside selected frame window: "
+                f"{current_shape} != {shape}"
+            )
+        tensors.append(
+            torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1)
+        )
+    if not tensors:
+        raise ValueError("decoded frame window is empty")
+    return torch.stack(tensors).to(dtype=torch.float32).div_(255.0)
+
+
+def decode_matched_video_last_contiguous_pyav(
+    video_path: str | os.PathLike,
+    *,
+    decoder_contract: Mapping[str, Any] | None = None,
+) -> DecodedMatchedVideo:
+    """Decode exactly frames 17..80 from a contractually 81-frame MP4."""
+
+    contract = dict(
+        decoder_contract
+        if decoder_contract is not None
+        else MATCHED_PYAV_LAST64_DECODER_CONTRACT
+    )
+    if contract.get("backend") != "pyav":
+        raise ValueError("decoder contract backend must be pyav")
+    if contract.get("selection") != "last_contiguous":
+        raise ValueError("decoder contract selection must be last_contiguous")
+    num_frames = int(contract.get("num_frames", 0))
+    expected_count = int(contract.get("expected_source_frame_count", 0))
+    if num_frames <= 0 or expected_count < num_frames:
+        raise ValueError("decoder contract frame geometry is invalid")
+    expected_bounds = [expected_count - num_frames, expected_count - 1]
+    if list(contract.get("expected_frame_indices", ())) != expected_bounds:
+        raise ValueError(
+            "decoder contract expected_frame_indices do not describe last64"
+        )
+    try:
+        import av  # Imported only for cache misses.
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError(
+            "PyAV is required for a matched-video teacher cache miss; "
+            "install the pinned av==13.1.0 dependency"
+        ) from exc
+    expected_version = str(contract.get("backend_version", ""))
+    actual_version = str(getattr(av, "__version__", ""))
+    if actual_version != expected_version:
+        raise RuntimeError(
+            "PyAV version differs from the cache decoder contract: "
+            f"{actual_version!r} != {expected_version!r}"
+        )
+
+    path = Path(video_path).expanduser().resolve()
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    selected: list[Any] = []
+    decoded_count = 0
+    start = expected_count - num_frames
+    with av.open(str(path), mode="r") as container:
+        streams = container.streams.video
+        if len(streams) != 1:
+            raise ValueError(
+                f"expected exactly one video stream in {path}, got {len(streams)}"
+            )
+        stream = streams[0]
+        for frame_index, frame in enumerate(container.decode(stream)):
+            decoded_count = frame_index + 1
+            if start <= frame_index < expected_count:
+                selected.append(frame)
+    if decoded_count != expected_count:
+        raise ValueError(
+            "matched video source frame count violates the last64 contract: "
+            f"sample={path.name}, actual={decoded_count}, "
+            f"expected={expected_count}"
+        )
+    if len(selected) != num_frames:
+        raise RuntimeError(
+            f"failed to decode frames {start}..{expected_count - 1} from {path}"
+        )
+    video01 = _frames_to_video01(selected)
+    if tuple(video01.shape[:2]) != (num_frames, 3):
+        raise RuntimeError(f"decoded video has invalid shape {tuple(video01.shape)}")
+    return DecodedMatchedVideo(
+        video01=video01,
+        frame_indices=torch.arange(start, expected_count, dtype=torch.long),
+        source_frame_count=decoded_count,
+    )
+
+
+DecodeFunction = Callable[[Path, Mapping[str, Any]], DecodedMatchedVideo]
+
+
+def _default_decode(
+    path: Path,
+    contract: Mapping[str, Any],
+) -> DecodedMatchedVideo:
+    return decode_matched_video_last_contiguous_pyav(
+        path, decoder_contract=contract
+    )
+
+
+class MatchedManifestVideoDataset(Dataset):
+    """Return metadata eagerly; teacher cache/decode only on explicit request."""
+
+    def __init__(
+        self,
+        manifest_path: str | os.PathLike,
+        video_root: str | os.PathLike,
+        cache_root: str | os.PathLike,
+        expected_model_sha256: str,
+        expected_feature_contract_sha256: str,
+        *,
+        expected_source_csv_sha256: str | None = None,
+        expected_count: int | None = None,
+        feature_contract: Mapping[str, Any] | None = None,
+        decoder_contract: Mapping[str, Any] | None = None,
+        decode_fn: DecodeFunction | None = None,
+    ) -> None:
+        self.manifest_path = Path(manifest_path).expanduser().resolve()
+        self.video_root = Path(video_root).expanduser().resolve()
+        if not self.video_root.is_dir():
+            raise NotADirectoryError(self.video_root)
+        self.entries, self.manifest_metadata = load_matched_video_manifest(
+            self.manifest_path,
+            expected_source_csv_sha256=expected_source_csv_sha256,
+            expected_count=expected_count,
+        )
+        if not self.entries:
+            raise ValueError("matched-video manifest is empty")
+        self.feature_cache = MatchedVideoFeatureCache(
+            cache_root,
+            expected_model_sha256=expected_model_sha256,
+            expected_feature_contract_sha256=expected_feature_contract_sha256,
+            feature_contract=(
+                feature_contract
+                if feature_contract is not None
+                else MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT
+            ),
+            decoder_contract=(
+                decoder_contract
+                if decoder_contract is not None
+                else MATCHED_PYAV_LAST64_DECODER_CONTRACT
+            ),
+        )
+        self.decoder_contract = self.feature_cache.decoder_contract
+        self.decode_fn = decode_fn if decode_fn is not None else _default_decode
+        self._cache_key_to_index: dict[str, int] = {}
+        self._sample_id_to_index: dict[str, int] = {}
+        for index, entry in enumerate(self.entries):
+            cache_key = self.feature_cache.key_for(entry)
+            if cache_key in self._cache_key_to_index:
+                raise ValueError(f"duplicate matched-video cache key: {cache_key}")
+            if entry.sample_id in self._sample_id_to_index:
+                raise ValueError(f"duplicate sample_id: {entry.sample_id}")
+            self._cache_key_to_index[cache_key] = index
+            self._sample_id_to_index[entry.sample_id] = index
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def _entry_for_sample(self, sample_id: str) -> MatchedVideoManifestEntry:
+        try:
+            return self.entries[self._sample_id_to_index[str(sample_id)]]
+        except KeyError as exc:
+            raise KeyError(f"unknown matched-video sample_id: {sample_id!r}") from exc
+
+    def _video_path_and_verify_stat(
+        self,
+        entry: MatchedVideoManifestEntry,
+    ) -> Path:
+        path = self.video_root / entry.video_name
+        if not path.is_file():
+            raise FileNotFoundError(path)
+        stat = path.stat()
+        actual = (int(stat.st_size), int(stat.st_mtime_ns))
+        expected = (entry.video_size_bytes, entry.video_mtime_ns)
+        if actual != expected:
+            raise ValueError(
+                "matched video is stale after manifest freeze for "
+                f"sample_id={entry.sample_id}: actual={actual}, "
+                f"expected={expected}"
+            )
+        return path
+
+    def cache_key_for_sample(self, sample_id: str) -> str:
+        return self.feature_cache.key_for(self._entry_for_sample(sample_id))
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Return only prompt/sample/path metadata; never touch teacher data."""
+
+        entry = self.entries[int(index)]
+        video_path = self._video_path_and_verify_stat(entry)
+        return {
+            "idx": entry.manifest_index,
+            "prompts": entry.caption,
+            "sample_id": entry.sample_id,
+            "teacher_cache_key": self.feature_cache.key_for(entry),
+            "video_name": entry.video_name,
+            "video_path": str(video_path),
+        }
+
+    def load_teacher_payload(self, sample_id: str) -> dict[str, Any]:
+        """Load/decode teacher input only after suffix replay selects a sample."""
+
+        entry = self._entry_for_sample(sample_id)
+        video_path = self._video_path_and_verify_stat(entry)
+        cache_key = self.feature_cache.key_for(entry)
+        cached = self.feature_cache.load(entry, cache_key=cache_key)
+        payload: dict[str, Any] = {
+            "sample_id": entry.sample_id,
+            "teacher_cache_key": cache_key,
+            "video_name": entry.video_name,
+            "video_path": str(video_path),
+        }
+        if cached is not None:
+            payload.update({
+                "current_teacher_cache_hit": True,
+                "current_teacher_decode_seconds": 0.0,
+                "current_teacher_features": cached,
+            })
+            return payload
+
+        decode_started = time.perf_counter()
+        decoded = self.decode_fn(video_path, self.decoder_contract)
+        decode_seconds = time.perf_counter() - decode_started
+        if not isinstance(decoded, DecodedMatchedVideo):
+            raise TypeError("decode_fn must return DecodedMatchedVideo")
+        expected_count = int(self.decoder_contract["expected_source_frame_count"])
+        expected_frames = int(self.decoder_contract["num_frames"])
+        expected_indices = torch.arange(
+            expected_count - expected_frames,
+            expected_count,
+            dtype=torch.long,
+        )
+        if (
+            int(decoded.source_frame_count) != expected_count
+            or not isinstance(decoded.frame_indices, torch.Tensor)
+            or decoded.frame_indices.dtype != torch.long
+            or not torch.equal(
+                decoded.frame_indices.detach().cpu(), expected_indices
+            )
+        ):
+            raise ValueError(
+                "decode_fn violated the strict last-contiguous 17..80 contract"
+            )
+        video = decoded.video01
+        if (
+            not isinstance(video, torch.Tensor)
+            or video.dtype != torch.float32
+            or video.ndim != 4
+            or tuple(video.shape[:2]) != (expected_frames, 3)
+            or not bool(torch.isfinite(video).all().item())
+            or float(video.min()) < 0.0
+            or float(video.max()) > 1.0
+        ):
+            raise ValueError(
+                "decode_fn returned invalid video01; expected finite FP32 "
+                f"[{expected_frames},3,H,W] in [0,1]"
+            )
+        payload.update({
+            "current_teacher_cache_hit": False,
+            "current_teacher_decode_seconds": decode_seconds,
+            "current_teacher_frame_indices": decoded.frame_indices,
+            "current_teacher_source_frame_count": expected_count,
+            "current_teacher_video01": video.contiguous(),
+        })
+        return payload
+
+    def write_feature_cache(
+        self,
+        cache_key: str,
+        feature: torch.Tensor,
+        *,
+        sample_id: str,
+    ) -> Path:
+        try:
+            index = self._cache_key_to_index[str(cache_key)]
+        except KeyError as exc:
+            raise KeyError(f"unknown teacher cache key: {cache_key!r}") from exc
+        entry = self.entries[index]
+        if entry.sample_id != str(sample_id):
+            raise ValueError(
+                "sample_id does not own teacher cache key: "
+                f"{sample_id!r} != {entry.sample_id!r}"
+            )
+        self._video_path_and_verify_stat(entry)
+        return self.feature_cache.write(
+            entry, feature, cache_key=str(cache_key)
+        )
+
+    @staticmethod
+    def collate_fn(batch: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
+        return collate_matched_video_metadata(batch)
+
+
+def collate_matched_video_metadata(
+    batch: Sequence[Mapping[str, Any]],
+) -> dict[str, Any]:
+    """Collate only metadata; teacher payloads intentionally have no collator."""
+
+    if not batch:
+        raise ValueError("cannot collate an empty matched-video batch")
+    forbidden = {
+        "current_teacher_features",
+        "current_teacher_video01",
+        "current_teacher_cache_hit",
+    }
+    for item in batch:
+        present = forbidden.intersection(item)
+        if present:
+            raise ValueError(
+                "metadata collator received eager teacher payload fields: "
+                f"{sorted(present)}"
+            )
+    return {
+        "idx": torch.tensor([int(item["idx"]) for item in batch]),
+        "prompts": [str(item["prompts"]) for item in batch],
+        "sample_ids": [str(item["sample_id"]) for item in batch],
+        "teacher_cache_keys": [str(item["teacher_cache_key"]) for item in batch],
+        "video_names": [str(item["video_name"]) for item in batch],
+        "video_paths": [str(item["video_path"]) for item in batch],
+    }
+
+
+__all__ = [
+    "DecodedMatchedVideo",
+    "MatchedManifestVideoDataset",
+    "collate_matched_video_metadata",
+    "decode_matched_video_last_contiguous_pyav",
+]

1.3B/code/utils/matched_video_feature_cache.py

@@ -0,0 +1,364 @@
+"""Strict safetensors cache for matched-manifest live-teacher features."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+import tempfile
+from typing import Any, Mapping
+
+from safetensors import safe_open
+from safetensors.torch import save_file
+import torch
+
+from utils.matched_video_manifest import MatchedVideoManifestEntry
+
+
+CACHE_FORMAT = "sf_matched_manifest_video_vjepa2_teacher_feature"
+CACHE_FORMAT_VERSION = 1
+CACHE_METADATA_KEY = "sf_matched_video_teacher_feature"
+CACHE_TENSOR_KEY = "features"
+
+# Independent from the fixed-reference central-64 protocol.  Both new student
+# and live teacher use the last 64 frames, i.e. 17..80 for the 81-frame source.
+MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT: dict[str, Any] = {
+    "adapter_contract_version": 1,
+    "contract_version": 1,
+    "crop_size": 256,
+    "expected_source_frame_count": 81,
+    "feature_dim": 1024,
+    "feature_type": "vjepa2_dense_latent",
+    "frame_indices": [17, 80],
+    "frame_sampling": "last_contiguous",
+    "interpolation": "bilinear_antialias",
+    "layer_index": 15,
+    "model_family": "vjepa2-vitl-fpc64-256",
+    "model_config_sha256": (
+        "3dec96fe962e94e569182d3a7b9ef0dd74b6b8c89c337a428e43e10d593e70c9"
+    ),
+    "normalization": "imagenet_default",
+    "num_frames": 64,
+    "num_tokens": 1024,
+    "output_dtype": "float32",
+    "output_shape": [1024, 1024],
+    "processor_config_sha256": (
+        "d2fab4418fc0390b62c4cd72ade56908a7929f80c62288adbe10dd8d23421227"
+    ),
+    "resize_short_side": 292,
+    "sample_axis": "video",
+    "token_alignment": "same_deterministic_indices",
+    "token_indices_sha256": (
+        "a77ee1f9d20cba6b2f2a9a971c3d36610ca7f8aa2a55c616192582665be69d10"
+    ),
+    "token_sampling": "uniform_linspace_round",
+    "transformers_version": "4.53.3",
+}
+
+MATCHED_PYAV_LAST64_DECODER_CONTRACT: dict[str, Any] = {
+    "backend": "pyav",
+    "backend_version": "13.1.0",
+    "contract_version": 1,
+    "expected_frame_indices": [17, 80],
+    "expected_source_frame_count": 81,
+    "frame_count_source": "container_decode_order_with_stream_frames_hint",
+    "num_frames": 64,
+    "output_dtype": "float32",
+    "output_layout": "TCHW",
+    "output_range": [0.0, 1.0],
+    "pixel_format": "rgb24",
+    "selection": "last_contiguous",
+}
+
+
+class MatchedVideoFeatureCacheValidationError(ValueError):
+    """An existing cache file is stale, corrupt, or identity-incompatible."""
+
+
+def canonical_json(value: Any) -> str:
+    return json.dumps(
+        value,
+        ensure_ascii=True,
+        sort_keys=True,
+        separators=(",", ":"),
+        allow_nan=False,
+    )
+
+
+def canonical_sha256(value: Any) -> str:
+    return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest()
+
+
+MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256 = canonical_sha256(
+    MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT
+)
+MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256 = canonical_sha256(
+    MATCHED_PYAV_LAST64_DECODER_CONTRACT
+)
+
+
+def _is_sha256(value: str) -> bool:
+    return len(value) == 64 and all(c in "0123456789abcdef" for c in value)
+
+
+def _entry_context(
+    entry: MatchedVideoManifestEntry | Mapping[str, Any],
+) -> dict[str, Any]:
+    if isinstance(entry, MatchedVideoManifestEntry):
+        result = {
+            "sample_id": entry.sample_id,
+            "video_mtime_ns": int(entry.video_mtime_ns),
+            "video_name": entry.video_name,
+            "video_size_bytes": int(entry.video_size_bytes),
+        }
+    elif isinstance(entry, Mapping):
+        required = {
+            "sample_id", "video_mtime_ns", "video_name", "video_size_bytes",
+        }
+        missing = required - set(entry)
+        if missing:
+            raise ValueError(f"cache sample context is missing {sorted(missing)}")
+        result = {
+            "sample_id": str(entry["sample_id"]),
+            "video_mtime_ns": int(entry["video_mtime_ns"]),
+            "video_name": str(entry["video_name"]),
+            "video_size_bytes": int(entry["video_size_bytes"]),
+        }
+    else:
+        raise TypeError("entry must be MatchedVideoManifestEntry or mapping")
+    if not result["sample_id"] or result["sample_id"] != result["video_name"]:
+        raise ValueError("cache sample_id must equal video_name")
+    if result["video_size_bytes"] <= 0 or result["video_mtime_ns"] <= 0:
+        raise ValueError("cache video stat must be positive")
+    return result
+
+
+def _tensor_sha256(value: torch.Tensor) -> str:
+    value = value.detach().cpu().contiguous()
+    return hashlib.sha256(value.numpy().tobytes(order="C")).hexdigest()
+
+
+class MatchedVideoFeatureCache:
+    """Lazy one-file-per-video FP32 cache with immutable identity checks."""
+
+    def __init__(
+        self,
+        cache_root: str | os.PathLike,
+        *,
+        expected_model_sha256: str,
+        expected_feature_contract_sha256: str,
+        feature_contract: Mapping[str, Any] | None = None,
+        decoder_contract: Mapping[str, Any] | None = None,
+        expected_shape: tuple[int, int] | None = None,
+    ) -> None:
+        self.cache_root = Path(cache_root).expanduser().resolve()
+        self.model_sha256 = str(expected_model_sha256).lower()
+        if not _is_sha256(self.model_sha256):
+            raise ValueError("expected_model_sha256 must be lowercase SHA-256")
+        self.feature_contract = json.loads(canonical_json(
+            feature_contract or MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT
+        ))
+        self.feature_contract_sha256 = canonical_sha256(self.feature_contract)
+        expected_hash = str(expected_feature_contract_sha256).lower()
+        if not _is_sha256(expected_hash):
+            raise ValueError(
+                "expected_feature_contract_sha256 must be lowercase SHA-256"
+            )
+        if self.feature_contract_sha256 != expected_hash:
+            raise ValueError(
+                "feature contract SHA-256 mismatch: "
+                f"expected={expected_hash}, actual={self.feature_contract_sha256}"
+            )
+        self.decoder_contract = json.loads(canonical_json(
+            decoder_contract or MATCHED_PYAV_LAST64_DECODER_CONTRACT
+        ))
+        self.decoder_contract_sha256 = canonical_sha256(self.decoder_contract)
+        contract_shape = tuple(
+            int(value) for value in self.feature_contract.get("output_shape", ())
+        )
+        self.expected_shape = (
+            tuple(int(value) for value in expected_shape)
+            if expected_shape is not None else contract_shape
+        )
+        if len(self.expected_shape) != 2 or min(self.expected_shape) <= 0:
+            raise ValueError("expected feature shape must be positive [N,D]")
+        if contract_shape and self.expected_shape != contract_shape:
+            raise ValueError(
+                "expected_shape must match feature_contract.output_shape"
+            )
+
+    def _key_payload(
+        self, entry: MatchedVideoManifestEntry | Mapping[str, Any],
+    ) -> dict[str, Any]:
+        return {
+            "decoder_contract": self.decoder_contract,
+            "decoder_contract_sha256": self.decoder_contract_sha256,
+            "feature_contract": self.feature_contract,
+            "feature_contract_sha256": self.feature_contract_sha256,
+            "format": CACHE_FORMAT,
+            "format_version": CACHE_FORMAT_VERSION,
+            "model_sha256": self.model_sha256,
+            "sample": _entry_context(entry),
+        }
+
+    def key_for(
+        self, entry: MatchedVideoManifestEntry | Mapping[str, Any],
+    ) -> str:
+        digest = hashlib.sha256()
+        digest.update(b"sf-matched-video-teacher-feature-cache\0")
+        digest.update(canonical_json(self._key_payload(entry)).encode("utf-8"))
+        return digest.hexdigest()
+
+    def path_for_key(self, cache_key: str) -> Path:
+        key = str(cache_key).lower()
+        if not _is_sha256(key):
+            raise ValueError("cache_key must be lowercase SHA-256")
+        return self.cache_root / key[:2] / key[2:4] / f"{key}.safetensors"
+
+    def _expected_metadata(
+        self,
+        entry: MatchedVideoManifestEntry | Mapping[str, Any],
+        cache_key: str,
+        tensor_sha256: str,
+    ) -> dict[str, Any]:
+        result = self._key_payload(entry)
+        result.update({
+            "cache_key": cache_key,
+            "tensor": {
+                "dtype": "float32", "key": CACHE_TENSOR_KEY,
+                "sha256": tensor_sha256, "shape": list(self.expected_shape),
+            },
+        })
+        return result
+
+    def _validate_tensor(self, feature: torch.Tensor) -> torch.Tensor:
+        if not isinstance(feature, torch.Tensor):
+            raise TypeError("teacher feature must be a torch.Tensor")
+        if feature.ndim == 3 and int(feature.shape[0]) == 1:
+            feature = feature[0]
+        if feature.ndim != 2 or tuple(feature.shape) != self.expected_shape:
+            raise ValueError(
+                f"teacher feature shape mismatch: {tuple(feature.shape)} "
+                f"!= {self.expected_shape}"
+            )
+        if not feature.is_floating_point():
+            raise TypeError("teacher feature must be floating point")
+        result = feature.detach().to(device="cpu", dtype=torch.float32).contiguous()
+        if not bool(torch.isfinite(result).all().item()):
+            raise FloatingPointError("teacher feature contains non-finite values")
+        return result
+
+    def load(
+        self,
+        entry: MatchedVideoManifestEntry | Mapping[str, Any],
+        *,
+        cache_key: str | None = None,
+    ) -> torch.Tensor | None:
+        key = self.key_for(entry)
+        if cache_key is not None and str(cache_key) != key:
+            raise MatchedVideoFeatureCacheValidationError(
+                "caller cache_key does not match sample/model/contracts"
+            )
+        path = self.path_for_key(key)
+        if not path.exists():
+            return None
+        if not path.is_file():
+            raise MatchedVideoFeatureCacheValidationError(
+                f"cache path is not a regular file: {path}"
+            )
+        try:
+            with safe_open(path, framework="pt", device="cpu") as handle:
+                if list(handle.keys()) != [CACHE_TENSOR_KEY]:
+                    raise MatchedVideoFeatureCacheValidationError(
+                        f"cache tensor keys mismatch in {path}"
+                    )
+                serialized = (handle.metadata() or {}).get(CACHE_METADATA_KEY)
+                if serialized is None:
+                    raise MatchedVideoFeatureCacheValidationError(
+                        f"cache metadata is missing in {path}"
+                    )
+                metadata = json.loads(serialized)
+                feature = handle.get_tensor(CACHE_TENSOR_KEY).clone()
+        except MatchedVideoFeatureCacheValidationError:
+            raise
+        except Exception as exc:
+            raise MatchedVideoFeatureCacheValidationError(
+                f"failed to read safetensors cache {path}: {exc}"
+            ) from exc
+        feature = self._validate_tensor(feature)
+        expected = self._expected_metadata(entry, key, _tensor_sha256(feature))
+        if metadata != expected:
+            raise MatchedVideoFeatureCacheValidationError(
+                f"cache metadata/identity mismatch (stale or corrupt) in {path}"
+            )
+        return feature
+
+    def write(
+        self,
+        entry: MatchedVideoManifestEntry | Mapping[str, Any],
+        feature: torch.Tensor,
+        *,
+        cache_key: str | None = None,
+        overwrite: bool = False,
+    ) -> Path:
+        key = self.key_for(entry)
+        if cache_key is not None and str(cache_key) != key:
+            raise MatchedVideoFeatureCacheValidationError(
+                "caller cache_key does not match sample/model/contracts"
+            )
+        feature = self._validate_tensor(feature)
+        tensor_hash = _tensor_sha256(feature)
+        destination = self.path_for_key(key)
+        if destination.exists() and not overwrite:
+            existing = self.load(entry, cache_key=key)
+            if existing is None or _tensor_sha256(existing) != tensor_hash:
+                raise MatchedVideoFeatureCacheValidationError(
+                    "refusing to replace a valid cache entry with different features: "
+                    f"{destination}"
+                )
+            return destination
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        metadata = self._expected_metadata(entry, key, tensor_hash)
+        descriptor, temporary_name = tempfile.mkstemp(
+            dir=destination.parent,
+            prefix=f".{destination.name}.", suffix=".tmp",
+        )
+        os.close(descriptor)
+        temporary = Path(temporary_name)
+        try:
+            temporary.unlink()
+            save_file(
+                {CACHE_TENSOR_KEY: feature}, str(temporary),
+                metadata={CACHE_METADATA_KEY: canonical_json(metadata)},
+            )
+            with temporary.open("rb") as handle:
+                os.fsync(handle.fileno())
+            if overwrite:
+                os.replace(temporary, destination)
+            else:
+                try:
+                    os.link(temporary, destination)
+                except FileExistsError:
+                    existing = self.load(entry, cache_key=key)
+                    if existing is None or _tensor_sha256(existing) != tensor_hash:
+                        raise MatchedVideoFeatureCacheValidationError(
+                            "concurrent cache writer produced different features "
+                            f"for {key}"
+                        )
+                temporary.unlink(missing_ok=True)
+        finally:
+            temporary.unlink(missing_ok=True)
+        self.load(entry, cache_key=key)
+        return destination
+
+
+__all__ = [
+    "CACHE_FORMAT", "CACHE_FORMAT_VERSION", "CACHE_METADATA_KEY",
+    "CACHE_TENSOR_KEY", "MATCHED_PYAV_LAST64_DECODER_CONTRACT",
+    "MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256",
+    "MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT",
+    "MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256",
+    "MatchedVideoFeatureCache", "MatchedVideoFeatureCacheValidationError",
+    "canonical_json", "canonical_sha256",
+]