[misc] correct fp32 checkpoint size to measured 741 GiB (was ~528 estimate)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

.agents/lessons/2026-08-14_h3-dmd2-training-preflight-rules.md

@@ -0,0 +1,120 @@
+---
+date: 2026-08-14
+experiment: feat/h3-dmd2-vsa — data-free DMD2 distillation of MiniMax-H3 (33B, video+audio) on NVL72 GB200
+category: training
+severity: critical
+---
+
+# Preflight Rules for Large-Model Distillation Runs (distilled from the H3 DMD2 v1–v5 lineage)
+
+Every rule below is a real failure we shipped, diagnosed, and fixed across
+five production launches. Rules are grouped; each states the check, the
+failure it prevents, and where it is now enforced in code (if it is).
+
+## Numerics
+
+1. **Master weights must be fp32.** In-place optimizer steps on bf16 params
+   round away updates below ~half an ulp of each weight's magnitude:
+   O(1)-magnitude norm gains froze **bit-identically for 1000 steps**
+   (205/211 tensors) at lr 2e-6, silently capping convergence. Enforced:
+   `fastvideo/train/trainer.py::_verify_master_weight_precision` and
+   `fastvideo/training/training_utils.py::verify_master_weight_precision`
+   hard-fail at train start (opt-outs:
+   `training.model.allow_low_precision_master_weights`,
+   `FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS=1`).
+2. **Optimizer second moments must be fp32 regardless of param dtype.**
+   `torch.zeros_like(p)` inherits the starved dtype; bf16 `exp_avg_sq` loses
+   g² increments below ~0.4% of its running magnitude. Enforced in
+   `AdamWBeta1Zero` (fp32 state + fp32 update math for non-fp32 params).
+3. **Do not copy inference precision defaults into training configs.**
+   `PipelineConfig.dit_precision` defaults to bf16 (inference); training
+   YAMLs that copy it inherit rule-1's failure. `dit_precision: fp32` costs
+   only sharded-master memory — FSDP `param_dtype=bf16` keeps compute speed
+   (measured: no step-time change at 33B).
+4. **Post-hoc freeze check (cheap, mapping-free):** base checkpoints are
+   stored bf16, so any fp32-updated tensor is almost surely no longer
+   exactly bf16-representable. Census a checkpoint's norm params for
+   `x != bf16(x)` fractions — near-zero fractions mean weights are not
+   moving. Verified on the v5 lineage: at checkpoint-500 with fp32 masters,
+   210/210 student norm tensors moved (median 99.9% of elements), vs
+   205/211 bit-frozen after 1000 bf16 steps.
+
+## Schedules and shifts
+
+5. **Recompute every timestep-space constant through the model's actual
+   shift map** (`sigma' = s*t / (1 + (s-1)t)`). The Wan-copied ladder
+   `[1000,757,522]` under H3's video shift 12 produced sigma
+   {1.0, 0.974, 0.929} — a near-no-op middle step and an 86% final jump.
+   Re-spaced to `[1000,667,333]` (uniform base-t, the same design rule the
+   Wan recipe uses pre-warp).
+6. **Check supervision *density* in sigma space, not t space.** Uniform-t
+   sampling under shift 12 put 57% of score-model draws at sigma_v > 0.9 and
+   ~0% below 0.2. Fix: `method.score_timestep_shift` (sample uniform in
+   shifted-sigma, invert to base-t). Verify any new sampler with a quantile
+   census before launching.
+7. **Shared timestep-ratio bounds hit modalities asymmetrically.** With
+   min_ratio 0.02, video's supervision floor was sigma 0.197 while audio's
+   was 0.057 (shift 12 vs 3). Derive per-modality sigma floors from the
+   ratios; H3 uses 0.005.
+8. **Know whether a step list is consumed in base-t or shifted space.**
+   `warp_denoising_step` maps the list through the (video) scheduler's
+   shifted grid; H3's adapter *also* shifts per modality internally —
+   enabling warp double-shifts. Wan needs warp on; H3 needs it off.
+
+## Packed multi-modality
+
+9. **Never take a single mean loss over a packed multi-modality sequence.**
+   H3 packs video:audio at ~270:1 elements; a packed mean silently trains
+   only video. Compute per-modality losses via `modality_slices()` and log
+   each stream (`*_video`, `*_audio`) separately.
+
+## Distillation specifics
+
+10. **Check whether the teacher is guidance-distilled before configuring
+    CFG.** H3's released checkpoint rejects CFG at inference; correct
+    setting is `real_score_guidance_scale: 1.0`, and DMD2 then skips the
+    unconditional teacher forward entirely (free speedup).
+11. **Read DMD scalars as game state, not quality.** Generator loss =
+    normalized critic–teacher disagreement at student samples (healthy band
+    ~0.2–1.0; it *rises* from ~0 as the critic specializes away from its
+    teacher init). Critic flow-matching loss = tracking meter (falling =
+    keeping up). The 10x-LR blowup signature: critic loss exploding (5.2)
+    plus student grad-norms pinned above clip (1.8–2.5 vs 1.0) — a run that
+    reports 10x LR but delivers ~4x, into a broken critic.
+12. **Grad-norm vs clip tells you the LR you actually delivered.** Isolated
+    spikes above clip that mean-revert are normal equilibration; sustained
+    saturation means the configured LR is a lie.
+13. **Sampler A/Bs need a teacher control and matched seeds.** The
+    stochastic-vs-deterministic verdict on a student is meaningless without
+    knowing the teacher's own behavior under both hop rules (teacher x0
+    overshoot at sigma≈1, std 2.33, dominated ours). File-size heuristics
+    on encoded video are not a quality signal — high bitrate can be noise.
+
+## Run ops
+
+14. **Requeue-safe by default:** `resume_from_checkpoint: latest` +
+    `sbatch --requeue`. DCP restores optimizer LRs over the YAML — use
+    `reset_lr_on_resume: true` for LR-change experiments.
+15. **Budget checkpoint disk before launch:** save size × retention vs
+    filesystem free. fp32 masters tripled saves to 741 GB (measured; student+critic params + fp32 second moments); lustre at 98%
+    would have killed the run at its 4th save. `checkpoints_total_limit`
+    is part of the launch math, not a detail.
+16. **Verify W&B auth (`wandb.Api().viewer` with the job's $HOME) before
+    every launch**; the tracker now degrades online→offline→none instead of
+    crashing 32 GPUs on a rotated key. Confirm step-0 media actually
+    uploaded — the uploader has died on the first video burst before.
+17. **Unique `run_name` per launch attempt.** Three runs sharing one name
+    made W&B show another run's config (the "DP=12" confusion).
+18. **Cluster quirks (Slinky/GB200):** compute pods have no /home (export
+    HOME=lustre in every sbatch; submit with a lustre --chdir); PENDing
+    jobs are what make the autoscaler create nodes (warmup-ladder + race
+    for multi-node); released nodes deregister in seconds; boot watchdog +
+    NCCL env pins for the silent IB comm-init wedge; `scontrol release`
+    after "launch failed requeued held"; Slurm requeues silently truncate
+    stdout (check elapsed vs log).
+19. **Exports and deletions:** cross-user hardlinks are blocked
+    (fs.protected_hardlinks) — use `--link-base` symlink exports (~66 GB vs
+    465 GB); always clear stale `*.safetensors.index.json` when rewriting a
+    module dir (a leftover base index shadows the fresh weights); before any
+    rm of run artifacts, verify no symlinks/hardlinks point into originals
+    and that the target is regenerable.

HANDOFF-h3-dmd2-vsa.md

@@ -0,0 +1,88 @@
+# Handoff: VSA-H3 DMD2 overfit — run directly on tray 1795370
+
+State as of 2026-08-09. Everything is staged on the tray; only the launch is
+left (my login-node srun path was blocked by per-user thread exhaustion —
+running inside the tray's container sidesteps it entirely).
+
+## The experiment
+
+Single-sample DMD2 overfit on `c17_v2` (your prompt-rewrite v2 generation,
+1344×768 / 5.2 s / with audio), parity-first: **student on the VSA-H3 backend
+at sparsity 0.0** (parity-proven ≤2e-4 vs dense in unit tests) vs **teacher +
+critic on FLASH_ATTN with FA4**. 3-step schedule `[1000, 757, 522]`,
+`generator_update_interval: 5` (TTUR: critic every step, student every 5th),
+4 GPUs, sp=4, pure-FSDP sharding (hsdp 1×4), all three models
+gradient-checkpointed, batch 1, 2000 steps, ckpt+validation every 250.
+
+## What is already staged
+
+- **Code**: `/mnt/fv-h3ssim` on branch `h3-dmd2-vsa` (= `feat/h3-dmd2-vsa`
+  @ d81f1f166 = DMD2-for-H3 + PR #1695 VSA backend + per-role backend wiring
+  + wandb credential guard). Also pushed to origin as `feat/h3-dmd2-vsa`.
+- **Data**: `/mnt/h3-dmd2-overfit/data/data_00000.parquet` — one preprocessed
+  t2va record (video+audio VAE latents + text embedding 225×5120).
+- **Validation prompts**: `/mnt/h3-dmd2-overfit/c17_validation.json` (the c17
+  v2 prompt, 3-step sampling) — the YAML may still point at the Wan
+  placeholder; step 2 below fixes it idempotently.
+- **Config**: `examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml`
+- Raw inputs (for reference): `/mnt/h3-dmd2-overfit/c17_v2.mp4`, `c17_v2_prompt.txt`.
+
+## Run it (on the tray)
+
+```bash
+# 0. enter the persistent container (created by my earlier runs)
+docker exec -it fv-dev-1795370-0 bash
+# (if it's gone: ~/docker.sh style run with /home/scratch.willlin_ent -> /mnt)
+
+# 1. clean any leftover trainer from my attempts + confirm GPUs are free
+ps -eo pid,args | grep entrypoint.train | grep -v grep   # kill any PIDs listed
+nvidia-smi --query-gpu=memory.used --format=csv,noheader  # expect ~0 MiB x4
+
+# 2. env + config fix (idempotent)
+cd /mnt/fv-h3ssim
+export PYTHONPATH=/mnt/studio8/pyoverride   # routes 'import fastvideo' to THIS checkout
+export FASTVIDEO_FA4=1                      # teacher/critic FLASH_ATTN -> FA4 path
+export WANDB_API_KEY=8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde
+sed -i "s|dataset_file: examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json|dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json|" \
+  examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml
+
+# 3. launch (interactive; or wrap in setsid nohup ... & for detach)
+torchrun --nproc_per_node 4 -m fastvideo.train.entrypoint.train \
+  --config examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml \
+  2>&1 | tee /mnt/h3-dmd2-overfit/logs/train.log
+```
+
+wandb: project `h3-dmd2-vsa`, run `dmd2_vsa0_overfit` — the URL prints within
+the first seconds of trainer init (`https://wandb.ai/...`). Logged per step:
+total/generator/fake-score loss, update_student flag, step_time_sec,
+vsa_sparsity.
+
+## What to expect / what parity means
+
+- Trio load: ~3×(33B DiT) FSDP-sharded + text encoder; first load on this
+  tray is warm for the text encoder (preprocess touched it) but cold for the
+  3 transformer copies — expect ~10-20 min before step 1.
+- Healthy signals: both losses finite from step 1; fake-score loss drops
+  fast (critic fitting one sample); generator loss moves every 5th step.
+- Parity claim to check: with sparsity 0.0 the student's VSA-H3 path should
+  train indistinguishably from a dense student (the sparsity-0 mask selects
+  every tile; kernel parity ≤2e-4). If it diverges/NaNs where a dense run
+  wouldn't, that's a VSA-H3-under-autograd bug — grab
+  `/mnt/h3-dmd2-overfit/logs/train.log` and the wandb run.
+- Next knob after parity holds: `training.vsa.sparsity` (e.g. 0.5 → 0.875
+  → 0.9) in the YAML — student-only; teacher/critic stay dense by
+  construction (`models.*.attention_backend`).
+
+## Known traps (all pre-paid)
+
+- `PYTHONPATH=/mnt/studio8/pyoverride` is mandatory — without it,
+  `import fastvideo` resolves to `/mnt/FastVideo` (the maintainer's checkout)
+  via the venv's editable-install finders, and cwd `/mnt/FastVideo` shadows too.
+- Don't set `FASTVIDEO_ATTENTION_BACKEND` globally — per-role backends come
+  from the YAML (`models.student/teacher/critic.attention_backend`).
+- The tile-buffer autograd fix is in this branch (grad paths get fresh
+  buffers); if you see "modified by an inplace operation" at backward, you're
+  on a stale checkout — `git -C /mnt/fv-h3ssim log --oneline -1` should show
+  d81f1f16.
+- Checkpoints land in `outputs/minimax_h3_dmd2_vsa0_overfit/` under the
+  checkout (NFS); validation videos under the same tree per callback config.

HANDOFF-slurm-32xgb200.md

@@ -0,0 +1,120 @@
+# Handoff: MiniMax-H3 data-free DMD2 on 8x 4xGB200 (Slurm)
+
+State as of 2026-08-11 on branch `feat/h3-dmd2-vsa`. Everything needed for
+the 32-GPU run is committed; what remains is cluster-side: fill in the
+`CLUSTER:` knobs in the sbatch, swap in the real dataset, and submit.
+
+## What this run is
+
+Data-free DMD2 distillation (`rollout_mode: simulate` — the student builds
+its own trajectory from pure noise through `[1000, 757, 522]`; no video/audio
+latents consumed, text conditioning only) of the 33B dual-modality H3 DiT.
+All three roles (student/teacher/critic) on dense FLASH_ATTN with FA4.
+**Global batch 64** = 8 data-parallel groups (one per tray) x
+`gradient_accumulation_steps: 8`.
+
+- Config: `examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_32gpu.yaml`
+- Launcher: `examples/train/slurm/dmd2_32xgb200.sbatch`
+- Docs: `examples/train/configs/distribution_matching/minimax_h3/README.md`
+
+## Everything below is already validated on a single 4xGB200 node
+
+- The exact simulate recipe trained stably (wandb `h3-dmd2-vsa/z2z91ruf`):
+  TTUR cadence correct, generator loss settles into the expected adversarial
+  0.2–1.0 band, no NaN/OOM. ~10 s/it critic steps, ~20 s/it generator steps
+  (every 5th) at 1344x768x124f, sp=4.
+- Effective-batch-2 via gradient accumulation over two distinct videos
+  (wandb `58dr571y`, 1800+ steps): both modalities' losses decline
+  monotonically (video 0.128->0.092, audio 0.179->0.053 by step 1000).
+  Accumulation is the batching mechanism — H3's packed pipeline is batch-1
+  per forward by design (variable caption lengths cannot stack).
+- Memory: three 33B roles fit at sp=4 / HSDP shard=4 with full gradient
+  checkpointing and `betas: [0.0, 0.999]` (buffer-free AdamW). Do NOT try
+  sp=2 or fewer than 4 GPUs per model replica: the backward working set
+  alone OOMs 184 GiB cards (measured 178 GiB before optimizer states exist).
+- Triton VSA kernel backward fix, bf16 autocast boundaries, validation
+  offload fix, sigma-exact validation — all on the branch (see git log).
+
+## Cluster pre-flight
+
+1. **Code**: checkout `feat/h3-dmd2-vsa` on the shared FS; venv with the
+   repo's CUDA install. `run.sh` sets PYTHONPATH to the repo root and cd's
+   into it — on trays that carry `/mnt/studio8/pyoverride`, make sure it is
+   NOT in PYTHONPATH or `import fastvideo` resolves to the wrong checkout.
+2. **Weights**: `/mnt/models/MiniMax-H3` must be visible from compute nodes
+   (all three roles `init_from` it).
+3. **Data**: the config ships pointing at the 2-row smoke-test set
+   (`/mnt/h3-dmd2-overfit/data_bs2`) — fine for the first submission, swap
+   `training.data.data_path` (and the validation `dataset_file`) for the
+   real preprocessed t2va parquet dir before a production run. Simulate mode
+   only reads text conditioning from the rows.
+4. **sbatch knobs**: partition/account (marked `CLUSTER:`), and if trays
+   require containers, wrap the srun payload (pyxis
+   `--container-image/--container-mounts=/mnt:/mnt,/dev/shm`) — hooks are
+   marked in the script.
+5. **W&B**: `~/.netrc` on a home visible from compute nodes, or export
+   `WANDB_MODE=offline` at submit time.
+6. **Login-node limits**: a previous srun attempt died to per-user thread
+   exhaustion on the login node — submit with `sbatch` (batch daemon runs
+   the script on the first compute node), don't launch interactive `srun`
+   pipelines from the login shell.
+
+## Launch
+
+```bash
+sbatch examples/train/slurm/dmd2_32xgb200.sbatch                 # 8 trays / 32 GPUs
+sbatch --nodes=2 examples/train/slurm/dmd2_32xgb200.sbatch       # scaled smoke test
+CONFIG=<other.yaml> sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+```
+
+The script derives world size / rendezvous / `hsdp_replicate_dim` from the
+actual allocation, so a smaller `--nodes` shrinks the DP width (and global
+batch) consistently. Recommended sequence: `--nodes=1` (reproduces the
+validated local run), then `--nodes=2`, then the full 8.
+
+## What to expect
+
+- **First boot per node**: ~10-13 min per 33B role from cold NFS (~35 min
+  worst case for the trio, plus text encoder/VAEs) — then each node writes
+  its weight cache to `/dev/shm/fastvideo-wcache`
+  (`FASTVIDEO_WEIGHT_SHARD_CACHE_PER_NODE=1` is set by the sbatch so every
+  tray caches, not just tray 0). **Relaunches: ~2-4 s per role.** Note
+  `/dev/shm` may be purged by node-local cleaners — a purged cache degrades
+  to a full load and rewrites itself, never fails the run.
+- **Step time**: per-tray compute is identical to the validated 4-GPU run;
+  with accum 8, expect roughly 8 x (10 s + 2 s amortized generator surcharge)
+  ≈ 90-100 s per optimizer step, ~4000 steps ≈ 4.5 days. Cut
+  `max_train_steps` or validation frequency to taste.
+- **Logs**: one dir per node under `examples/train/logs/<jobid>-node<k>/`;
+  wandb project `h3-dmd2-vsa`, run `dmd2_simulate_fa4_32gpu` (rank 0 only).
+- **Checkpointing** is ON every 500 steps (DCP; large — ~120 GiB+ per save
+  for two trainable roles). Point `output_dir` somewhere with headroom or
+  set `training_state_checkpointing_steps: 0` to disable.
+- **Validation** every 250 steps samples every prompt in the validation
+  json at the training sigmas; with the smoke-test json that is 2 videos,
+  ~3 min.
+
+## Known traps (all pre-paid on the branch)
+
+- Don't set `FASTVIDEO_ATTENTION_BACKEND` globally — per-role backends come
+  from the YAML.
+- `FASTVIDEO_FA4=1` is required (sbatch sets it); without it FLASH_ATTN
+  falls back to a slower path.
+- `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` is required (sbatch
+  sets it); without it fragmentation OOMs appear after a few steps.
+- If a VSA student is used later (`attention_backend: VIDEO_SPARSE_ATTN_H3`
+  + `training.vsa.sparsity`): the Triton backward bugfix is on this branch —
+  do not run VSA training from an older checkout. Sparsity ~0.9+ matches
+  dense wall-clock in training and ~2x in validation sampling; the win is
+  inference-side.
+- NCCL fabric env (IB HCA / socket ifname) is cluster-specific — hook in the
+  sbatch.
+
+## If something breaks
+
+Grab the failing node's `examples/train/logs/<jobid>-node<k>/` log and the
+wandb run. The historically likely failure classes and their signatures:
+OOM ledger lines (memory — check betas are [0.0, 0.999] and grad ckpt full),
+`attempting to assign a gradient with device type` (a module was CPU-offloaded
+mid-training — should be fixed, see validation.py), rendezvous timeout at
+boot (MASTER_ADDR/port reachability between trays).

examples/distill/MiniMax-H3/distill_dmd.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+# DMD2 distillation for MiniMax H3 (joint video + audio) on the modular
+# trainer (fastvideo/train). Unlike the legacy SFWan scripts, the modern
+# stack is YAML-driven: the {student, teacher, critic} trio, DMD2 method
+# knobs, and callbacks all live in the run config.
+
+set -euo pipefail
+
+export TOKENIZERS_PARALLELISM=false
+export MASTER_PORT=${MASTER_PORT:-29513}
+# H3 training pins the dense TORCH_SDPA backend; do not export a sparse
+# attention backend here.
+
+NUM_GPUS=${NUM_GPUS:-4}
+CONFIG=${CONFIG:-examples/train/configs/distribution_matching/minimax_h3/dmd2_t2va.yaml}
+DATA_DIR=${DATA_DIR:-data/crush-smol_h3_t2va_single_sample_preprocessed}
+OUTPUT_DIR=${OUTPUT_DIR:-outputs/minimax_h3_dmd2_3steps}
+
+torchrun \
+  --nnodes 1 \
+  --master_port "$MASTER_PORT" \
+  --nproc_per_node "$NUM_GPUS" \
+  -m fastvideo.train.entrypoint.train \
+  --config "$CONFIG" \
+  --training.distributed.num_gpus "$NUM_GPUS" \
+  --training.distributed.sp_size "$NUM_GPUS" \
+  --training.distributed.hsdp_shard_dim "$NUM_GPUS" \
+  --training.data.data_path "$DATA_DIR" \
+  --training.checkpoint.output_dir "$OUTPUT_DIR"

examples/distill/MiniMax-H3/overfit_vsa0.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+# Single-sample VSA-H3 DMD2 overfit for MiniMax H3 (sparsity 0.0, parity vs
+# dense). Two steps:
+#   (a) preprocess ONE mp4 (with audio) + prompt into the t2va parquet row
+#       the H3 dataloader consumes (768x1344, 124 frames @ 24 fps, 32 kHz);
+#   (b) train on 4 GPUs: student on VIDEO_SPARSE_ATTN_H3, teacher/critic on
+#       FLASH_ATTN with the FA4 path (FASTVIDEO_FA4=1).
+#
+# Usage:
+#   VIDEO=/path/clip.mp4 PROMPT="a caption" bash examples/distill/MiniMax-H3/overfit_vsa0.sh
+#   SKIP_PREPROCESS=1 bash examples/distill/MiniMax-H3/overfit_vsa0.sh   # data already written
+
+set -euo pipefail
+
+export TOKENIZERS_PARALLELISM=false
+export MASTER_PORT=${MASTER_PORT:-29513}
+
+NUM_GPUS=${NUM_GPUS:-4}
+CONFIG=${CONFIG:-examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml}
+DATA_DIR=${DATA_DIR:-/mnt/h3-dmd2-overfit/data}
+OUTPUT_DIR=${OUTPUT_DIR:-outputs/minimax_h3_dmd2_vsa0_overfit}
+MODEL_PATH=${MODEL_PATH:-data/models/MiniMax-H3}
+
+# ---------------------------------------------------------------- preprocess
+if [ -z "${SKIP_PREPROCESS:-}" ]; then
+  VIDEO=${VIDEO:?set VIDEO=/path/to/clip.mp4 (with soundtrack) or SKIP_PREPROCESS=1}
+  PROMPT=${PROMPT:?set PROMPT="text prompt for the clip"}
+  python -m fastvideo.pipelines.preprocess.preprocess_minimax_h3_overfit \
+    --video "$VIDEO" \
+    --prompt "$PROMPT" \
+    --output-dir "$DATA_DIR" \
+    --model-path "$MODEL_PATH"
+fi
+
+# --------------------------------------------------------------------- train
+# FA4 flash-attention path for the dense teacher/critic FLASH_ATTN roles.
+export FASTVIDEO_FA4=1
+
+torchrun \
+  --nnodes 1 \
+  --master_port "$MASTER_PORT" \
+  --nproc_per_node "$NUM_GPUS" \
+  -m fastvideo.train.entrypoint.train \
+  --config "$CONFIG" \
+  --training.distributed.num_gpus "$NUM_GPUS" \
+  --training.distributed.sp_size "$NUM_GPUS" \
+  --training.distributed.hsdp_shard_dim "$NUM_GPUS" \
+  --training.data.data_path "$DATA_DIR" \
+  --training.checkpoint.output_dir "$OUTPUT_DIR"

examples/train/configs/distribution_matching/minimax_h3/README.md

@@ -0,0 +1,87 @@
+# MiniMax-H3 training configs (DMD2 distillation + SFT overfit)
+
+Configs for training the 33B dual-modality (video+audio) MiniMax-H3 DiT with
+the modular trainer (`fastvideo/train/`), validated on 4x GB200 (192 GB,
+aarch64). SFT overfit configs live in
+[`../../fine_tuning/minimax_h3/`](../../fine_tuning/minimax_h3/); the 8-tray
+Slurm launch is [`../../../slurm/dmd2_32xgb200.sbatch`](../../../slurm/dmd2_32xgb200.sbatch).
+
+## Configs
+
+| Config | What it is |
+|---|---|
+| `dmd2_vsa0_overfit.yaml` | Single-sample DMD2, student on VIDEO_SPARSE_ATTN_H3 at sparsity 0 (parity vs dense), teacher/critic on FLASH_ATTN/FA4, `rollout_mode: data_latent` |
+| `dmd2_simulate_fa4_overfit.yaml` | **Data-free** DMD2 (`rollout_mode: simulate`): student rolls out from pure noise through `dmd_denoising_steps`; parquet VAE latents ignored, only text conditioning used. All roles dense FA4 |
+| `dmd2_simulate_fa4_32gpu.yaml` | The simulate recipe scaled to 8 trays x 4 GPUs, global batch 64 |
+| `../../fine_tuning/minimax_h3/sft_*.yaml` | SFT overfit family: FA4 dense control, VSA at sparsity 0 / 0.8 / 0.9 / 0.95 / 0.97, and `sft_fa4_bs2_overfit.yaml` (effective batch 2 over two videos) |
+
+## Launch
+
+```bash
+# single node (4x GB200)
+FASTVIDEO_FA4=1 \
+FASTVIDEO_WEIGHT_SHARD_CACHE=/dev/shm/fastvideo-wcache \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+bash examples/train/run.sh examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_overfit.yaml
+
+# 8 trays / 32 GPUs from a Slurm login node (global batch 64)
+sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+```
+
+`FASTVIDEO_FA4=1` selects the FA4 path inside the FLASH_ATTN backend and is
+required for the dense roles. Per-role backends come from
+`models.<role>.attention_backend` — do not set `FASTVIDEO_ATTENTION_BACKEND`
+globally.
+
+## Topology and batch semantics
+
+- `sp_size: 4` — one packed `[text|cond|audio|video]` document is
+  sequence-parallel across the 4 GPUs of a node; HSDP mesh
+  `(replicate=nodes, shard=4)` keeps parameters sharded only inside a node
+  (~145 GiB/rank for DMD2's three 33B roles with full gradient checkpointing).
+- **The H3 packed pipeline is batch-1 per forward by design** (documents with
+  different caption lengths cannot stack). Batching composes through
+  `training.loop.gradient_accumulation_steps`, which the trainer loops
+  natively — global batch = DP groups x accumulation rounds. The 32-GPU
+  config's global 64 = 8 DP groups x 8 rounds; validated locally as
+  effective-batch-2 (`sft_fa4_bs2_overfit.yaml`, two videos, both losses
+  declining, ~10 s/step = 2 micro fwd+bwd).
+- DMD2 TTUR: critic every step, generator every
+  `generator_update_interval`-th. `rollout_mode: simulate` costs ~4 dense
+  forwards per critic step and ~7 + 2 backwards per generator step
+  (~10/20 s at 1344x768x124f on GB200).
+
+## Measured reference points (1344x768, 124 frames, sp=4)
+
+- SFT step time: FA4 dense 4.9 s/it; VSA 90% 4.8; 95% 4.49; 97% 4.33;
+  VSA-at-density-1 11.2 (VSA machinery overhead — attention is only ~1/3 of
+  step FLOPs, so sparsity is mostly an inference win: 3-step validation
+  sample ~120 s dense vs ~59 s at 95%).
+- 2x GB200 does NOT fit this model/sequence without CPU offload: the first
+  backward alone allocates ~178 GiB at sp=2 (params 31 + grads 31 + ~116
+  working set) before optimizer states exist.
+- `betas: [0.0, 0.999]` auto-selects a buffer-free AdamW
+  (`fastvideo/train/utils/optimizer.py`), saving one full set of optimizer
+  state (~15.4 GiB/rank/model) — required for DMD2's two trainable roles.
+
+## Weight shard cache
+
+`FASTVIDEO_WEIGHT_SHARD_CACHE=<tmpfs dir>` caches each rank's post-shard
+DTensor chunks after the first full load
+(`fastvideo/models/loader/shard_cache.py`); relaunches rebuild each 33B role
+in ~2-4 s instead of ~10 min of cold-NFS reads (DMD2's three roles share one
+entry). On multi-node runs with node-local cache dirs set
+`FASTVIDEO_WEIGHT_SHARD_CACHE_PER_NODE=1` (the sbatch does) so every node
+writes its own copy. Any validation failure degrades to the normal full load.
+
+## Data
+
+Single/two-sample overfit parquets are produced by
+`fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py` (one mp4
+with soundtrack + prompt -> one t2va row; >=124 frames at 24 fps and an audio
+track are required). The preprocessor writes exactly one
+`data_00000.parquet` per output dir — build multi-sample sets by copying rows
+into one dir as `data_00000.parquet`, `data_00001.parquet`, ... The
+validation callback samples every prompt in `callbacks.validation.dataset_file`
+at the method's exact training sigmas (`dmd_denoising_steps` is injected into
+the validation pipeline).

examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_32gpu.yaml

@@ -0,0 +1,123 @@
+# 32x GB200 (8 trays x 4 GPUs) data-free DMD2 on MiniMax H3, all-dense FA4.
+#
+# Topology: sp_size=4 within a tray (one packed document per SP group), HSDP
+# mesh (replicate=8, shard=4) — params sharded only inside a tray (keeps the
+# proven ~145 GiB/rank profile from the 4-GPU runs), gradients replicated
+# across trays. 32/4 = 8 data-parallel groups.
+#
+# Global batch 64 = 8 DP groups x gradient_accumulation_steps 8 x micro-batch
+# 1. H3's packed-sequence pipeline is batch-1 per forward by design (one
+# interleaved [text|cond|audio|video] document); batching composes through
+# accumulation, which the trainer loops natively.
+#
+# Launch: sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+#
+# NOTE: data_path below still points at the 2-row overfit set — good for a
+# 32-GPU smoke test only. Swap in the real preprocessed t2va dataset before a
+# production run.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  cfg_uncond:
+    text: zero
+
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 32
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 8
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data_bs2  # smoke-test set; swap for real data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    # 8 accumulation rounds x 8 DP groups = global batch 64 per optimizer step.
+    gradient_accumulation_steps: 8
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_dmd2_simulate_fa4_32gpu
+    # Multi-node runs should checkpoint; DCP saves are large — point
+    # output_dir at a volume with headroom or set back to 0 to disable.
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_simulate_fa4_32gpu
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_bs2_validation.json
+    every_steps: 250
+    run_at_start: true
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_32gpu_vidprom.yaml

@@ -0,0 +1,124 @@
+# 32x GB200 (8 trays x 4 GPUs) data-free DMD2 on MiniMax H3 over the VidProM
+# prompt set — the production variant of dmd2_simulate_fa4_32gpu.yaml.
+#
+# Identical validated topology and hyperparameters; only the data source and
+# cluster paths differ:
+#   - training data: 63,136 text-only conditioning rows (VidProM-H3-T2VA
+#     prompts encoded by preprocess_minimax_h3_text_only.py — Qwen3-VL
+#     layer-50, fp32) instead of the 2-row t2va smoke set. Latents are never
+#     read in simulate mode; text_only rows carry exactly what the run uses.
+#   - paths point at /mnt/lustre (nvl-gb200 cluster layout).
+#
+# Launch: REPO=/mnt/lustre/vlm-wlsaidhi/fastvideo/FastVideo \
+#         CONFIG=<this file> sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  cfg_uncond:
+    text: zero
+
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 32
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 8
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/lustre/vlm-wlsaidhi/fastvideo/data/vidprom_h3_text_only
+    preprocessed_data_type: text_only
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    # 8 accumulation rounds x 8 DP groups = global batch 64 per optimizer step.
+    gradient_accumulation_steps: 8
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-wlsaidhi/fastvideo/outputs/minimax_h3_dmd2_simulate_fa4_32gpu_vidprom
+    # DCP saves are ~120 GiB+ per checkpoint for the two trainable roles;
+    # lustre is the only shared volume with that headroom (monitor `df`).
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_simulate_fa4_32gpu_vidprom
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    # 64 held-out prompts (Wan2.1 DMD2 validation_64 rewritten into H3's
+    # three-field document format per MiniMax's prompt guide; verified
+    # disjoint from the VidProM training captions). 8 DP groups render 8
+    # videos each per pass. Repo copy: validation_wan64_h3.json next to
+    # this config.
+    dataset_file: /mnt/lustre/vlm-wlsaidhi/fastvideo/data/h3_validation_wan64.json
+    every_steps: 100
+    run_at_start: true
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_overfit.yaml

@@ -0,0 +1,130 @@
+# Single-sample DMD2 DATA-FREE overfit: rollout_mode=simulate on MiniMax H3
+# with FLASH_ATTN (FA4) for every role — no VSA in the loop.
+#
+# - simulate rollout: the student builds its own trajectory from pure noise
+#   through dmd_denoising_steps (no-grad prefix steps, one with-grad step at a
+#   sampled target index). VAE latents in the parquet are ignored (zeros
+#   latents sized from training.data dims); only the text conditioning is
+#   consumed.
+# - Student/teacher/critic: FLASH_ATTN (run with FASTVIDEO_FA4=1). DMD2 still
+#   passes attn_kind="vsa" for the student rollout, which degrades to dense
+#   because attn_metadata_vsa stays None for non-VSA backends.
+# - Data: same ONE preprocessed t2va parquet row (768x1344, 124 frames @ 24
+#   fps) as dmd2_vsa0_overfit.yaml; see examples/distill/MiniMax-H3/.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  # H3 has no negative-prompt encoder at training time; unconditional teacher
+  # forwards zero the text embeddings instead.
+  cfg_uncond:
+    text: zero
+
+  # Critic optimizer (required — no fallback to training.optimizer)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    # Preprocessed-latents parquet dir; simulate mode only reads the text
+    # conditioning from it (latents are zeros sized from the dims below).
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_dmd2_simulate_fa4_overfit
+    # 0 disables training-state checkpointing (no resume; avoids ~400+ GiB
+    # DCP saves on the 87%-full volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_simulate_fa4_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0 (costs the validation pipeline
+    # build upfront: text-encoder + VAE loads before the first train step).
+    run_at_start: true
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+  # EMA disabled: its GPU shadow costs ~15.4 GiB/rank (OOM on 184 GiB cards).
+  # ema:
+  #   _target_: fastvideo.train.callbacks.ema.EMACallback
+  #   decay: 0.98
+  #   start_iter: 0
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_sp1_fsdp52_test.yaml

@@ -0,0 +1,105 @@
+# MEMORY/TOPOLOGY TEST: data-free DMD2 on 13 trays (52x GB200, single rack),
+# sp_size=1 (no sequence parallelism — one full packed document per GPU) with
+# FULL FSDP sharding across all 52 ranks (mesh replicate=1 x shard=52).
+#
+# Hypothesis under test: with params/grads/optimizer sharded 52-way (~7 GiB/rank
+# for all three 33B roles) and NVL72 intra-rack NVLink carrying the all-gathers,
+# does the sp=1 activation working set fit in 184 GiB? Prior measurement says
+# ~116 GiB at sp=2 with full checkpointing → ~230 GiB extrapolated at sp=1.
+#
+# Test shape: accum=1 (global batch = 52 = one sample per GPU per step),
+# validation and checkpointing disabled, 200 steps.
+#
+# Launch (13 trays, rack-2):
+#   SP_SIZE=1 HSDP_REPLICATE=1 HSDP_SHARD=52 \
+#   CONFIG=<this file> sbatch --nodes=13 -p hpc-rack-2 \
+#     examples/train/slurm/dmd2_32xgb200.sbatch
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  cfg_uncond:
+    text: zero
+
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 52
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 52
+
+  data:
+    data_path: /mnt/lustre/vlm-wlsaidhi/fastvideo/data/vidprom_h3_text_only
+    preprocessed_data_type: text_only
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 200
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-wlsaidhi/fastvideo/outputs/minimax_h3_dmd2_sp1_fsdp52_test
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 1
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_sp1_fsdp52_test
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_sp1_fsdp40_vidprom.yaml

@@ -0,0 +1,157 @@
+# Production data-free DMD2 on VidProM prompts — sp=1 / full-FSDP topology.
+#
+# v5 (2026-08-14): fresh run folding in the precision + schedule fixes:
+#   * fp32 master weights (bf16 masters bit-froze norm gains for 1000 steps)
+#   * dmd_denoising_steps re-spaced for H3's shift-12 video axis:
+#     [1000, 667, 333] -> sigma_video {1.0, 0.960, 0.857}, audio {1.0, 0.857,
+#     0.60} (the old Wan-derived [1000, 757, 522] gave video {1.0, 0.974,
+#     0.929} — a near-no-op middle step). Same uniform-base-t design rule as
+#     the proven Wan recipe ([1000,750,500,250] + warp at shift 5). Keep
+#     warp_denoising_step OFF: the H3 adapter shifts per modality internally.
+#   * score supervision sampled uniform in sigma_video (score_timestep_shift
+#     12.0) instead of uniform base-t, which put 57% of draws at
+#     sigma_v > 0.9; min_timestep_ratio 0.005 drops the video supervision
+#     floor 0.197 -> 0.057 (audio parity).
+#   * LRs back to 1x — the fp32 fix changes what earlier LR curves meant.
+#
+# Topology: sp_size=1 (one packed document per GPU), mesh (replicate=1,
+# shard=world). num_gpus/hsdp_shard_dim below are defaults; the sbatch
+# overrides them to the allocated world size at launch. 40 GPUs bf16
+# measured ~89 GiB/GPU (job 1770); fp32 masters add ~8-13 GiB/rank
+# depending on shard width — fits with wide margin.
+#
+# Validation: all 64 held-out Wan-derived H3 prompts every 100 steps (+ at
+# start). Checkpoints: every 100 steps starting at step 500, keep last 4
+# (741 GiB each measured at fp32 masters — budget disk accordingly).
+#
+# Launch (after warmup ladder or onto trays held by a previous job):
+#   SP_SIZE=1 HSDP_REPLICATE=1 HSDP_SHARD=<4*nodes> \
+#   CONFIG=<this file> sbatch --nodes=<N> -p hpc-rack-2 \
+#     examples/train/slurm/dmd2_32xgb200.sbatch
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/lustre/vlm-k1kong/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  # 1.0: H3's released checkpoint is guidance-distilled (its inference stack
+  # rejects CFG), so the conditional teacher prediction already carries
+  # guidance — and DMD2 skips the unconditional forward entirely at 1.0.
+  real_score_guidance_scale: 1.0
+  # Base-t ladder; the H3 adapter maps each entry through the per-modality
+  # shifts (video 12 / audio 3). Do NOT enable warp_denoising_step here.
+  dmd_denoising_steps: [1000, 667, 333]
+  # Score-supervision window and density, tuned for shift 12 (see header).
+  min_timestep_ratio: 0.005
+  max_timestep_ratio: 0.98
+  score_timestep_shift: 12.0
+  cfg_uncond:
+    text: zero
+
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 40
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 40
+
+  data:
+    data_path: /mnt/lustre/vlm-wlsaidhi/fastvideo/data/vidprom_h3_text_only
+    preprocessed_data_type: text_only
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: /mnt/lustre/vlm-wlsaidhi/fastvideo/outputs/minimax_h3_dmd2_sp1_vidprom_v5_fp32_sched
+    # Requeue-safe: fresh output_dir -> first boot starts from scratch; after
+    # a pod flap the requeued job resumes from the newest checkpoint here.
+    resume_from_checkpoint: latest
+    training_state_checkpointing_steps: 100
+    checkpointing_start_step: 500
+    checkpoints_total_limit: 4
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_sp1_vidprom_v5_fp32_sched
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  # fp32 sharded master weights (compute stays bf16 via FSDP param_dtype).
+  # bf16 masters silently round away updates below ~half an ulp — norm-scale
+  # params were bit-frozen for 1000 steps (verified against checkpoint-1000).
+  # The trainer hard-checks this at startup.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  latent_vis:
+    _target_: fastvideo.train.callbacks.latent_vis.LatentVisCallback
+    # Decodes rank 0's generator/real-score/fake-score x0 estimates at their
+    # training timesteps — the DMD dynamics view, complementing validation's
+    # full 3-step samples at the same cadence.
+    every_steps: 100
+    fps: 24
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    # 64 held-out prompts (Wan2.1 DMD2 validation_64 rewritten into H3's
+    # three-field format; disjoint from the VidProM training captions).
+    dataset_file: /mnt/lustre/vlm-wlsaidhi/fastvideo/data/h3_validation_wan64.json
+    every_steps: 100
+    run_at_start: true
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_t2va.yaml

@@ -0,0 +1,121 @@
+# DMD2 distillation: MiniMax H3 T2VA (teacher 50-step -> student 3-step).
+#
+# - Teacher: frozen pretrained MiniMax H3
+# - Student: trainable, initialized from the same pretrained weights
+# - Critic:  trainable, initialized from the same pretrained weights
+# - Both modality streams (video + stereo audio) are distilled jointly through
+#   the packed-latent adapter (MiniMaxH3DMDModel): one shared base timestep,
+#   per-modality scheduler shifts (video 12.0, audio 3.0).
+# - Validation: standard MiniMaxH3Pipeline at 3 sampling steps (no dedicated
+#   H3 DMD few-step pipeline yet).
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: TORCH_SDPA
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: TORCH_SDPA
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: TORCH_SDPA
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  # The preprocessed t2va parquet rows carry paired VAE latents, so perturb
+  # those directly; switch to `simulate` for data-free rollout from noise.
+  rollout_mode: data_latent
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  # H3 has no negative-prompt encoder at training time; unconditional teacher
+  # forwards zero the text embeddings instead.
+  cfg_uncond:
+    text: zero
+
+  # Critic optimizer (required — no fallback to training.optimizer)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/crush-smol_h3_t2va_single_sample_preprocessed
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_dmd2_3steps
+    training_state_checkpointing_steps: 20
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_minimax_h3
+    run_name: minimax_h3_dmd2_3steps
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json
+    every_steps: 50
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+  ema:
+    _target_: fastvideo.train.callbacks.ema.EMACallback
+    decay: 0.98
+    start_iter: 0
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml

@@ -0,0 +1,133 @@
+# Single-sample DMD2 overfit: MiniMax H3 with the VSA-H3 student backend at
+# sparsity 0.0 (parity-vs-dense experiment; the backend's sparsity-0 path is
+# parity-proven <=2e-4 vs dense).
+#
+# - Student:  VIDEO_SPARSE_ATTN_H3 backend; DMD2 forwards it attn_kind="vsa"
+#   with real packed-sequence metadata built in MiniMaxH3DMDModel.prepare_batch.
+# - Teacher/critic: FLASH_ATTN (run with FASTVIDEO_FA4=1 for the FA4 path);
+#   their forwards stay dense (attn_metadata=None).
+# - Data: ONE preprocessed t2va parquet row (768x1344, 124 frames @ 24 fps)
+#   produced by fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py
+#   — see examples/distill/MiniMax-H3/overfit_vsa0.sh.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: false
+    disable_custom_init_weights: true
+    attention_backend: FLASH_ATTN
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    disable_custom_init_weights: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: data_latent
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  # H3 has no negative-prompt encoder at training time; unconditional teacher
+  # forwards zero the text embeddings instead.
+  cfg_uncond:
+    text: zero
+
+  # Critic optimizer (required — no fallback to training.optimizer)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    # Preprocessed-latents parquet dir; the cluster run produces it.
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.0 = every tile selected, dense parity).
+  vsa:
+    sparsity: 0.0
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_dmd2_vsa0_overfit
+    # 0 disables training-state checkpointing (no resume; avoids ~400+ GiB
+    # DCP saves on the 87%-full volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: dmd2_vsa0_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0 (costs the validation pipeline
+    # build upfront: text-encoder + VAE loads before the first train step).
+    run_at_start: true
+    sampling_steps: [3]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+  # EMA disabled: its GPU shadow costs ~15.4 GiB/rank, which put the student
+  # update step ~1 GiB over the 184 GiB cards (OOM at step 5/10). Re-enable
+  # once the shadow can live on CPU.
+  # ema:
+  #   _target_: fastvideo.train.callbacks.ema.EMACallback
+  #   decay: 0.98
+  #   start_iter: 0
+
+pipeline: {}

examples/train/configs/distribution_matching/minimax_h3/validation_wan64_h3.json

@@ -0,0 +1,196 @@
+{
+ "data": [
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A close-up frames a woman elegantly showcasing her earrings, raising one hand to touch the intricate silver design with her fingertips. Ambient purple and pink lighting bathes her delicate features, enhancing the vivid tones of her lipstick and eye makeup, and her smoothly styled hair frames her jawline and cheekbones. Behind her, blurred neon lights drift out of focus. The camera pushes in with small amplitude at slow speed toward the earring as she tilts her head. overall_soundscape: Soft rustle of fabric as she raises her arm; a faint electrical hum from the neon signs; quiet room tone. non_diegetic_music: A downtempo electronic groove with soft synth pads and a muted kick pulses at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a lone rider in a classic blue shirt and wide-brimmed hat guiding a chestnut horse across an expansive open field at sunset, horse and rider silhouetted against the warm evening glow. The horse walks gracefully, mane and tail flowing with each step, while towering trees border the field, shadows lengthen across the grass, and a few horses graze near wooden fences by the distant hills. The camera tracks alongside the pair at slow speed. overall_soundscape: Rhythmic hoofbeats on soft earth, leaves rustling in the breeze, a distant whinny, and faint evening birdsong. non_diegetic_music: A gentle acoustic guitar and warm strings play a slow pastoral melody."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A slow push in with small amplitude frames a mysterious pink bottle labeled \"Authentic 100% organic POISON\" standing in the foreground of a dimly lit, eerie room, green fog swirling lightly around its base. Behind it, a shadowy golden bottle adorned with a spider emblem emerges from the gloom, while dim candles flicker faint light across both bottles. overall_soundscape: Low candle-flame flutter, a faint hiss of drifting fog, and hollow room tone with occasional wooden creaks. non_diegetic_music: A sparse minor-key harpsichord line over a low sustained drone at a very slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A static shot frames two large textured tree trunks in the foreground, opening between them onto a calm muddy river shimmering under gentle sunlight in a dense forest. Sunbeams filter through the canopy, casting shifting patterns of light and shadow on the bank, where small bushes and vibrant red and orange foliage transition into deep green leafy plants; the water drifts slowly past. overall_soundscape: A soft river current lapping the bank, layered birdsong from the canopy, and leaves rustling in a light breeze. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a martial artist in a traditional white uniform with a black belt against a stark black background, transitioning gracefully between stances at a deliberate pace. He begins in an open-hand pose with feet firmly grounded, then shifts smoothly into a powerful forward punch, soft lighting creating subtle highlights on the folds of the uniform. The camera holds a static shot as each limb traces a precise, controlled trajectory. overall_soundscape: Crisp fabric snaps with each strike, controlled sharp exhales, and bare feet pivoting on the floor. non_diegetic_music: A solitary taiko drum strikes a slow, spacious rhythm beneath a low string drone."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. An aerial drone view glides over a quiet stretch of sandy shoreline where gentle waves kiss the shore under a clear blue sky, the camera tracking forward at slow speed. Traditional thatched-roof huts nestle among lush palm trees, two white umbrellas stand on the sand, and two small figures walk leisurely along the water's edge leaving faint footprints. overall_soundscape: Rolling waves washing softly onto sand, palm leaves rustling in the sea breeze, and distant gull cries. non_diegetic_music: A breezy ukulele and soft shaker groove at a relaxed tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a lone figure in white sneakers, shorts, and a loose plaid shirt standing on a large moss-covered rock beside a rushing stream, the shirt stirring in the breeze. Behind them a waterfall cascades into the stream, dramatic light carving shadows across the rock's texture and the moving water below. The camera arcs around the figure with small amplitude at slow speed. overall_soundscape: The steady rush of the waterfall, water burbling over stones, and wind moving through nearby foliage. non_diegetic_music: A soft ambient pad with slow piano notes echoes sparsely."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a person leaning casually against a metal railing in an industrial interior, wearing a vibrant patterned jacket over a simple white crop top. Warm ambient light casts soft shadows across concrete walls where intricate wiring and pipes form the backdrop; they hold a relaxed posture and a direct, engaging gaze. The camera pushes in slowly with small amplitude. overall_soundscape: Low ventilation hum, a faint metallic tick from cooling pipes, and the soft scuff of a boot on concrete. non_diegetic_music: A gritty trip-hop beat with a deep bassline plays at a moderate tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a man in a sleeveless black shirt and hooded vest energetically stretching his arms and shoulders in an open-air venue, rows of vibrant red seats curving behind him. The building's geometric glass panels and overlapping metallic beams rise as a dynamic backdrop, with greenery edging the frame; he rolls his shoulders and swings his arms through a warm-up routine. The camera trucks right slowly. overall_soundscape: Fabric rustling with each stretch, measured breathing, birdsong, and a light wind moving through the open venue. non_diegetic_music: An energetic electronic track with a driving four-on-the-floor kick builds steadily."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young woman with loose dark hair, in a green plaid shirt, seated on the floor before a plush beige tufted couch, sorting a stack of papers. She places sheets onto a small round white table, lifting individual pages to examine them, her expression shifting subtly with concentration; nested tables hold documents, a small plant in a gray pot, and a calculator, against a dark paneled wall. The camera holds a static shot. overall_soundscape: Paper sheets sliding and crisply turning, a soft tap as pages meet the table, and quiet room tone. non_diegetic_music: A light lo-fi piano loop with soft vinyl crackle plays at a low volume."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A slow push in frames a heavily rusted metal gate locked by a thick old chain looping around two vertical bars, a worn black iron padlock hanging at its center. The chain's coarse surface shows shades of orange and brown against soft blurred greenery beyond, where tall trees rise and a mossy stone pathway leads away; branches sway gently, their shadows playing across the metalwork. overall_soundscape: Wind moving through distant leaves, a faint metallic creak from the swaying chain, and sparse birdsong. non_diegetic_music: A lone cello plays long, low notes at a very slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames three people mid-session in a softly lit yoga studio, each holding an upward-facing stretch on mats over a warm wooden floor. At center a woman with shoulder-length brown hair in a light cropped top and green leggings holds the pose with grace; to her right a woman in purple mirrors it, and on her left a person with a bun steadies their hands on yoga blocks, all framed by a pastel mural, a gong, and a Buddha statue as natural light floods in from a large window. The camera pans right with small amplitude at slow speed. overall_soundscape: Slow synchronized breathing, mats compressing softly underfoot, and faint birdsong through the window. non_diegetic_music: A meditative drone with soft singing-bowl tones rings slowly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. An underwater wide shot frames two dolphins gliding effortlessly through deep blue water, sunlight filtering down in shafts that scatter shimmering caustics across their sleek skin. They move with fluid grace, exchanging gentle nudges as they curve through the frame. The camera tracks alongside them at slow speed. overall_soundscape: Muffled underwater ambience, soft dolphin clicks and whistles, and the low whoosh of displaced water. non_diegetic_music: A weightless ambient synth pad swells and recedes slowly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young woman with long light-brown hair standing against a vibrant graffiti mural, deeply engrossed in her smartphone, her expression mixing focus with subtle satisfaction. She wears a black floral-patterned top against the mural's teal, orange, and black geometric swirls and pop-art humanoid figures, and like notifications bloom on her screen as she taps. The camera holds a static shot. overall_soundscape: Muted city traffic, soft fingertip taps on glass, gentle notification chimes from the phone, and a passing breeze. non_diegetic_music: An upbeat pop instrumental with bright synth plucks bounces at a medium tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young mother in a sleeveless top and jeans sitting on a bed with her baby, who wears an animal-print onesie amid scattered colorful toys, a plush toy, and a board book. Warm light from a hanging lamp glows over pillows propped against the headboard as she leans in, animatedly holding up a soft white cushion with red and blue accents to catch the baby's curiosity, a small bottle beside her. The camera pushes in slowly with small amplitude. overall_soundscape: The baby's soft coos and giggles, gentle rustling of bedding, and a quiet toy rattle. non_diegetic_music: A tender music-box melody plays slowly over soft strings."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young girl with long brown hair seated at a round wooden table, focused on a laptop whose screen glows a vivid uniform green. To her left a doll in a yellow and white outfit lies atop a stack of books; sheer curtains diffuse soft natural light behind her, and a dark couch blurs in the foreground. The camera holds a static shot as she types and leans toward the screen. overall_soundscape: Light keyboard taps, the laptop fan's faint whir, and muffled daytime ambience through the window. non_diegetic_music: A playful pizzicato string tune skips along at a light tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide aerial shot frames a calm bay where a fleet of sailboats lies anchored in a regimented line toward the horizon, the serene blue water reflecting the soft early-morning sky. A single larger vessel cuts a distinct path, its graceful wake trailing as subtle ripples drift behind the anchored boats; modern high-rise buildings cluster on one shore against the hazy distant coastline. The camera pans left at slow speed. overall_soundscape: Low lapping of water against hulls, distant rigging clinking in the breeze, and far-off gull cries. non_diegetic_music: A calm piano and soft strings drift at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a person standing centered in a pitch-black featureless space, a single spotlight isolating their white gi and black belt in stark contrast. Their expression is serious and focused as their hands transition gracefully through martial arts stances, the crisp fabric folding and catching the light with each precise movement. The camera pushes in with small amplitude at slow speed. overall_soundscape: Controlled breathing, the whisper of fabric through still air, and bare feet gripping the floor. non_diegetic_music: A low sustained drone with sparse struck-metal accents keeps a slow pulse."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a woman in a white hoodie seated in a gaming chair in a dark room washed in neon purple and blue, her large cat-ear headphones glowing softly. Her long hair cascades over her shoulders as her hands fly across a backlit keyboard, eyes locked on the bright screen before her. The camera arcs slowly around her with small amplitude. overall_soundscape: Rapid mechanical keyboard clacks, mouse clicks, the PC fan's low hum, and muffled game audio leaking from headphones. non_diegetic_music: A synthwave track with pulsing arpeggios and a steady electronic beat drives at a medium tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide interior shot frames five riders seated along the worn seats of a dimly lit bus, bathed in red and green overhead light. On the left a woman with curly hair under a patterned scarf sits relaxed in plaid and neon socks; beside her a person in a denim jacket rests their head on a hand deep in thought, and further back a figure in a bucket hat and oversized yellow attire gazes across the aisle. The camera holds a static shot as the bus sways gently. overall_soundscape: The low diesel rumble of the engine, rattling windows over road bumps, and the hiss of air brakes at a stop. non_diegetic_music: A mellow downtempo beat with warm electric piano chords plays softly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. An overhead aerial shot frames two tennis players mid-game on a deep blue court bordered by vibrant green, one in white at the left baseline and one in red at the right, rackets poised. Long late-afternoon shadows stretch across the textured surface as one player serves and the other splits into a ready stance, the dark net dividing crisp white lines. The camera holds a static shot from above. overall_soundscape: The sharp pop of a served ball, shoes squeaking on acrylic, and a bounce echoing across the open court. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, stylized. A mid shot frames a lone dancer moving energetically against deep blue and purple hues, long-exposure streaks smearing their sweeping arms into trails of colored light. Neon glow accentuates their form, bursts of orange and red abstractly illuminating the face as swirling light trails intertwine with the body, patterns rippling across the space with each turn. The camera orbits in an arc shot with small amplitude at slow speed. overall_soundscape: Fabric whipping through air, quick breaths landing between movements, and feet brushing the floor. non_diegetic_music: An intense electronic score with surging synth swells and a pounding kick accelerates."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a photographer in a denim jacket aiming a professional camera at a model in soft pink attire who stands gracefully against a minimalist backdrop, adjusting her long flowing hair with delicate movements. Diffused studio light casts a gentle glow over both as he crouches slightly, refocusing between frames. The camera trucks left slowly with small amplitude. overall_soundscape: Rapid shutter clicks, a strobe's soft recycle whine, quiet footsteps on the studio floor, and fabric rustling. non_diegetic_music: A chic house groove with a smooth bassline plays at a moderate tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a serene expanse of gently sloping hills dotted with tall dense-canopied trees, patches of dry golden grass contrasting the lush green foliage. Sunlight bathes the terrain, revealing a narrow dirt path meandering between the trees while distant hills blur into atmospheric haze. The camera pans right at slow speed. overall_soundscape: A steady breeze moving through leaves, layered distant birdsong, and faint insect chirps in the grass. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a bustling ski slope where skiers in colorful winter gear carve distinct weaving paths down pristine snow, flanked by towering snow-draped evergreens under a clear blue sky. Figures stand silhouetted at the hilltop as a chair lift on the right continuously delivers more skiers, sunlight and shadow accentuating the myriad tracks etched into the slope. The camera tilts down slowly from the hilltop to the descending skiers. overall_soundscape: Skis hissing and edges biting through packed snow, wind rushing past, distant whoops, and the rhythmic clank of the chair lift. non_diegetic_music: A bright orchestral piece with driving strings and brass swells energetically."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide interior shot frames three young riders on a dimly lit bus in muted green and red light. On the left a person with tied-back hair in a jacket and jeans rests their head on a hand; at center another sits upright in a plaid blazer and oversized hoop earrings with intense focus, and on the right a person in a bucket hat and striped shirt leans back, adjusting the hat with a nonchalant gesture. The camera pushes in slowly with small amplitude. overall_soundscape: The engine's low idle and rev, seats creaking with the bus's sway, and a soft chime as a stop is requested. non_diegetic_music: A hazy ambient hip-hop instrumental with a slow boom-bap drum loop simmers quietly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a determined climber in a teal shirt and dark pants scaling a massive sandy-colored rock face, secured by rope and harness, each hold measured and deliberate. The rugged texture of the stone dominates the frame while other large formations and sparse vegetation stretch below a bright overcast sky. The camera tilts up slowly, following the ascent. overall_soundscape: Fingers scraping on gritty rock, carabiners clinking, the rope sliding taut, and wind gusting across the face. non_diegetic_music: A tense percussive score with a ticking rhythm and low strings builds gradually."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a woman in a navy blue jumpsuit standing confidently before a large array of solar panels anchored on a grassy slope, warm sunlight casting gentle shadows across the grid-like structure. Her expression is calm and focused, eyes ahead, as a light breeze stirs the lush green grass at her feet and foliage edges the left of the frame. The camera pulls out slowly with small amplitude to reveal more of the array. overall_soundscape: Wind moving over open grass, a faint electrical hum from the inverters, and distant birdsong. non_diegetic_music: An optimistic minimal piano motif over soft electronic pulses plays at a moderate tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames two colleagues at a wooden desk with an iMac: one in a white knit sweater guides the wireless mouse with the right hand, left hand resting on the sleek white keyboard, while the second sits close by, observing and pointing occasionally at the screen. The monitor displays a tidy array of files and folders, and sticky notes with handwritten reminders cling to the monitor stand. The camera holds a static shot over their shoulders. overall_soundscape: Soft mouse clicks and sparse key taps, a chair creaking lightly, and low office room tone. non_diegetic_music: A neutral corporate lo-fi track with muted keys plays quietly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a man in front of a modern glass facade pulling a dark hoodie up over his head to reveal a gray tank top, arms lifted high in one fluid motion. The lighting highlights the contours of his arms and shoulders while the reflective glass panels behind him double the scene in soft mirror tones. The camera holds a static shot. overall_soundscape: Fabric sliding and snapping free, a soft exhale, and muted street ambience reflecting off the glass. non_diegetic_music: A sparse trap-influenced beat with a deep 808 pulse plays at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, macro. A close-up frames fiery orange flames dancing against a stark black background, tendrils twisting and intertwining into symmetrical swirling patterns that expand and contract rhythmically. Each tongue of fire moves with hypnotic fluidity, its glow casting subtle shifting highlights that deepen the texture of the frame. The camera holds a static shot. overall_soundscape: A low roar and crackle of burning gas, punctuated by soft whooshes as the flames surge. non_diegetic_music: A dark cinematic drone with slow tribal drum hits pulses beneath the flames."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a person seated in a dimly lit wood-paneled recording studio, sunglasses obscuring part of their face, holding several wooden drumsticks in both hands. Their colorful orange-and-blue patterned shirt stands out against the dark background as warm ambient light rakes across the sticks; they tap the sticks together, twirl one between their fingers, and settle into a ready grip. The camera pushes in slowly with small amplitude on the hands. overall_soundscape: The bright click of stick on stick, a soft twirl-and-catch slap into the palm, and the room's dense acoustic hush. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a man lounging on a sofa with a TV remote in one hand, reaching over a table crowded with takeout containers of appetizers with the other, and taking a bite of a sandwich. Colorful patterned cushions and warm ambient lighting fill the background as his attention drifts between the food and the unseen screen. The camera holds a static shot. overall_soundscape: Rustling takeout containers and paper wrappers, an unhurried bite and chew, and the muffled murmur of a TV. non_diegetic_music: A relaxed jazzy sitcom-style cue with brushed drums and upright bass plays softly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a tranquil ocean horizon at dusk, the sky a vibrant tapestry of oranges, pinks, and purples as tall palm trees sway in silhouette at the frame's edges. The shimmering water reflects the sunset while a small boat sits centered on the horizon and gentle waves lap patterns onto the sandy beach in the foreground, scattered clouds catching the last light. The camera pushes in very slowly toward the horizon. overall_soundscape: Waves washing rhythmically onto sand, palm fronds rustling overhead, and a soft offshore breeze. non_diegetic_music: A warm slack-key guitar melody drifts at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a man in a simple gray t-shirt hunched on a couch, head bowed into his hands, elbows on his knees, his face obscured. Gentle light filters through sheer curtains behind him, casting a soft glow that contrasts his static form against the hazy brightness, the muted room emphasizing his stillness as his shoulders rise with one slow breath. The camera pushes in with small amplitude at very slow speed. overall_soundscape: Heavy room silence, a long uneven breath, the couch creaking faintly as his weight shifts. non_diegetic_music: A solo piano plays sparse minor chords at a very slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A close-up frames a pair of hands in a gray T-shirt carefully unrolling a large sheet of crisp white paper across a dark wooden table, the faint beginnings of a colorful printed map emerging as the edges unfurl. Warm light highlights the paper's texture and the wood grain while fingers deftly smooth the sheet flat without creases, a hint of red curtain softening the background. The camera holds a static shot from above the table. overall_soundscape: Paper crackling and sliding over wood, palms smoothing the sheet, and quiet interior room tone. non_diegetic_music: A curious plucked-string motif with light hand percussion ambles at a gentle tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young woman on a vibrant green bus seat, lit by soft pink and blue neon, wearing a neon pink top under a dark-sleeved jacket with jeans and large hoop earrings that catch the light as she shifts. Her gaze drifts thoughtfully to the side while a metallic pole beside her reflects the kaleidoscope of hues against a clean futuristic white interior. The camera holds a static shot as the bus glides. overall_soundscape: A smooth electric hum of the bus, soft cabin rattle, and the hiss of doors at a brief stop. non_diegetic_music: A dreamy synth-pop instrumental with airy pads floats at a medium-slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a silver SUV driving along a winding snow-covered mountain road flanked by dense snow-laden pines, moving smoothly toward a bend where it slips from view; a darker SUV follows onto the snow-dusted asphalt. A towering snow-covered rock formation rises on the right against the deep green of the forest. The camera tracks the vehicles from a high angle at slow speed. overall_soundscape: Tires crunching and hissing over packed snow, engines humming past, and wind sighing through the pines. non_diegetic_music: A spacious ambient score with slow piano and soft strings unfolds."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a vibrant urban skyline at twilight, tall cylindrical industrial structures and equipment grids in the foreground against modern glass high-rises catching the warm hues of the setting sun. A prominent structure labeled \"CITY OF AUSTIN POWER PLANT\" stands out as city lights begin to pierce the dusk, building shadows adding depth and scale. The camera pans right at slow speed across the skyline. overall_soundscape: A low industrial hum, distant traffic wash, and a faint gust across open rooftops. non_diegetic_music: A moody electronic score with deep pads and a slow-blooming synth lead simmers."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames enormous smooth yellow architectural arches towering into the night, soft ambient light gliding along their majestic curves. In the foreground, blurred streaks of headlights and taillights sweep past, their motion contrasting the static monument, while a lone streetlamp and a small tree stand to the left and the clear night sky peeks through the gaps, arch shadows patterning the ground. The camera holds a static long-exposure-style shot. overall_soundscape: A continuous wash of passing traffic, tires on asphalt, and the faint buzz of the streetlamp. non_diegetic_music: A nocturnal downtempo groove with a deep sub-bass pulse plays at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a tranquil marina under golden sunset light, gleaming yachts moored in neat rows, their masts reflected on the gently rippling water. A broad wooden pier juts into the harbor at center while colorful waterfront shops line the left, rolling hillsides and a modern skyline layering the distance as elongated shadows stretch across the docks. The camera pans left at slow speed across the harbor. overall_soundscape: Halyards clinking against masts, water lapping at hulls, gulls calling, and a low murmur of the waterfront. non_diegetic_music: A smooth jazz guitar and soft brushes sway at a relaxed tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a confident person standing atop a structure against a clear blue sky in a striking yellow button-up tied at the waist over a white top, high-waisted white jeans with a black belt, and a bold red scarf at the neck. Yellow-tinted sunglasses reflect the sun as one hand grazes the headphones resting over their ears and the other rests on a railing, hair pulled back neatly, the pose easy and self-assured. The camera arcs around them with small amplitude at slow speed. overall_soundscape: Steady rooftop wind, a distant city murmur far below, and the scarf snapping lightly in the gusts. non_diegetic_music: A feel-good funk groove with a bright bassline and clean rhythm guitar bounces at a medium tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a ballerina spinning across a studio washed in continuous soft pink, her shimmering white tutu and bodice catching the light as she turns. Her arms extend with precise elegance through the pirouette into a poised arabesque, her focused expression sharpening the dreamlike monochrome space. The camera tracks her at slow speed across the floor. overall_soundscape: Pointe shoes brushing and tapping the floor, the whisper of tulle in motion, and controlled breaths between turns. non_diegetic_music: A delicate solo piano waltz turns in slow three-four time."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames two distant figures in hiking gear crossing a serene meadow of tall golden grass swaying in the breeze, the low sun casting a warm, diffused glow. Their silhouettes move deliberately against the lush green of surrounding trees, whose leaves catch the last rays as light and shadow quilt the landscape. The camera holds a static shot as they walk toward the horizon. overall_soundscape: Grass swishing against boots, a soft evening wind, and sparse birdsong settling with the dusk. non_diegetic_music: A folk acoustic guitar fingerpicks a wistful melody at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a large cargo ship docked at an industrial port, its white superstructure rising over a deep green and yellow deck detailed with pipes, railings, and equipment. Calm deep-blue harbor water reflects the vessel's imposing hull in the foreground while storage facilities, green patches, and parked vehicles spread across the paved yard behind. The camera trucks right slowly along the waterline. overall_soundscape: Water slapping the hull, a distant crane motor whirring, gull cries, and a low harbor wind. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a lone climber in a pink shirt and gray pants ascending a towering rock face, harness and chalk bag at the waist, their expression set with determination. Natural pockets and crevices pepper the rugged stone as sunlight casts soft shadows across the cliff, sparse vegetation clinging between distant rocky formations under a clear blue sky. The camera tilts up slowly with the climb. overall_soundscape: Chalked hands slapping stone, boots scraping for purchase, gear jangling on the harness, and wind across the cliff. non_diegetic_music: A patient post-rock build of clean guitar and soft cymbal swells rises slowly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, black-and-white, cinematic. A close-up frames practiced fingers turning the tuning pegs of a guitar headstock, the frame's monochrome palette drawing out the contrast between the strings' sheen and the dark grain of the neck. The player's textured shirt softens the background as the hand makes small, precise adjustments and plucks to test the pitch. The camera holds a static shot. overall_soundscape: Single plucked notes bending gradually into tune, the creak of a peg, fingertips squeaking faintly along a string. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a musician in a dark cap and matching suit playing a large brass instrument, the words \"Brass Band\" clearly visible on its bell, against a vibrant yellow backdrop. He is deeply focused, cheeks working steadily as the instrument's intricate tubing gleams, dramatic lighting carving shadows and metallic highlights while he plays a bold, rolling melody. The camera pushes in slowly with small amplitude. overall_soundscape: The room's close reflections around each sustained brass note and a quiet breath drawn between phrases. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a street musician in a light shirt, dark pants, and flat cap playing an accordion before a grand cathedral, its twin spires and intricate stonework towering behind a lively central fountain. He sways gently in rhythm on the sunlit cobblestones, water jets splashing and catching glimmers of light under a clear azure sky as a warm waltz unspools from the bellows. The camera arcs around him with small amplitude at slow speed. overall_soundscape: Fountain water splashing steadily, pigeons fluttering off the cobbles, and a soft plaza murmur beneath the accordion's melody. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a person seated in a meditative pose on a gentle hillside, back to the camera, silhouetted against a dawning sky of soft oranges and deep purples. Their right hand rests on a knee in a relaxed mudra as the brightening light spills over the lush green landscape, a river reflecting the sky's colors as it meanders past distant city outlines on the left. The camera pushes in at very slow speed. overall_soundscape: A faint morning breeze in the grass, early birdsong waking across the valley, and the distant murmur of the river. non_diegetic_music: A soft ambient drone with slow bansuri flute phrases floats over the sunrise."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide panoramic shot frames a vast valley blanketed in gentle mist, undulating hills lush with dense greenery bordering the left of the frame. The mist weaves through the landscape like a soft ethereal blanket while dark mountain peaks emerge in the distance against a pale blue sky, wispy clouds drifting as sunlight filters warmly through the haze. The camera pans right at very slow speed. overall_soundscape: A hushed highland wind, faint birdsong carrying across the valley, and leaves stirring in the nearer trees. non_diegetic_music: A slow cinematic pad with distant French horn swells breathes with the mist."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a bearded man in a dark buttoned shirt holding a smartphone at an elevated vantage point, the setting sun flaring golden across the panoramic cityscape behind him. His hands cradle the phone delicately but purposefully, eyes engaged on the screen as the sunlight glimmers off its surface in a striking lens flare. The camera holds a static shot. overall_soundscape: A soft rooftop breeze, faint city hum rising from below, and a fingertip tapping the glass. non_diegetic_music: A reflective downtempo beat with warm Rhodes chords plays at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a solitary performer in dark fitted clothing in a vast industrial hall of towering columns, one leg bent forward, arms outstretched in a horizontal arc with intense flames tracing along both arms. The fire forms a brilliant orange halo, flickering shadows across weathered walls and worn tile as dim daylight filters through lining windows, the dancer holding balance amid the heat. The camera arcs around the performer with small amplitude at slow speed. overall_soundscape: The whoosh and roar of the flames as the arms sweep, fuel hissing softly, and footfalls echoing in the empty hall. non_diegetic_music: A slow tribal percussion pattern with a dark drone deepens the ritual pace."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A close-up frames a man in a dark shirt, a subtle necklace at his chest, playing the violin with his eyes narrowed in concentration and passion. His fingers dance along the strings as the bow moves smoothly, warm room light reflecting off the instrument's rich varnish while the background falls into a soft blur, a soaring melody pouring from the strings. The camera pushes in with small amplitude at slow speed. overall_soundscape: The faint brush of the bow between phrases and a quiet intake of breath at each swell beneath the violin's melody. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames two players in an impromptu soccer duel inside a dimly lit parking garage, neon light reflecting off patches of standing water. The first, in a light grey shirt and black three-striped pants, maneuvers the ball with precise footwork while a second figure in dark clothing shifts, poised to intercept between stark yellow floor lines and orange bollards; the ball glides across the wet concrete, kicking up droplets. The camera tracks the ball at slow speed. overall_soundscape: The ball skimming and splashing over wet concrete, shoe squeaks and quick steps echoing off low ceilings, and a metallic reverb after each kick. non_diegetic_music: A gritty electronic beat with syncopated percussion drives at a medium tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a lone climber in a light-colored shirt and jeans gripping a towering vertical rock face, navigating rough textures and crevices with precision. The sheer cliff spans hues from light tan to deep gray, dwarfing the figure while scattered greenery and rugged terrain spread far below in wild isolation. The camera pedestals up slowly alongside the ascent. overall_soundscape: Deliberate breaths, fingertips finding gritty holds, a scatter of pebbles falling away, and wind sweeping the rock face. non_diegetic_music: A sparse tension score of ticking percussion and low sustained strings holds its breath."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a lush meadow of vibrant yellow wildflowers swaying gently in the breeze, a towering broad-branched tree standing majestically at the right beneath a bright blue sky of fluffy clouds. On the left, dense trees form a natural corridor toward the horizon, green grass contrasting the golden distant fields as light and shadow ripple with the wind. The camera pushes in at slow speed through the flowers. overall_soundscape: A soft breeze combing through grass and petals, bees humming close by, and easy afternoon birdsong. non_diegetic_music: A pastoral flute and acoustic guitar duet lilts at a gentle tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a solitary boat gliding across a mirror-still teal lake, its gentle wake spreading delicate ripples that blend into a calm gradient sky of soft orange near the shore and pale blue above. Faint silhouettes of distant trees line the horizon while a few slender poles emerge from the water, remnants of an old structure, as the boat presses steadily toward the open beyond. The camera holds a static shot. overall_soundscape: A small outboard motor puttering softly, water parting at the bow, and early birds calling across the lake. non_diegetic_music: A minimal piano theme with long sustained notes drifts at a very slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a woman with dark wavy hair in a denim jacket seated alone at a rustic wooden table in a cozy dim caf\u00e9, leaning over an open book in deep focus. Hanging pendant lights cast a warm glow over shelves lined with jars and coffee gear behind her, a small cup of coffee within reach beside a glass dome sheltering a solitary pastry; she turns a page slowly. The camera pushes in with small amplitude at slow speed. overall_soundscape: A page turning softly, the espresso machine hissing in the background, cups clinking, and a low caf\u00e9 murmur. non_diegetic_music: A cozy bossa nova guitar pattern sways quietly."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a solitary musician at a small keyboard rig in a vast deserted landscape under the night sky, strategically placed lights washing the setup in orange and blue. A large circular light glows to the left as their hands move across the keys among cabled electronic equipment, the immense darkness pressing in around the intimate performance island; layered synth chords ring out into the open air. The camera pulls out slowly with small amplitude. overall_soundscape: A soft desert wind crossing open ground and the faint electrical buzz of the rig beneath the synthesizer's chords. non_diegetic_music: N/A"
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames cars streaking past a weathered building with a faded yellow and brown graffiti-covered facade, their motion blurred against the static urban grit. A solitary figure in a gray top and mustard trousers stands to the side gazing into the street, and a bright yellow taxi slows as it approaches them under a soft overcast sky, a vertical pole cutting the frame. The camera holds a static shot. overall_soundscape: A continuous wash of passing traffic, an engine decelerating, a distant horn, and the ambient chatter of the street. non_diegetic_music: A muted jazz trumpet over a shuffling brushed-drum rhythm plays at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a young woman seated on a curb in a tranquil park, a collie resting calmly beside her, both bathed in the golden hue of the setting sun that ignites the dog's fur with a serene glow. Her hand strokes gently along the collie's back as tall trees cast elongated shadows across the leaf-strewn ground. The camera pushes in at very slow speed. overall_soundscape: Leaves rustling overhead, the dog's relaxed breathing and a contented huff, and distant park ambience. non_diegetic_music: A warm acoustic guitar and soft strings play a tender, unhurried melody."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A wide shot frames a grand elephant standing in an open sunlit field, its massive form dominating the frame, skin a tapestry of earthy tones and rough textured wrinkles dusted with dirt. Its powerful trunk sways gently as light plays across patches of dust on its hide, the background a soft blur of trees and shrubs. The camera arcs slowly around the animal with small amplitude. overall_soundscape: A deep rumbling exhale, the trunk brushing dry grass, insects droning in the heat, and distant savanna birds. non_diegetic_music: A stately orchestral theme with low strings and soft horns moves at a slow tempo."
+  },
+  {
+   "caption": "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A mid shot frames a fluffy dog with brown patches lying relaxed on a plain white background, intently nudging and playfully biting a bright red fire-hydrant-shaped toy with a yellow and orange rope attached. Its ears perk with curiosity and its eyes stay fixed on the toy as a paw pins the rope. The camera holds a static shot. overall_soundscape: Soft squeaks from the toy, the rope dragging across the floor, an excited snuffle, and a light tail thump. non_diegetic_music: A bouncy ukulele and glockenspiel jingle plays at a playful tempo."
+  }
+ ]
+}

examples/train/configs/fine_tuning/minimax_h3/README.md

@@ -0,0 +1,13 @@
+# MiniMax-H3 SFT overfit configs
+
+Single- and two-sample SFT overfit runs used to validate the VSA-H3 backend,
+sparsity sweep (0 / 0.8 / 0.9 / 0.95 / 0.97), the FA4 dense control, and
+effective-batch-2 via gradient accumulation (`sft_fa4_bs2_overfit.yaml`).
+
+Launch, topology, batch semantics, measured step times, and the weight shard
+cache are documented in
+[`../../distribution_matching/minimax_h3/README.md`](../../distribution_matching/minimax_h3/README.md).
+
+Note: `sft_vsa95_overfit_2gpu.yaml` / `sft_vsa97_overfit_2gpu.yaml` are kept
+as negative results — 2x GB200 cannot fit this model/sequence without CPU
+offload (backward-pass working set alone OOMs 184 GiB cards at sp=2).

examples/train/configs/fine_tuning/minimax_h3/sft_fa4_bs2_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# FA4 dense control (student on FLASH_ATTN, run with FASTVIDEO_FA4=1) for
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data_bs2
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1  # per-micro-batch; effective batch 2 via grad accumulation
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.0 = every tile selected, dense parity).
+  vsa:
+    sparsity: 0.0
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 2
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_fa4_bs2_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_fa4_bs2_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_bs2_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_fa4_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# FA4 dense control (student on FLASH_ATTN, run with FASTVIDEO_FA4=1) for
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: FLASH_ATTN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.0 = every tile selected, dense parity).
+  vsa:
+    sparsity: 0.0
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_fa4_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_fa4_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa0_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.0 — the fine-tuning counterpart of
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.0 = every tile selected, dense parity).
+  vsa:
+    sparsity: 0.0
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa0_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa0_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa80_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.8 (80% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.8 = top-20% of tiles attended).
+  vsa:
+    sparsity: 0.8
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa80_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa80_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa90_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.9 (90% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.9 = top-10% of tiles attended).
+  vsa:
+    sparsity: 0.9
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa90_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa90_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa95_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.95 (95% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.9 = top-10% of tiles attended).
+  vsa:
+    sparsity: 0.95
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa95_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa95_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa95_overfit_2gpu.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.95 (95% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 2
+    sp_size: 2
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 2
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.95 = top-5% of tiles attended).
+  vsa:
+    sparsity: 0.95
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.0, 0.999]  # beta1=0 -> buffer-free AdamW; std betas OOM at sp=2
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa95_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa95_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa97_overfit.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.97 (97% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.9 = top-10% of tiles attended).
+  vsa:
+    sparsity: 0.97
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa97_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa97_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/configs/fine_tuning/minimax_h3/sft_vsa97_overfit_2gpu.yaml

@@ -0,0 +1,93 @@
+# Single-sample SFT overfit: MiniMax H3 with the VSA-H3 backend at sparsity
+# 0.97 (97% of video tiles pruned) — sparse variant of sft_vsa0_overfit.yaml;
+# distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml (same c17_v2
+# preprocessed record, same validation prompt). Optimizer recipe follows the
+# crush-smol H3 T2VA overfit (overfit_minimax_h3_t2va.yaml), scaled to one
+# 4-GPU node with pure-FSDP sharding.
+
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3Model
+    init_from: /mnt/models/MiniMax-H3
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: VIDEO_SPARSE_ATTN_H3
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 2
+    sp_size: 2
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 2
+
+  data:
+    data_path: /mnt/h3-dmd2-overfit/data
+    preprocessed_data_type: t2va
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 37
+    num_height: 768
+    num_width: 1344
+    num_frames: 124
+
+  # Student VSA sparsity (0.97 = top-3% of tiles attended).
+  vsa:
+    sparsity: 0.97
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.0, 0.999]  # beta1=0 -> buffer-free AdamW; std betas OOM at sp=2
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/minimax_h3_sft_vsa97_overfit
+    # 0 disables training-state checkpointing (no resume; avoids large DCP
+    # saves on the shared volume).
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: h3-dmd2-vsa
+    run_name: sft_vsa97_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: bf16
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline
+    dataset_file: /mnt/h3-dmd2-overfit/c17_validation.json
+    every_steps: 50
+    # Sample the untrained baseline at step 0.
+    run_at_start: true
+    # SFT is not a few-step generator; sample on the full 50-step schedule
+    # (mirrors the crush-smol overfit validation).
+    sampling_steps: [50]
+    guidance_scale: 1.0
+    num_frames: 124
+    num_videos_per_prompt: 1
+    use_validation_media_conditioning: false
+    offload_training_state: true
+    text_encoder_cpu_offload: true
+    vae_cpu_offload: true
+
+pipeline: {}

examples/train/slurm/dmd2_32xgb200.sbatch

@@ -0,0 +1,115 @@
+#!/bin/bash
+# 8-tray (8 x 4x GB200 = 32 GPU) MiniMax-H3 DMD2 launch from a Slurm login node.
+#
+#   sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+#   sbatch --nodes=2 examples/train/slurm/dmd2_32xgb200.sbatch          # smaller sweep
+#   CONFIG=path/to/other.yaml sbatch examples/train/slurm/dmd2_32xgb200.sbatch
+#
+# The paired config (dmd2_simulate_fa4_32gpu.yaml) assumes 8 nodes:
+# hsdp_replicate_dim must equal the node count and num_gpus the world size —
+# this script forwards overrides derived from the actual allocation, so
+# shrinking --nodes "just works" (global batch shrinks with it).
+#
+# Cluster-specific knobs (partition/account/container) are marked CLUSTER.
+
+#SBATCH --job-name=h3-dmd2-32xgb200
+#SBATCH --nodes=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gpus-per-node=4
+#SBATCH --exclusive
+#SBATCH --time=48:00:00
+#SBATCH --output=slurm-%x-%j.out
+##SBATCH --partition=<CLUSTER: partition>
+##SBATCH --account=<CLUSTER: account>
+
+set -euo pipefail
+
+# ── paths (shared FS visible on every tray) ──────────────────────────────
+REPO="${REPO:-/mnt/FastVideo}"
+VENV="${VENV:-${REPO}/.venv}"
+CONFIG="${CONFIG:-${REPO}/examples/train/configs/distribution_matching/minimax_h3/dmd2_simulate_fa4_32gpu.yaml}"
+
+# ── topology derived from the allocation ─────────────────────────────────
+NNODES="${SLURM_JOB_NUM_NODES}"
+GPUS_PER_NODE="${SLURM_GPUS_PER_NODE:-4}"
+GPUS_PER_NODE="${GPUS_PER_NODE##*:}"   # "4" or "gb200:4" -> 4
+WORLD_SIZE=$((NNODES * GPUS_PER_NODE))
+# Topology overrides (default: sp across a tray, HSDP replicate=nodes/shard=4).
+# e.g. SP_SIZE=1 HSDP_REPLICATE=1 HSDP_SHARD=$WORLD_SIZE for pure full-shard FSDP.
+SP_SIZE="${SP_SIZE:-${GPUS_PER_NODE}}"
+HSDP_REPLICATE="${HSDP_REPLICATE:-${NNODES}}"
+HSDP_SHARD="${HSDP_SHARD:-${GPUS_PER_NODE}}"
+MASTER_ADDR="$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1)"
+MASTER_PORT="${MASTER_PORT:-29501}"
+
+echo "=== H3 DMD2 multi-node launch ==="
+echo "Nodes:        ${NNODES} (${SLURM_JOB_NODELIST})"
+echo "World size:   ${WORLD_SIZE}"
+echo "Topology:     sp=${SP_SIZE} hsdp=(replicate=${HSDP_REPLICATE}, shard=${HSDP_SHARD})"
+echo "Rendezvous:   ${MASTER_ADDR}:${MASTER_PORT}"
+echo "Config:       ${CONFIG}"
+
+# ── per-rank environment (srun propagates exported vars) ─────────────────
+# CLUSTER: if compute nodes lack the login /home (k8s worker pods), submit
+# with LUSTRE_HOME=<shared writable dir> — triton's JIT cache (touched at
+# fastvideo import), HF caches, and ~/.netrc all resolve through $HOME.
+export HOME="${LUSTRE_HOME:-${HOME}}"
+export FASTVIDEO_FA4=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export TOKENIZERS_PARALLELISM=false
+# Node-local weight cache: relaunches rebuild all three 33B roles from tmpfs
+# in seconds. PER_NODE=1 makes every tray write its own /dev/shm copy
+# (single-writer semantics would leave 7 of 8 trays permanently cold).
+export FASTVIDEO_WEIGHT_SHARD_CACHE=/dev/shm/fastvideo-wcache
+export FASTVIDEO_WEIGHT_SHARD_CACHE_PER_NODE=1
+# W&B: credentials come from ~/.netrc on the shared home; flip to offline if
+# compute nodes have no egress.
+export WANDB_MODE="${WANDB_MODE:-online}"
+# CLUSTER: NCCL fabric tuning. These three come from the cluster onboarding
+# guide's multi-node NCCL-test recipe (GB200 NVL72 + CX7 IB cross-rack):
+# MNNVL is multi-node NVLink within a rack's NVL72 domain.
+export NCCL_CUMEM_ENABLE="${NCCL_CUMEM_ENABLE:-1}"
+export NCCL_MNNVL_ENABLE="${NCCL_MNNVL_ENABLE:-2}"
+export UCX_NET_DEVICES="${UCX_NET_DEVICES:-eth0}"
+# WARN costs nothing and turns silent transport hangs at communicator init
+# into named errors in the log.
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+
+export NNODES MASTER_ADDR MASTER_PORT
+export NUM_GPUS="${GPUS_PER_NODE}"
+
+# Boot watchdog: NCCL communicator init can wedge silently (all ranks parked
+# in ib_uverbs_event_read/poll, no WARN output, GPUs at 0%). If no weight
+# loading has been logged within BOOT_TIMEOUT_MIN minutes, abort the job so
+# the allocation isn't squatted on. 0 disables.
+BOOT_TIMEOUT_MIN="${BOOT_TIMEOUT_MIN:-20}"
+if [ "${BOOT_TIMEOUT_MIN}" -gt 0 ]; then
+  (
+    sleep $((BOOT_TIMEOUT_MIN * 60))
+    if ! grep -qs "Loading weights took" "${REPO}/examples/train/logs/${SLURM_JOB_ID}-node0"/*.log; then
+      echo "BOOT WATCHDOG: no weight loading after ${BOOT_TIMEOUT_MIN} min — aborting job ${SLURM_JOB_ID}"
+      scancel "${SLURM_JOB_ID}"
+    fi
+  ) &
+  WATCHDOG_PID=$!
+fi
+
+# CLUSTER: if trays require a container, wrap the srun payload accordingly,
+# e.g. add `--container-image=... --container-mounts=/mnt:/mnt` (pyxis) so
+# /mnt (repo, models, data) and /dev/shm are available inside.
+srun --kill-on-bad-exit=1 --ntasks-per-node=1 bash -c "
+  set -euo pipefail
+  export NODE_RANK=\${SLURM_NODEID}
+  # One log dir per node — parallel tee writers on one NFS file garble it.
+  export LOG_DIR=examples/train/logs/\${SLURM_JOB_ID}-node\${SLURM_NODEID}
+  source '${VENV}/bin/activate'
+  cd '${REPO}'
+  bash examples/train/run.sh '${CONFIG}' \
+    --training.distributed.num_gpus ${WORLD_SIZE} \
+    --training.distributed.hsdp_replicate_dim ${HSDP_REPLICATE} \
+    --training.distributed.hsdp_shard_dim ${HSDP_SHARD} \
+    --training.distributed.sp_size ${SP_SIZE}
+"
+SRUN_EXIT=$?
+[ -n "${WATCHDOG_PID:-}" ] && kill "${WATCHDOG_PID}" 2>/dev/null
+exit "${SRUN_EXIT}"

examples/train/slurm/encode_h3_text_only_1xgb200.sbatch

@@ -0,0 +1,78 @@
+#!/bin/bash
+# Encode the VidProM-H3 prompt list into text-only conditioning parquet on one
+# GB200 tray: 4 GPUs = 4 independent shard processes of
+# preprocess_minimax_h3_text_only.py (each loads the full Qwen3-VL encoder).
+#
+#   sbatch examples/train/slurm/encode_h3_text_only_1xgb200.sbatch
+#
+# Shards are resume-safe: resubmitting after a timeout/preemption continues
+# after the rows already on disk. Paths assume the nvl-gb200 cluster layout
+# (login /home is NOT on compute nodes — everything lives on /mnt/lustre);
+# override via the environment, e.g. PROMPTS=... MODEL=... sbatch ...
+#
+# On this autoscaled cluster a PENDing job is what creates node demand — leave
+# the job queued and the slurm-operator scales a tray up.
+
+#SBATCH --job-name=h3-textenc
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:4
+#SBATCH --exclusive
+#SBATCH --partition=all
+#SBATCH --time=08:00:00
+#SBATCH --output=/mnt/lustre/vlm-wlsaidhi/fastvideo/logs/h3-textenc-%j.out
+
+set -uo pipefail
+
+# Compute pods have no /home/<user>: triton's JIT cache (touched at fastvideo
+# import), HF caches, and W&B netrc all resolve through $HOME — point it at
+# lustre or every fastvideo import dies with PermissionError.
+export HOME="${LUSTRE_HOME:-/mnt/lustre/vlm-wlsaidhi}"
+
+REPO="${REPO:-/mnt/lustre/vlm-wlsaidhi/fastvideo/FastVideo}"
+VENV="${VENV:-${REPO}/.venv}"
+PROMPTS="${PROMPTS:-/mnt/lustre/vlm-wlsaidhi/fastvideo/data/vidprom_h3_v1_63136_20260811.txt}"
+MODEL="${MODEL:-/mnt/lustre/vlm-k1kong/models/MiniMax-H3}"
+OUT="${OUT:-/mnt/lustre/vlm-wlsaidhi/fastvideo/data/vidprom_h3_text_only}"
+NUM_SHARDS="${NUM_SHARDS:-4}"
+LIMIT="${LIMIT:-}"   # e.g. LIMIT=8 for a smoke run
+
+mkdir -p "$OUT"
+
+echo "node=$(hostname) repo=$REPO prompts=$PROMPTS model=$MODEL out=$OUT shards=$NUM_SHARDS limit=${LIMIT:-none}"
+cd "$REPO"
+
+pids=()
+for k in $(seq 0 $((NUM_SHARDS - 1))); do
+  CUDA_VISIBLE_DEVICES=$k MASTER_PORT=$((29531 + k)) \
+    "$VENV/bin/python" fastvideo/pipelines/preprocess/preprocess_minimax_h3_text_only.py \
+      --prompts-file "$PROMPTS" \
+      --model-path "$MODEL" \
+      --output-dir "$OUT" \
+      --shard-index "$k" \
+      --num-shards "$NUM_SHARDS" \
+      ${LIMIT:+--limit "$LIMIT"} \
+      > "$OUT/shard_${k}.log" 2>&1 &
+  pids+=($!)
+done
+
+fail=0
+for k in "${!pids[@]}"; do
+  if ! wait "${pids[$k]}"; then
+    echo "shard $k FAILED (see $OUT/shard_${k}.log)"
+    fail=1
+  fi
+done
+
+# The training dataloader caches the parquet file list under the dataset root;
+# drop any stale cache so the next training start rescans fresh shards.
+rm -rf "$OUT/map_style_cache"
+
+"$VENV/bin/python" - <<EOF
+import pyarrow.parquet as pq
+from pathlib import Path
+files = sorted(Path("$OUT").rglob("*.parquet"))
+rows = sum(pq.ParquetFile(f).metadata.num_rows for f in files)
+print(f"TOTAL: {rows} rows across {len(files)} parquet files under $OUT")
+EOF
+
+exit $fail

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py

@@ -237,7 +237,12 @@ def _attn_bwd_dkdv(
         # Load m before computing qk to reduce pipeline stall.
         offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
         m = tl.load(M + offs_m)
-        qkT = tl.dot(k, qT)
+        # Recompute logits exactly as the forward does: raw bf16 operands into
+        # the dot, fp32 scale after accumulation. A bf16 pre-scaled K perturbs
+        # the recomputed logits relative to the saved M by an error
+        # proportional to |logit|, which exp2 amplifies into arbitrarily wrong
+        # probabilities at large activations.
+        qkT = tl.dot(k, qT) * (sm_scale * 1.4426950408889634)
         pT = tl.math.exp2(qkT - m[None, :])
         mask = tl.arange(0, BLOCK_N1) < block_size
         pT = tl.where(mask[:, None], pT, 0.0)
@@ -268,6 +273,7 @@ def _attn_bwd_dq(
         do,
         m,
         D,
+        sm_scale,
         # shared by Q/K/V/DO.
         q2k_index,
         q2k_num,
@@ -315,7 +321,7 @@ def _attn_bwd_dq(
         block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
-        qk = tl.dot(q, kT)
+        qk = tl.dot(q, kT) * (sm_scale * 1.4426950408889634)
         p = tl.math.exp2(qk - m)
         offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
         mask = offs_in_block < block_size
@@ -324,8 +330,7 @@ def _attn_bwd_dq(
         dp = tl.dot(do, vT).to(tl.float32)
         ds = p * (dp - Di[:, None])
         ds = ds.to(tl.bfloat16)
-        # Compute dQ.
-        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        # Compute dQ (kT is raw; the caller applies sm_scale once at the end).
         dq += tl.dot(ds, tl.trans(kT))
         # Increment pointers.
     return dq
@@ -453,6 +458,7 @@ def _attn_bwd(
         do,
         m,
         D,  #
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -470,7 +476,7 @@ def _attn_bwd(
     )
     # Write back dQ.
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq *= LN2
+    dq *= sm_scale
     tl.store(dq_ptrs, dq)
 
 
@@ -591,6 +597,7 @@ def _attn_bwd_dq_kernel(
         Q,
         K,
         V,
+        sm_scale,
         DO,  #
         DQ,
         M,
@@ -663,6 +670,7 @@ def _attn_bwd_dq_kernel(
         do,
         m,
         D,
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -680,7 +688,7 @@ def _attn_bwd_dq_kernel(
     )
 
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
+    dq_acc *= sm_scale
     tl.store(dq_ptrs, dq_acc)
 
 
@@ -748,9 +756,11 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
     dv = torch.empty_like(v)
     BATCH, N_HEAD = q.shape[:2]
     BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
-    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+    # K stays raw: the backward kernels apply sm_scale in fp32 after the dot,
+    # matching the forward's rounding exactly. (A bf16 pre-scaled K perturbs
+    # the recomputed logits vs the saved M; exp2 turns that into unboundedly
+    # wrong probabilities at large activations.)
     arg_k = k
-    arg_k = arg_k * (sm_scale * RCP_LN2)
     PRE_BLOCK = 64
     assert Tq % PRE_BLOCK == 0
     pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
@@ -813,6 +823,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
         q,
         arg_k,
         v,
+        sm_scale,
         do,
         dq,
         M,

fastvideo/attention/backends/abstract.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
 
+import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field, fields
 from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
@@ -10,6 +11,22 @@ if TYPE_CHECKING:
 
 import torch
 
+_LAYER_IDX_RE = re.compile(r"blocks\.(\d+)")
+
+
+def layer_idx_from_prefix(prefix: str, default: int | None = None) -> int:
+    """Parse the transformer-block index out of a layer prefix.
+
+    Shared by backends that key per-layer behavior off the block number
+    (vmoba, VSA-H3). Raises when unparsable unless ``default`` is given.
+    """
+    match = _LAYER_IDX_RE.search(prefix)
+    if match:
+        return int(match.group(1))
+    if default is None:
+        raise ValueError(f"cannot parse a layer index from attention prefix {prefix!r}")
+    return default
+
 
 class AttentionBackend(ABC):
     """Abstract class for attention backends."""

fastvideo/attention/backends/video_sparse_attn.py

@@ -61,6 +61,7 @@ def construct_variable_block_sizes(
     dit_seq_shape: tuple[int, int, int],
     num_tiles: tuple[int, int, int],
     device: torch.device,
+    tile_size: tuple[int, int, int] = VSA_TILE_SIZE,
 ) -> torch.LongTensor:
     """
     Compute the number of valid (non‑padded) tokens inside every
@@ -73,7 +74,7 @@ def construct_variable_block_sizes(
     """
     # unpack
     t, h, w = dit_seq_shape
-    ts_t, ts_h, ts_w = VSA_TILE_SIZE
+    ts_t, ts_h, ts_w = tile_size
     n_t, n_h, n_w = num_tiles
 
     def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
@@ -157,10 +158,35 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     cache_tile_buf: bool = True
 
 
+def compute_topk(sparsity: float, num_blocks: int) -> int:
+    """Blocks to keep for a sparsity level, clamped to [1, num_blocks]."""
+    return max(1, min(math.ceil((1 - sparsity) * num_blocks), num_blocks))
+
+
 def _compute_cur_topk(attn_metadata: VideoSparseAttentionMetadata) -> int:
-    num_kv_blocks = attn_metadata.variable_block_sizes.numel()
-    cur_topk = math.ceil((1 - attn_metadata.VSA_sparsity) * num_kv_blocks)
-    return max(1, min(cur_topk, num_kv_blocks))
+    return compute_topk(attn_metadata.VSA_sparsity, attn_metadata.variable_block_sizes.numel())
+
+
+def scatter_into_tile_buf(
+    x: torch.Tensor,
+    target_shape: tuple[int, ...],
+    dst_index: torch.Tensor,
+    buf: torch.Tensor | None,
+    src_index: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Zero-padded tile scatter shared by the VSA backends.
+
+    Allocates (zeros) when ``buf`` is missing or mismatched; otherwise reuses
+    it — pad slots are never written and every non-pad slot is fully
+    overwritten per call, so a reused buffer stays valid. Callers own the
+    buffer's lifetime (per-metadata for Wan VSA, per-builder for VSA-H3) and
+    its aliasing contract: the result is only valid until the next call with
+    the same buffer.
+    """
+    if (buf is None or buf.shape != target_shape or buf.dtype != x.dtype or buf.device != x.device):
+        buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
+    buf[:, dst_index] = x if src_index is None else x[:, src_index]
+    return buf
 
 
 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
@@ -243,23 +269,15 @@ class VideoSparseAttentionImpl(AttentionImpl):
         target_shape = (x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1])
 
         if not attn_metadata.cache_tile_buf:
-            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
-            buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
-            return buf
+            return scatter_into_tile_buf(x, target_shape, attn_metadata.non_pad_index, None,
+                                         attn_metadata.tile_partition_indices)
 
-        # Reuse the per-step buffer stashed on metadata (lazily allocated
-        # on the first VSA layer's call within a denoising step).  Pad
-        # positions are zero from the initial torch.zeros and never
-        # written to.  Scoping to metadata makes reuse safe across
-        # concurrent requests and keeps the "pad positions are zero"
-        # invariant trivially true: ``non_pad_index`` is fixed within
-        # a single metadata instance.
-        buf = attn_metadata.tile_buf
-        if (buf is None or buf.shape != target_shape or buf.dtype != x.dtype or buf.device != x.device):
-            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
-            attn_metadata.tile_buf = buf
-
-        buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
+        # Buffer scoped to the per-step metadata (lazily allocated on the
+        # first VSA layer's call within a denoising step), which keeps reuse
+        # safe across concurrent requests.
+        buf = scatter_into_tile_buf(x, target_shape, attn_metadata.non_pad_index, attn_metadata.tile_buf,
+                                    attn_metadata.tile_partition_indices)
+        attn_metadata.tile_buf = buf
         return buf
 
     def untile(self, x: torch.Tensor, untile_combined_index: torch.LongTensor) -> torch.Tensor:

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+"""VSA for MiniMax H3's packed mixed-modality self-attention.
+
+H3 runs one joint bidirectional attention over
+``[text | condition keyframes | audio | generated video]``, so this
+backend differs from the Wan-tuned ``video_sparse_attn``:
+
+- Tiles are ``[segment-pure prefix chunks] + [3D (4,8,8) video tiles]``;
+  prefix tiles never straddle segment boundaries.
+- Selection is pure Python on pooled tile scores; the block-sparse kernel
+  consumes an explicit bool mask, so no kernel changes are needed.
+- The compression branch is gated by ``to_gate_compress``, which the H3
+  checkpoint does not carry: the loader zero-initializes it, so untrained
+  inference is exactly pure sparse and finetuning can learn the gate.
+- Non-video *queries* are always dense. Non-video *keys* are either
+  always-selected for every query ("exempt", default) or compete in
+  top-k under a FLOP-matched budget ("compete") — the ablation axis,
+  switched per request via ``generate_video(..., vsa_mode=...)``
+  (default: exempt). Per-request scheduling knobs
+  (``vsa_dense_first_n_steps``, ``vsa_dense_layers``) let mixed schedules
+  run the diffuse steps/layers dense while pushing the rest harder.
+
+Targets sm10.x through the FA4 CuTe 256-tile path
+(``FASTVIDEO_VSA_CUTEDSL=1``); the Triton 256→64 expansion is the
+fallback and keeps identical mask semantics.
+"""
+
+import functools
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+try:
+    from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_256_bshd
+except ImportError:
+    block_sparse_attn_256_bshd = None
+
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
+                                                   AttentionMetadataBuilder, layer_idx_from_prefix)
+from fastvideo.attention.backends.video_sparse_attn import (compute_topk, construct_variable_block_sizes,
+                                                            get_non_pad_index, get_tile_partition_indices,
+                                                            scatter_into_tile_buf)
+from fastvideo.attention.backends.video_sparse_attn_h3_probe import probe_enabled, record_probe
+from fastvideo.logger import init_logger
+
+VSA_H3_TILE_SIZE = (4, 8, 8)  # 256 elements -> FA4 CuTe fastpath on sm10.x
+_TILE_ELEMS = math.prod(VSA_H3_TILE_SIZE)
+
+logger = init_logger(__name__)
+
+
+def token_tile_and_valid(variable_block_sizes: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per padded-token tile id and pad-validity mask.
+
+    The single encoding of the padding contract, shared by the probe and the
+    test oracle so they cannot drift from the backend's tile geometry.
+    """
+    device = variable_block_sizes.device
+    token_tile = torch.arange(variable_block_sizes.numel(), device=device).repeat_interleave(_TILE_ELEMS)
+    token_valid = (torch.arange(_TILE_ELEMS, device=device)[None, :] < variable_block_sizes[:, None]).reshape(-1)
+    return token_tile, token_valid
+
+
+@functools.lru_cache(maxsize=10)
+def _h3_tile_geometry(
+    prefix_segments: tuple[int, ...],
+    dit_seq_shape: tuple[int, int, int],
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+    """Tile the packed sequence: segment-pure prefix chunks, then video tiles.
+
+    Returns (tile_partition_indices, variable_block_sizes,
+    untile_combined_index, num_prefix_tiles, num_video_tiles).
+    """
+    prefix_len = sum(prefix_segments)
+
+    prefix_sizes: list[int] = []
+    for segment in prefix_segments:
+        full, rem = divmod(segment, _TILE_ELEMS)
+        prefix_sizes.extend([_TILE_ELEMS] * full)
+        if rem:
+            prefix_sizes.append(rem)
+    num_prefix_tiles = len(prefix_sizes)
+
+    ts_t, ts_h, ts_w = VSA_H3_TILE_SIZE
+    t, h, w = dit_seq_shape
+    num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w))
+    video_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device, VSA_H3_TILE_SIZE)
+    num_video_tiles = int(video_sizes.numel())
+
+    video_indices = get_tile_partition_indices(dit_seq_shape, VSA_H3_TILE_SIZE, device) + prefix_len
+    tile_partition_indices = torch.cat([
+        torch.arange(prefix_len, device=device, dtype=torch.long),
+        video_indices,
+    ])
+    # cat promotes the int32 helper output to int64 alongside the prefix sizes
+    variable_block_sizes = torch.cat([
+        torch.tensor(prefix_sizes, dtype=torch.long, device=device),
+        video_sizes,
+    ])
+
+    # get_non_pad_index is lru-cached on tensor identity; variable_block_sizes
+    # is itself cached by this function, so the identity stays stable.
+    non_pad_index = get_non_pad_index(variable_block_sizes, _TILE_ELEMS)
+
+    untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
+    return (tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles)
+
+
+class MiniMaxH3VSABackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_name() -> str:
+        return "VIDEO_SPARSE_ATTN_H3"
+
+    @staticmethod
+    def get_impl_cls() -> type["MiniMaxH3VSAImpl"]:
+        return MiniMaxH3VSAImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["MiniMaxH3VSAMetadata"]:
+        return MiniMaxH3VSAMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxH3VSAMetadataBuilder"]:
+        return MiniMaxH3VSAMetadataBuilder
+
+
+@dataclass
+class MiniMaxH3VSAMetadata(AttentionMetadata):
+    total_seq_length: int
+    num_prefix_tiles: int
+    num_video_tiles: int
+    exempt: bool
+    variable_block_sizes: torch.Tensor
+    untile_combined_index: torch.Tensor
+    # layers forced dense regardless of sparsity (probe-guided opt-outs)
+    dense_layers: tuple[int, ...] = ()
+    # Single-slot holder for the padded tile buffer, owned by the BUILDER so
+    # one buffer serves the whole denoising loop (pad slots stay zero and
+    # every non-pad slot is fully overwritten per tile(), so cross-step reuse
+    # is valid; saves a ~1.4 GB alloc+memset per step at 720p). VSA-H3 runs
+    # eager today — revisit the reuse if it ever goes under cudagraphs.
+    tile_buf_holder: list = None  # type: ignore[assignment]
+
+
+class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        self._tile_buf_holder: list = [None]
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore
+            self,
+            current_timestep: int,
+            raw_latent_shape: tuple[int, int, int],
+            patch_size: tuple[int, int, int],
+            VSA_sparsity: float,
+            prefix_segments: tuple[int, ...],
+            device: torch.device,
+            exempt: bool = True,
+            dense_layers: tuple[int, ...] = (),
+            **kwargs: dict[str, Any],
+    ) -> MiniMaxH3VSAMetadata:
+        dit_seq_shape = (raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1],
+                         raw_latent_shape[2] // patch_size[2])
+        prefix_segments = tuple(int(s) for s in prefix_segments if s > 0)
+        total_seq_length = sum(prefix_segments) + math.prod(dit_seq_shape)
+
+        (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles,
+         num_video_tiles) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device)
+
+        return MiniMaxH3VSAMetadata(
+            current_timestep=current_timestep,
+            VSA_sparsity=VSA_sparsity,
+            total_seq_length=total_seq_length,
+            num_prefix_tiles=num_prefix_tiles,
+            num_video_tiles=num_video_tiles,
+            exempt=exempt,
+            variable_block_sizes=variable_block_sizes,
+            untile_combined_index=untile_combined_index,
+            dense_layers=tuple(int(layer) for layer in dense_layers),
+            tile_buf_holder=self._tile_buf_holder,
+        )
+
+
+def _pool_tiles(x: torch.Tensor, variable_block_sizes: torch.Tensor) -> torch.Tensor:
+    """fp32 mean over each 256-token tile. x: [B, S_pad, H, D] -> [B, H, n_tiles, D].
+
+    Pad positions in the tile buffer are guaranteed zero (zeros-init, never
+    written), so a plain sum with fp32 accumulation needs no validity mask
+    and no materialized fp32 temp; dividing by the true tile size makes it
+    the masked mean exactly.
+    """
+    batch, seq_len, heads, dim = x.shape
+    n_tiles = seq_len // _TILE_ELEMS
+    pooled = x.view(batch, n_tiles, _TILE_ELEMS, heads, dim).sum(dim=2, dtype=torch.float32)
+    pooled = pooled / variable_block_sizes.view(1, -1, 1, 1)
+    return pooled.permute(0, 2, 1, 3)
+
+
+def _build_block_mask(
+    scores: torch.Tensor,
+    num_prefix_tiles: int,
+    num_video_tiles: int,
+    VSA_sparsity: float,
+    exempt: bool,
+) -> torch.Tensor:
+    """scores: [B, H, n_tiles, n_tiles] -> bool mask, same shape."""
+    n_tiles = scores.shape[-1]
+    k_vid = compute_topk(VSA_sparsity, num_video_tiles)
+    if k_vid == num_video_tiles:
+        return torch.ones_like(scores, dtype=torch.bool)
+    mask = torch.zeros_like(scores, dtype=torch.bool)
+    if exempt or num_prefix_tiles == 0:
+        video_cols = scores[..., num_prefix_tiles:]
+        idx = video_cols.topk(k_vid, dim=-1).indices + num_prefix_tiles
+        mask.scatter_(-1, idx, True)
+        mask[..., :num_prefix_tiles] = True
+    else:
+        k_total = min(k_vid + num_prefix_tiles, n_tiles)
+        idx = scores.topk(k_total, dim=-1).indices
+        mask.scatter_(-1, idx, True)
+    mask[:, :, :num_prefix_tiles, :] = True
+    return mask
+
+
+class MiniMaxH3VSAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        self.layer_idx = layer_idx_from_prefix(prefix, default=-1)
+
+    def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
+        """Scatter rows into the padded tile buffer (pad positions stay zero).
+
+        The returned tensor aliases the builder-owned buffer; callers must
+        consume it before the next ``tile()`` (both call sites in
+        ``forward()`` read it immediately).
+        """
+        if x.shape[1] != attn_metadata.total_seq_length:
+            raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, "
+                             f"got {x.shape[1]}. A non-packed sequence (e.g. the token refiner) is "
+                             "routed to the VSA-H3 backend; exclude it from the supported backends.")
+        n_tiles = attn_metadata.variable_block_sizes.numel()
+        target_shape = (x.shape[0], n_tiles * _TILE_ELEMS, x.shape[-2], x.shape[-1])
+
+        if torch.is_grad_enabled() and x.requires_grad:
+            # Grad-tracking forwards (VSA-H3 training students) get a fresh
+            # buffer per call: autograd saves views of the tiled buffer for
+            # backward (layer.py chunks q/k/v out of it), so a later layer's
+            # in-place overwrite of a shared buffer raises the saved-tensor
+            # version error at backward. Mirrors Wan training's
+            # vsa_cache_tile_buf=False default.
+            return scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, None)
+
+        # single scatter: untile_combined_index maps original row i to its
+        # padded slot, so this is exactly the inverse of postprocess_output
+        holder = attn_metadata.tile_buf_holder
+        holder[0] = scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, holder[0])
+        return holder[0]
+
+    def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
+        # Cast before tiling: autocast leaves the packed qkv fp32 (norm
+        # outputs), but the kernel computes in bf16 regardless — tiling in
+        # fp32 just doubles the dominant grad-path transient (the fresh
+        # per-forward tile buffer) and everything chunked out of it.
+        if qkv.dtype is not torch.bfloat16:
+            qkv = qkv.to(torch.bfloat16)
+        return self.tile(qkv, attn_metadata)
+
+    def postprocess_output(self, output: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
+        return output[:, attn_metadata.untile_combined_index]
+
+    def forward(  # type: ignore[override]
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        gate_compress: torch.Tensor | None,
+        attn_metadata: MiniMaxH3VSAMetadata,
+    ) -> torch.Tensor:
+        if block_sparse_attn_256_bshd is None:
+            raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed")
+
+        # The VSA-256 kernel computes its PV dot in bf16 and accepts bf16
+        # only; training paths can leak fp32 activations into attention
+        # (fp32 master weights — the same phenomenon FLASH_ATTN guards
+        # against). Cast through bf16 and restore on output.
+        orig_dtype = query.dtype
+        if orig_dtype is not torch.bfloat16:
+            if not torch.compiler.is_compiling():
+                logger.warning_once(f"VIDEO_SPARSE_ATTN_H3 received {orig_dtype} inputs; casting to "
+                                    "bfloat16 for the kernel and restoring on output.")
+            query = query.to(torch.bfloat16)
+            key = key.to(torch.bfloat16)
+            value = value.to(torch.bfloat16)
+            if gate_compress is not None:
+                gate_compress = gate_compress.to(torch.bfloat16)
+
+        # probe-guided per-layer opt-out: diffuse layers run dense (all-True
+        # mask) while the rest keep the configured sparsity
+        layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity
+        probe_dir = probe_enabled()
+
+        scores = None
+        if layer_sparsity > 0.0 or gate_compress is not None or probe_dir is not None:
+            q_pooled = _pool_tiles(query, attn_metadata.variable_block_sizes)
+            k_pooled = _pool_tiles(key, attn_metadata.variable_block_sizes)
+            scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / (query.shape[-1]**0.5)
+            if probe_dir is not None:
+                record_probe(probe_dir, self.layer_idx, query, key, scores, attn_metadata)
+
+        if scores is None:
+            n_tiles = attn_metadata.variable_block_sizes.numel()
+            mask = torch.ones(query.shape[0], query.shape[2], n_tiles, n_tiles, dtype=torch.bool, device=query.device)
+        else:
+            mask = _build_block_mask(
+                scores,
+                attn_metadata.num_prefix_tiles,
+                attn_metadata.num_video_tiles,
+                layer_sparsity,
+                attn_metadata.exempt,
+            )
+
+        out, _ = block_sparse_attn_256_bshd(query, key, value, mask, attn_metadata.variable_block_sizes)
+
+        if gate_compress is not None:
+            # Wan-style compression branch: dense attention over pooled tiles,
+            # broadcast to each tile's rows, scaled by the learned gate
+            # (zero-initialized for H3 => branch contributes nothing until
+            # finetuned; the model layer skips it entirely for all-zero gates).
+            v_pooled = _pool_tiles(value, attn_metadata.variable_block_sizes)
+            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)  # [B, H, n_tiles, D]
+            out_c = out_c.permute(0, 2, 1, 3).to(out.dtype)  # [B, n_tiles, H, D]
+            batch, _, heads, dim = out.shape
+            n_tiles = attn_metadata.variable_block_sizes.numel()
+            out.view(batch, n_tiles, _TILE_ELEMS, heads,
+                     dim).addcmul_(out_c.unsqueeze(2), gate_compress.view(batch, n_tiles, _TILE_ELEMS, heads, dim))
+        if out.dtype != orig_dtype:
+            out = out.to(orig_dtype)
+        return out

fastvideo/attention/backends/video_sparse_attn_h3_probe.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Attention-mass probe for VSA-H3 selection quality.
+
+Enabled via ``FASTVIDEO_H3_VSA_PROBE=<out_dir>`` on a sparsity-0 run (mask
+selects everything, so the model follows its exact dense trajectory while
+we measure what any top-k WOULD have captured). Per (step, layer) it
+records, from the pooled tile scores:
+
+- ``recall_mean``: pooled softmax mass of video->video attention captured
+  by the top-k video tiles, as a function of k/V, per head;
+- ``prefix_mass``: pooled mass video queries place on text/cond/audio tiles
+  (decides exempt-vs-compete);
+- ``true_recall@{frac}``: token-true validation on sampled query rows —
+  exact attention row aggregated per tile, mass captured by the tiles the
+  POOLED ranking would pick (catches pooling-proxy failure);
+- ``perfect_recall@{frac}``: the selection ceiling — mass captured when
+  ranking by the true tile masses themselves.
+
+One .pt file per (step, layer, rank); aggregate offline.
+"""
+
+import os
+
+import torch
+
+_FRACS = (0.05, 0.10, 0.125, 0.25, 0.50, 0.75)
+_TRUE_ROWS = 128  # sampled query rows per layer for the token-true check
+
+
+def probe_enabled() -> str | None:
+    return os.environ.get("FASTVIDEO_H3_VSA_PROBE") or None
+
+
+@torch.no_grad()
+def record_probe(
+    out_dir: str,
+    layer: int,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    scores: torch.Tensor,
+    attn_metadata,
+) -> None:
+    """scores: [B, H, n, n] scaled pooled scores; query/key tiled BSHD."""
+    step = int(attn_metadata.current_timestep)
+    P = attn_metadata.num_prefix_tiles
+    V = attn_metadata.num_video_tiles
+
+    probs = torch.softmax(scores.float(), dim=-1)  # pooled attention, all tiles
+    vid_q = probs[:, :, P:, :]  # video-query rows
+    prefix_mass = vid_q[..., :P].sum(-1).mean(dim=(0, 2))  # [H]
+
+    vid_vid = vid_q[..., P:]
+    vid_vid = vid_vid / vid_vid.sum(-1, keepdim=True).clamp_min(1e-12)
+    sorted_mass = vid_vid.sort(dim=-1, descending=True).values.cumsum(-1)  # [B,H,Vq,V]
+    ks = [max(1, min(int(round(f * V)), V)) for f in _FRACS]
+    recall = torch.stack([sorted_mass[..., k - 1] for k in ks])  # [F,B,H,Vq]
+    recall_mean = recall.mean(dim=(1, 3))  # [F,H]
+
+    # token-true check on sampled rows (exact row softmax, tile-aggregated)
+    gen = torch.Generator(device="cpu").manual_seed(step * 1000 + layer)
+    # sample among video rows in the PADDED/tiled domain that are non-pad
+    from fastvideo.attention.backends.video_sparse_attn_h3 import token_tile_and_valid
+    token_tile, token_valid = token_tile_and_valid(attn_metadata.variable_block_sizes)
+    video_rows = torch.nonzero((token_tile >= P) & token_valid, as_tuple=False).flatten()
+    idx = video_rows[torch.randint(0, video_rows.numel(), (_TRUE_ROWS, ), generator=gen).to(query.device)]
+
+    q_s = query[:, idx].float()  # [B, R, H, D]
+    logits = torch.einsum("brhd,bshd->bhrs", q_s, key.float()) / (query.shape[-1]**0.5)
+    logits = logits.masked_fill(~token_valid.view(1, 1, 1, -1), float("-inf"))
+    true_probs = torch.softmax(logits, dim=-1)
+    n_tiles = attn_metadata.variable_block_sizes.numel()
+    tile_mass = torch.zeros(*true_probs.shape[:3], n_tiles, device=query.device, dtype=true_probs.dtype)
+    tile_mass.index_add_(3, token_tile, true_probs)  # [B,H,R,n]
+    tm_vid = tile_mass[..., P:]
+    tm_vid = tm_vid / tm_vid.sum(-1, keepdim=True).clamp_min(1e-12)
+    # pooled ranking per sampled row's q-tile
+    pooled_rank = scores[:, :, P:, P:].argsort(dim=-1, descending=True)  # [B,H,Vq,V]
+    row_qtile = (token_tile[idx] - P).view(1, 1, -1, 1).expand(pooled_rank.shape[0], pooled_rank.shape[1], -1, 1)
+    row_rank = pooled_rank.gather(2, row_qtile.expand(-1, -1, -1, pooled_rank.shape[-1]))  # [B,H,R,V]
+    perfect_sorted = tm_vid.sort(dim=-1, descending=True).values.cumsum(-1)  # selection ceiling
+    true_recall = {}
+    perfect_recall = {}
+    for f, k in zip(_FRACS, ks, strict=True):
+        captured = tm_vid.gather(-1, row_rank[..., :k]).sum(-1)  # [B,H,R]
+        true_recall[f] = captured.mean().item()
+        perfect_recall[f] = perfect_sorted[..., k - 1].mean().item()
+
+    os.makedirs(out_dir, exist_ok=True)
+    payload = dict(step=step,
+                   layer=layer,
+                   P=P,
+                   V=V,
+                   fracs=_FRACS,
+                   prefix_mass=prefix_mass.cpu(),
+                   recall_mean=recall_mean.cpu(),
+                   true_recall=true_recall,
+                   perfect_recall=perfect_recall)
+    # under SP each rank holds a distinct head subset; keep every rank's stats
+    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+    torch.save(payload, os.path.join(out_dir, f"probe_step{step:03d}_layer{layer:03d}_r{rank}.pt"))

fastvideo/attention/backends/vmoba.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import re
 from dataclasses import dataclass
 
 import torch
@@ -8,7 +7,7 @@ from einops import rearrange
 
 from fastvideo_kernel import (moba_attn_varlen, process_moba_input, process_moba_output)
 from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
-                                                   AttentionMetadataBuilder)
+                                                   AttentionMetadataBuilder, layer_idx_from_prefix)
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -132,10 +131,7 @@ class VMOBAAttentionImpl(AttentionImpl):
         self.pad_input = pad_input
 
     def _get_layer_idx(self, prefix: str) -> int | None:
-        match = re.search(r"blocks\.(\d+)", prefix)
-        if not match:
-            raise ValueError(f"Invalid prefix: {prefix}")
-        return int(match.group(1))
+        return layer_idx_from_prefix(prefix)
 
     def forward(
         self,

fastvideo/attention/layer.py

@@ -207,11 +207,21 @@ class DistributedAttention_VSA(DistributedAttention):
         ctx_attn_metadata = forward_context.attn_metadata
 
         batch_size, seq_len, num_heads, head_dim = q.shape
-        # Stack QKV
-        qkvg = torch.cat([q, k, v, gate_compress], dim=0)  # [4*batch, seq_len, num_heads, head_dim]
+        # VSA kernels are bf16-only; cast before the cat/all-to-all so the
+        # packed qkvg comm buffers and tile scatter move half-width tensors
+        # instead of carrying autocast's fp32 norm outputs through the
+        # sequence-parallel exchange (the dominant grad-path transient).
+        if q.dtype is not torch.bfloat16:
+            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
+            if gate_compress is not None:
+                gate_compress = gate_compress.to(torch.bfloat16)
+        # Stack QKV (a caller with a structurally-zero gate passes None and
+        # skips the gate's share of the all-to-all/tile traffic)
+        stack = [q, k, v] if gate_compress is None else [q, k, v, gate_compress]
+        qkvg = torch.cat(stack, dim=0)  # [3or4*batch, seq_len, num_heads, head_dim]
 
         # Redistribute heads across sequence dimension
-        # Before: [4*batch, shard_seq_len, num_heads, head_dim]
+        # Before: [3or4*batch, shard_seq_len, num_heads, head_dim]
         # After:  [4*batch, full_seq_len, shard_num_heads, head_dim]
         qkvg = sequence_model_parallel_all_to_all_4D(qkvg, scatter_dim=2, gather_dim=1)
 
@@ -225,7 +235,10 @@ class DistributedAttention_VSA(DistributedAttention):
 
         qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata)
 
-        q, k, v, gate_compress = qkvg.chunk(4, dim=0)
+        if gate_compress is None:
+            q, k, v = qkvg.chunk(3, dim=0)
+        else:
+            q, k, v, gate_compress = qkvg.chunk(4, dim=0)
         output = self.attn_impl.forward(q, k, v, gate_compress, ctx_attn_metadata)  # type: ignore[call-arg]
 
         # Redistribute back if using sequence parallelism

fastvideo/configs/models/dits/minimax_h3.py

@@ -29,6 +29,7 @@ class MiniMaxH3ArchConfig(DiTArchConfig):
         # sm_12x). Enabled for speed experiments; output quality against the
         # SSIM references is not yet validated.
         AttentionBackendEnum.ATTN_QAT_INFER,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,
     )
 
     param_names_mapping: dict = field(

fastvideo/entrypoints/video_generator.py

@@ -71,6 +71,11 @@ _BATCH_EXTRA_PASSTHROUGH_KEYS: tuple[str, ...] = (
     "ltx2_audio_denoise_mask",
     "audio_num_frames",
     "video_position_offset_sec",
+    # MiniMax-H3 VSA per-request knobs (read by the H3 denoising stage;
+    # sparsity itself flows through the existing ForwardBatch.VSA_sparsity)
+    "vsa_mode",
+    "vsa_dense_first_n_steps",
+    "vsa_dense_layers",
 )
 
 _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({

fastvideo/models/dits/minimax_h3.py

@@ -11,6 +11,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from fastvideo.attention import DistributedAttention
+from fastvideo.attention.layer import DistributedAttention_VSA
+from fastvideo.attention.selector import get_attn_backend
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
 from fastvideo.distributed.communication_op import (
     sequence_model_parallel_all_gather_with_unpad,
@@ -23,6 +25,7 @@ from fastvideo.layers.quantization import QuantizationConfig
 from fastvideo.layers.visual_embedding import Timesteps
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.utils import get_compute_dtype
 
 MINIMAX_H3_MODALITY_NUM = 3
 _CFG = MiniMaxH3Config()
@@ -131,13 +134,58 @@ class MiniMaxH3Attention(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.to_out",
         )
-        self.distributed_attention = DistributedAttention(
+        # VSA carries a learned gate on its pooled-compression branch. The H3
+        # checkpoint has no such weight, so the loader zero-initializes it
+        # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled
+        # until finetuned. Built only when VSA-H3 actually resolves, keeping
+        # the FLASH/SDPA paths and their state_dict untouched.
+        resolved_backend = get_attn_backend(attention_head_dim,
+                                            get_compute_dtype(),
+                                            supported_attention_backends=supported_attention_backends)
+        use_vsa = resolved_backend.get_name() == "VIDEO_SPARSE_ATTN_H3"
+        attention_cls = DistributedAttention_VSA if use_vsa else DistributedAttention
+        self.distributed_attention = attention_cls(
             num_heads=num_attention_heads,
             head_size=attention_head_dim,
             causal=False,
             supported_attention_backends=supported_attention_backends,
             prefix=prefix,
         )
+        self.to_gate_compress: ReplicatedLinear | None = None
+        # None = unchecked; the first forward tests the loaded weight once and
+        # skips the gate branch entirely while it is structurally zero.
+        self._gate_compress_active: bool | None = None
+        if use_vsa:
+            self.to_gate_compress = ReplicatedLinear(
+                hidden_size,
+                inner_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_gate_compress",
+            )
+
+    def _gate_active(self) -> bool:
+        """True when the compression gate can contribute.
+
+        The zero-initialized (not-yet-finetuned) gate makes the whole branch
+        a guaranteed zero — a full GEMM plus a third more all-to-all/tile
+        traffic per layer for nothing — so test the loaded weight once and
+        cache the answer. Under grad the branch always runs (gradients must
+        reach a zero gate for it to ever train).
+        """
+        if torch.is_grad_enabled():
+            # Training can turn the zero-init gate nonzero; drop the cached
+            # answer so the next no-grad forward (validation sampling)
+            # re-tests the weight instead of skipping a branch that has
+            # started contributing.
+            self._gate_compress_active = None
+            return True
+        if self._gate_compress_active is None:
+            weight = self.to_gate_compress.weight
+            # bool() on a DTensor reduction resolves collectively, so every
+            # rank caches the same answer.
+            self._gate_compress_active = bool((weight != 0).any())
+        return self._gate_compress_active
 
     @staticmethod
     def _apply_rotary_emb(
@@ -176,12 +224,18 @@ class MiniMaxH3Attention(nn.Module):
 
         # H3 rotates only 96/128 channels, which the generic `freqs_cis`
         # branch cannot express. Apply it above, then pass no RoPE here.
+        extra_attention_kwargs = {}
+        if self.to_gate_compress is not None and self._gate_active():
+            gate_compress, _ = self.to_gate_compress(hidden_states)
+            extra_attention_kwargs["gate_compress"] = gate_compress.unflatten(
+                -1, (self.num_attention_heads, self.attention_head_dim))
         hidden_states, _ = self.distributed_attention(
             query,
             key,
             value,
             original_seq_len=original_seq_len,
             freqs_cis=None,
+            **extra_attention_kwargs,
         )
         hidden_states = hidden_states.flatten(2, 3).type_as(query)
         hidden_states, _ = self.to_out(hidden_states)
@@ -521,7 +575,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             arch.norm_eps,
             arch.qk_norm_eps,
             arch.final_norm_eps,
-            self.supported_attention_backends,
+            # The refiner attends over the text stream only; the packed-sequence
+            # VSA backend must never be selected for it.
+            tuple(backend for backend in self.supported_attention_backends
+                  if backend != AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3),
             config.quant_config,
             prefix=f"{config.prefix}.token_refiner",
         )

fastvideo/models/loader/fsdp_load.py

@@ -19,6 +19,11 @@ from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, MixedPrecision
 from torch.nn.modules.module import _IncompatibleKeys
 
 from fastvideo.logger import init_logger
+from fastvideo.models.loader.shard_cache import (
+    shard_cache_context,
+    try_load_from_shard_cache,
+    write_shard_cache,
+)
 from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict)
 from fastvideo.models.loader.weight_utils import safetensors_weights_iterator
 from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
@@ -201,17 +206,37 @@ def maybe_load_fsdp_model(
                     fsdp_shard_conditions=model._fsdp_shard_conditions,
                     pin_cpu_memory=pin_cpu_memory)
 
-    weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
-    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
-    load_model_from_full_model_state_dict(
-        model,
-        weight_iterator,
-        device,
-        default_dtype,
-        strict=strict,
-        cpu_offload=cpu_offload,
-        param_names_mapping=param_names_mapping_fn,
-    )
+    # Sharded base-weight cache (opt-in via FASTVIDEO_WEIGHT_SHARD_CACHE):
+    # rebuild local DTensor chunks from tmpfs instead of re-reading and
+    # re-scattering the full checkpoint on every relaunch. Any miss or
+    # validation failure falls through to the full load below.
+    shard_cache_ctx = None
+    if use_fsdp and not cpu_offload:
+        shard_cache_ctx = shard_cache_context(
+            weight_dir_list=weight_dir_list,
+            device_mesh=device_mesh,
+            hsdp_replicate_dim=hsdp_replicate_dim,
+            hsdp_shard_dim=hsdp_shard_dim,
+            default_dtype=default_dtype,
+            param_dtype=param_dtype,
+            param_names_mapping=model.param_names_mapping,
+        )
+    cache_hit = (shard_cache_ctx is not None
+                 and try_load_from_shard_cache(model, shard_cache_ctx, device, strict=strict))
+    if not cache_hit:
+        weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
+        param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
+        load_model_from_full_model_state_dict(
+            model,
+            weight_iterator,
+            device,
+            default_dtype,
+            strict=strict,
+            cpu_offload=cpu_offload,
+            param_names_mapping=param_names_mapping_fn,
+        )
+        if shard_cache_ctx is not None:
+            write_shard_cache(model, shard_cache_ctx)
     if hasattr(model, "materialize_non_persistent_buffers"):
         model.materialize_non_persistent_buffers(device=device, dtype=default_dtype)
     for n, p in chain(model.named_parameters(), model.named_buffers()):
@@ -440,10 +465,14 @@ def load_model_from_full_model_state_dict(
                 sharded_tensor = full_tensor
         else:
             full_tensor = full_tensor.to(device=device, dtype=target_dtype)
+            # Every rank read the identical full tensor from the checkpoint, so
+            # each can slice its own shard locally; the default src_data_rank=0
+            # scatters from rank 0 and throws away the other ranks' copies.
             sharded_tensor = distribute_tensor(
                 full_tensor,
                 meta_sharded_param.device_mesh,
                 meta_sharded_param.placements,
+                src_data_rank=None,
             )
             if cpu_offload:
                 sharded_tensor = sharded_tensor.cpu()

fastvideo/models/loader/shard_cache.py

@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-rank sharded base-weight cache for fast training relaunches.
+
+After a full checkpoint load, each shard rank persists its local DTensor
+chunks (post-rename, post-cast — exactly what ``assign=True`` installed) as
+one safetensors file in a cache directory (typically tmpfs). Subsequent
+launches with an identical (checkpoint, mesh layout, dtype, name-mapping)
+tuple rebuild the model from those chunks via ``DTensor.from_local`` —
+skipping the full-tensor reads, per-rank H2D of the whole checkpoint, and
+the distribute/scatter step.
+
+Opt-in via ``FASTVIDEO_WEIGHT_SHARD_CACHE=<dir>`` (e.g. ``/dev/shm/fv-wcache``).
+Any validation failure or exception degrades to the normal full load — the
+cache can never fail a run.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.tensor import DTensor
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+_FORMAT_VERSION = 1
+_ENV_DIR = "FASTVIDEO_WEIGHT_SHARD_CACHE"
+_ENV_MAX_GB = "FASTVIDEO_WEIGHT_SHARD_CACHE_MAX_GB"
+# Set to "1" when the cache dir is node-local (tmpfs) on a multi-node run so
+# every replicate rank writes its own node's copy.
+_ENV_PER_NODE = "FASTVIDEO_WEIGHT_SHARD_CACHE_PER_NODE"
+# Mirrors fsdp_load's zero-init allowance for params absent from checkpoints.
+_ALLOWED_NEW_PARAM_PATTERNS = ("gate_compress", "proj_l")
+_WRITE_MARGIN_BYTES = 5 << 30
+
+
+@dataclass
+class ShardCacheContext:
+    entry_dir: Path
+    key: str
+    shard_index: int
+    num_shards: int
+    is_writer: bool  # replicate-coordinate 0 writes; replicas are identical
+
+
+def _expand_weight_files(weight_dir_list: list[str]) -> list[str]:
+    files: list[str] = []
+    for entry in weight_dir_list:
+        if os.path.isdir(entry):
+            files.extend(str(p) for p in sorted(Path(entry).glob("*.safetensors")))
+        else:
+            files.append(entry)
+    return files
+
+
+def _shard_file(entry_dir: Path, shard_index: int, num_shards: int) -> Path:
+    return entry_dir / f"shard{shard_index}-of-{num_shards}.safetensors"
+
+
+def shard_cache_context(
+    *,
+    weight_dir_list: list[str],
+    device_mesh: Any,
+    hsdp_replicate_dim: int,
+    hsdp_shard_dim: int,
+    default_dtype: torch.dtype,
+    param_dtype: torch.dtype,
+    param_names_mapping: dict[str, str] | None,
+) -> ShardCacheContext | None:
+    root = os.environ.get(_ENV_DIR)
+    if not root:
+        return None
+    try:
+        files = _expand_weight_files(weight_dir_list)
+        if not files:
+            return None
+        stats = sorted((os.path.basename(p), os.stat(p).st_size, os.stat(p).st_mtime_ns) for p in files)
+        mapping_items = sorted((param_names_mapping or {}).items())
+        key_material = json.dumps(
+            [
+                _FORMAT_VERSION,
+                stats,
+                int(hsdp_replicate_dim),
+                int(hsdp_shard_dim),
+                str(default_dtype),
+                str(param_dtype),
+                mapping_items,
+            ],
+            sort_keys=True,
+        )
+        key = hashlib.sha256(key_material.encode()).hexdigest()[:16]
+        coordinate = device_mesh.get_coordinate()
+        if coordinate is None:
+            return None
+        # mesh dims are ("replicate", "shard")
+        replicate_index, shard_index = int(coordinate[0]), int(coordinate[1])
+        # With a shared cache dir, replicate-coordinate 0 writes and replicas
+        # (identical shards) skip to avoid same-file collisions. With a
+        # node-local dir (e.g. /dev/shm on a multi-node HSDP run), every
+        # replica must write its own node's copy — same-name collisions are
+        # impossible across nodes, and a single-writer rule would leave every
+        # non-zero replica's node permanently cold (all-reduce MIN then turns
+        # that into a global miss).
+        per_node_root = os.environ.get(_ENV_PER_NODE, "0") == "1"
+        return ShardCacheContext(
+            entry_dir=Path(root) / key,
+            key=key,
+            shard_index=shard_index,
+            num_shards=int(hsdp_shard_dim),
+            is_writer=per_node_root or replicate_index == 0,
+        )
+    except Exception as exc:  # noqa: BLE001 - cache must never fail a load
+        logger.warning("shard cache disabled for this load (context error): %s", exc)
+        return None
+
+
+def _all_ranks_agree(local_ok: bool, device: torch.device) -> bool:
+    if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1):
+        return local_ok
+    flag = torch.tensor([1 if local_ok else 0], device=device, dtype=torch.int32)
+    dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+    return bool(flag.item())
+
+
+def _validate_entry(name: str, entry: dict[str, Any], meta_param: torch.Tensor) -> bool:
+    if entry["dtype"] != str(meta_param.dtype):
+        return False
+    if list(entry["global_shape"]) != list(meta_param.shape):
+        return False
+    is_dtensor = isinstance(meta_param, DTensor)
+    if entry["kind"] != ("dtensor" if is_dtensor else "tensor"):
+        return False
+    if is_dtensor:
+        if entry["placements"] != [str(p) for p in meta_param.placements]:
+            return False
+        if list(entry["local_shape"]) != list(meta_param.to_local().shape):
+            return False
+    return True
+
+
+def try_load_from_shard_cache(
+    model: nn.Module,
+    ctx: ShardCacheContext,
+    device: torch.device,
+    *,
+    strict: bool = True,
+) -> bool:
+    """Assemble the model's state dict from cached local shards. Returns False
+    (without mutating the model) on any mismatch."""
+    try:
+        manifest_path = ctx.entry_dir / "manifest.json"
+        shard_path = _shard_file(ctx.entry_dir, ctx.shard_index, ctx.num_shards)
+        local_ok = manifest_path.is_file() and shard_path.is_file()
+        manifest: dict[str, Any] = {}
+        if local_ok:
+            manifest = json.loads(manifest_path.read_text())
+            local_ok = (manifest.get("format_version") == _FORMAT_VERSION and manifest.get("key") == ctx.key)
+        meta_sd = model.state_dict()
+        if local_ok:
+            params_table = manifest["params"]
+            for name, meta_param in meta_sd.items():
+                entry = params_table.get(name)
+                if entry is None:
+                    if not any(pattern in name for pattern in _ALLOWED_NEW_PARAM_PATTERNS):
+                        local_ok = False
+                        break
+                    continue
+                if not _validate_entry(name, entry, meta_param):
+                    local_ok = False
+                    break
+        if not _all_ranks_agree(local_ok, device):
+            if local_ok:
+                logger.info("shard cache: another rank missed entry %s; falling back to full load", ctx.key)
+            return False
+
+        from safetensors import safe_open
+
+        named_buffers = dict(model.named_buffers())
+        dtype_selector = getattr(model, "_get_parameter_dtype", None)
+        sharded_sd: dict[str, Any] = {}
+        with safe_open(str(shard_path), framework="pt", device=str(device)) as f:
+            cached_keys = set(f.keys())
+            for name, meta_param in meta_sd.items():
+                if name in cached_keys:
+                    local = f.get_tensor(name)
+                    if isinstance(meta_param, DTensor):
+                        tensor: torch.Tensor = DTensor.from_local(
+                            local,
+                            meta_param.device_mesh,
+                            meta_param.placements,
+                            run_check=False,
+                            shape=meta_param.shape,
+                            stride=meta_param.stride(),
+                        )
+                    else:
+                        tensor = local
+                else:
+                    # Zero-init new params exactly like the full-load path.
+                    target_dtype = meta_param.dtype
+                    if callable(dtype_selector):
+                        target_dtype = dtype_selector(name, target_dtype)
+                    if isinstance(meta_param, DTensor):
+                        local = torch.zeros(meta_param.to_local().shape, device=device, dtype=target_dtype)
+                        tensor = DTensor.from_local(
+                            local,
+                            meta_param.device_mesh,
+                            meta_param.placements,
+                            run_check=False,
+                            shape=meta_param.shape,
+                            stride=meta_param.stride(),
+                        )
+                    else:
+                        tensor = torch.zeros(meta_param.shape, device=device, dtype=target_dtype)
+                sharded_sd[name] = tensor if name in named_buffers else nn.Parameter(tensor)
+
+        reverse_map = manifest.get("reverse_param_names_mapping", {})
+        model.reverse_param_names_mapping = {k: tuple(v) for k, v in reverse_map.items()}
+        model.load_state_dict(sharded_sd, strict=strict, assign=True)
+        # Freshen mtimes so mtime-based tmpfs cleaners (and our own LRU GC)
+        # treat actively used entries as recent.
+        for p in (shard_path, manifest_path):
+            try:
+                os.utime(p)
+            except OSError:
+                pass
+        logger.info(
+            "shard cache HIT %s: %d tensors from %s",
+            ctx.key,
+            len(sharded_sd),
+            shard_path,
+        )
+        return True
+    except Exception as exc:  # noqa: BLE001 - cache must never fail a load
+        logger.warning("shard cache load failed (%s); falling back to full load", exc)
+        return False
+
+
+def write_shard_cache(model: nn.Module, ctx: ShardCacheContext) -> None:
+    """Persist this rank's local shards after a successful full load."""
+    try:
+        from safetensors.torch import save_file
+
+        tensors: dict[str, torch.Tensor] = {}
+        params_table: dict[str, Any] = {}
+        for name, value in model.state_dict().items():
+            if isinstance(value, DTensor):
+                local = value.to_local().detach().to("cpu", copy=True).contiguous()
+                params_table[name] = {
+                    "kind": "dtensor",
+                    "dtype": str(value.dtype),
+                    "global_shape": list(value.shape),
+                    "placements": [str(p) for p in value.placements],
+                    "local_shape": list(local.shape),
+                }
+            else:
+                local = value.detach().to("cpu", copy=True).contiguous()
+                params_table[name] = {
+                    "kind": "tensor",
+                    "dtype": str(value.dtype),
+                    "global_shape": list(value.shape),
+                }
+            tensors[name] = local
+
+        ctx.entry_dir.mkdir(parents=True, exist_ok=True)
+        needed = sum(t.numel() * t.element_size() for t in tensors.values())
+        free = shutil.disk_usage(ctx.entry_dir).free
+        if free < needed + _WRITE_MARGIN_BYTES:
+            logger.warning(
+                "shard cache: skipping write (%.1f GiB needed, %.1f GiB free at %s)",
+                needed / 2**30,
+                free / 2**30,
+                ctx.entry_dir,
+            )
+            _barrier_if_initialized()
+            return
+
+        if ctx.is_writer:
+            shard_path = _shard_file(ctx.entry_dir, ctx.shard_index, ctx.num_shards)
+            tmp_path = shard_path.with_suffix(".safetensors.tmp")
+            save_file(tensors, str(tmp_path))
+            os.replace(tmp_path, shard_path)
+        _barrier_if_initialized()
+
+        rank = dist.get_rank() if (dist.is_available() and dist.is_initialized()) else 0
+        if rank == 0:
+            reverse_map = {
+                k: list(v)
+                for k, v in getattr(model, "reverse_param_names_mapping", {}).items()
+            }
+            manifest = {
+                "format_version": _FORMAT_VERSION,
+                "key": ctx.key,
+                "num_shards": ctx.num_shards,
+                "params": params_table,
+                "reverse_param_names_mapping": reverse_map,
+            }
+            manifest_tmp = ctx.entry_dir / "manifest.json.tmp"
+            manifest_tmp.write_text(json.dumps(manifest))
+            os.replace(manifest_tmp, ctx.entry_dir / "manifest.json")
+            logger.info("shard cache WRITE %s: %d tensors -> %s", ctx.key, len(tensors), ctx.entry_dir)
+            _gc_cache_root(ctx.entry_dir.parent, keep=ctx.entry_dir.name)
+    except Exception as exc:  # noqa: BLE001 - cache must never fail a run
+        logger.warning("shard cache write failed (non-fatal): %s", exc)
+        _barrier_if_initialized()
+
+
+def _barrier_if_initialized() -> None:
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        dist.barrier()
+
+
+def _gc_cache_root(root: Path, keep: str) -> None:
+    """Drop least-recently-written entries beyond the size cap."""
+    try:
+        max_bytes = float(os.environ.get(_ENV_MAX_GB, "300")) * 2**30
+        entries = []
+        for child in root.iterdir():
+            if not child.is_dir():
+                continue
+            manifest = child / "manifest.json"
+            size = sum(f.stat().st_size for f in child.glob("*") if f.is_file())
+            mtime = manifest.stat().st_mtime if manifest.is_file() else 0.0
+            entries.append((mtime, size, child))
+        total = sum(size for _, size, _ in entries)
+        for mtime, size, child in sorted(entries):
+            if total <= max_bytes:
+                break
+            if child.name == keep:
+                continue
+            shutil.rmtree(child, ignore_errors=True)
+            total -= size
+            logger.info("shard cache GC: evicted %s (%.1f GiB)", child, size / 2**30)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("shard cache GC failed (non-fatal): %s", exc)

fastvideo/models/schedulers/scheduling_minimax_h3.py

@@ -51,6 +51,11 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin):
             raise ValueError(f"`shift` must be positive, got {shift}.")
         self._shift = float(shift)
 
+    def shift_sigmas(self, base_sigmas: torch.Tensor) -> torch.Tensor:
+        """Warp unshifted base noise amounts onto this scheduler's shifted grid."""
+        base = torch.as_tensor(base_sigmas, dtype=torch.float32)
+        return self._shift * base / (1 + (self._shift - 1) * base)
+
     def set_timesteps(
         self,
         num_inference_steps: int | None = None,
@@ -62,7 +67,7 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin):
                 raise ValueError("`set_timesteps` requires explicit `sigmas` or "
                                  f"`num_inference_steps` >= 2, got {num_inference_steps}.")
             base = torch.linspace(1.0, 0.0, int(num_inference_steps), dtype=torch.float32)
-            sigma_tensor = self._shift * base / (1 + (self._shift - 1) * base)
+            sigma_tensor = self.shift_sigmas(base)
             sigma_tensor = torch.unique_consecutive(sigma_tensor)
         else:
             sigma_tensor = torch.as_tensor(sigmas, dtype=torch.float32).flatten().cpu()

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py

@@ -4,13 +4,16 @@
 from __future__ import annotations
 
 import contextlib
+import os
 from typing import Any
 
 import torch
 
+from fastvideo.attention.selector import component_attention_backend, get_attn_backend
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
 from fastvideo.profiler import get_global_controller
 from fastvideo.hooks.activation_trace import trace_step
 from fastvideo.pipelines.basic.minimax_h3.packing import (
@@ -23,6 +26,42 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.utils import get_compute_dtype
+
+logger = init_logger(__name__)
+
+
+def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> Any:
+    """Builder instance when the transformer resolved to VSA-H3, else None.
+
+    Resolves through the same selector record the attention layers used
+    (``component_attention_backend``) instead of introspecting module
+    internals, mirroring the generic DenoisingStage.
+    """
+    dit_config = fastvideo_args.pipeline_config.dit_config
+    backend = get_attn_backend(
+        head_size=dit_config.attention_head_dim,
+        dtype=get_compute_dtype(),
+        supported_attention_backends=dit_config._supported_attention_backends,
+        requested=component_attention_backend(transformer),
+    )
+    if backend.get_name() != "VIDEO_SPARSE_ATTN_H3":
+        return None
+    return backend.get_builder_cls()()
+
+
+def _h3_vsa_prefix_segments(layout: MiniMaxH3PackedLayout, patch_size: tuple[int, int, int]) -> tuple[int, ...]:
+    """Segment sizes preceding the generated-video tail, validated against the layout."""
+    n_text = int(layout.text_indices.numel())
+    n_cond = int(layout.num_condition_video_rows)
+    n_audio = int(layout.audio_indices.numel())
+    n_video = ((layout.num_video_latent_frames // patch_size[0]) * (layout.latent_height // patch_size[1]) *
+               (layout.latent_width // patch_size[2]))
+    if n_text + n_cond + n_audio + n_video != layout.sequence_length:
+        raise ValueError("VSA-H3 supports the standard [text|cond|audio|video] packing only; "
+                         f"segments ({n_text}, {n_cond}, {n_audio}) + video {n_video} do not sum to "
+                         f"sequence length {layout.sequence_length}.")
+    return n_text, n_cond, n_audio
 
 
 class MiniMaxH3DenoisingStage(PipelineStage):
@@ -69,8 +108,32 @@ class MiniMaxH3DenoisingStage(PipelineStage):
             batch.latents = batch.latents.to(device)
             batch.audio_latents = batch.audio_latents.to(device)
 
-        self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
-        self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device)
+        dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
+        if not dmd_steps:
+            env_steps = os.environ.get("FASTVIDEO_DMD_DENOISING_STEPS", "").strip()
+            if env_steps:
+                dmd_steps = [int(s) for s in env_steps.split(",") if s.strip()]
+        # Stochastic backward simulation (reference-DMD2 sampling): re-noise
+        # each hop's x0 estimate with FRESH noise instead of the deterministic
+        # Euler carry — matching the training rollout's input manifold. Only
+        # meaningful with explicit DMD steps.
+        stochastic_renoise = bool(dmd_steps) and (
+            bool(getattr(fastvideo_args.pipeline_config, "dmd_stochastic_renoise", False))
+            or os.environ.get("FASTVIDEO_DMD_STOCHASTIC_RENOISE", "0").strip().lower() in ("1", "true", "yes"))
+        if stochastic_renoise:
+            logger.info("MiniMax-H3 DMD denoising with stochastic fresh-noise re-noising.")
+        if dmd_steps:
+            # DMD method timesteps are integers on the shared 1000-step grid;
+            # normalize to base time and let each scheduler apply its own shift
+            # (mirroring MiniMaxH3DMDModel._noise_amounts), so few-step sampling
+            # visits exactly the trained jump points — one forward per step.
+            base = torch.tensor([float(step) / 1000.0 for step in dmd_steps] + [0.0], dtype=torch.float32)
+            self.scheduler.set_timesteps(sigmas=self.scheduler.shift_sigmas(base), device=device)
+            self.audio_scheduler.set_timesteps(sigmas=self.audio_scheduler.shift_sigmas(base), device=device)
+            logger.info("MiniMax-H3 denoising with explicit DMD steps %s.", list(dmd_steps))
+        else:
+            self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
+            self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device)
         video_timesteps = self.scheduler.timesteps
         audio_timesteps = self.audio_scheduler.timesteps
         if video_timesteps is None or audio_timesteps is None:
@@ -99,6 +162,26 @@ class MiniMaxH3DenoisingStage(PipelineStage):
         text_indices = layout.text_indices.to(device)
         prompt_embeds = batch.prompt_embeds[0].to(device)
 
+        vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args)
+        if vsa_metadata_builder is not None:
+            vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
+            vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size)
+            # Per-request knobs (sweeps flip these between generate_video calls
+            # without respawning workers); mode None defers to the env default.
+            vsa_mode = batch.extra.get("vsa_mode", "exempt")
+            if vsa_mode not in ("exempt", "compete"):
+                raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.")
+            vsa_exempt = vsa_mode == "exempt"
+            vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ()))
+            vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0))
+            # A per-request batch sparsity wins when set; otherwise fall back
+            # to the run-level args value (training validation and the CLI set
+            # fastvideo_args.VSA_sparsity — nothing populates the batch field
+            # on those paths, and the batch's 0.0 default silently sampled
+            # dense while training ran sparse).
+            vsa_sparsity_base = (float(batch.VSA_sparsity)
+                                 if float(batch.VSA_sparsity) > 0.0 else float(fastvideo_args.VSA_sparsity))
+
         controller = get_global_controller()
         denoise_region = (controller.region("profiler_region_inference_denoising")
                           if controller is not None else contextlib.nullcontext())
@@ -107,6 +190,23 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                 for index, (video_timestep,
                             audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)):
                     unique_timesteps, timestep_indices = row_timestep_plan[index]
+                    attn_metadata = None
+                    if vsa_metadata_builder is not None:
+                        # Optional schedule: run the first N steps dense (sparsity 0
+                        # selects every tile — parity-proven ≡ dense ≤2e-4); early
+                        # steps set global structure and are the most damage-prone.
+                        vsa_sparsity = 0.0 if index < vsa_dense_first_n else vsa_sparsity_base
+                        attn_metadata = vsa_metadata_builder.build(
+                            current_timestep=index,
+                            raw_latent_shape=(layout.num_video_latent_frames, layout.latent_height,
+                                              layout.latent_width),
+                            patch_size=vsa_patch_size,
+                            VSA_sparsity=vsa_sparsity,
+                            prefix_segments=vsa_prefix_segments,
+                            device=device,
+                            exempt=vsa_exempt,
+                            dense_layers=vsa_dense_layers,
+                        )
                     # Under torch.compile(mode="reduce-overhead") each denoising
                     # step must be marked, or cudagraph trees flag cross-step
                     # reuse of pooled outputs as "accessing tensor output of
@@ -115,7 +215,7 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                     torch.compiler.cudagraph_mark_step_begin()
                     with trace_step(index), set_forward_context(
                             current_timestep=index,
-                            attn_metadata=None,
+                            attn_metadata=attn_metadata,
                             forward_batch=batch,
                     ):
                         video_velocity, audio_velocity = self.transformer(
@@ -133,18 +233,52 @@ class MiniMaxH3DenoisingStage(PipelineStage):
 
                     video_start = layout.num_condition_video_rows
                     audio_start = layout.num_condition_audio_rows
-                    batch.latents[video_start:] = self.scheduler.step(
-                        video_velocity[0, video_start:].float(),
-                        video_timestep,
-                        batch.latents[video_start:],
-                        return_dict=False,
-                    )[0]
-                    batch.audio_latents[audio_start:] = self.audio_scheduler.step(
-                        audio_velocity[0, audio_start:].float(),
-                        audio_timestep,
-                        batch.audio_latents[audio_start:],
-                        return_dict=False,
-                    )[0]
+                    if os.environ.get("FASTVIDEO_DMD_DEBUG_STATS", "0") == "1":
+                        with torch.no_grad():
+                            for tag, lat, vel, st, sig in (
+                                ("video", batch.latents, video_velocity, video_start, self.scheduler.sigmas),
+                                ("audio", batch.audio_latents, audio_velocity, audio_start,
+                                 self.audio_scheduler.sigmas),
+                            ):
+                                s = float(sig[index])
+                                xin = lat[st:].float()
+                                x0dbg = xin + s * vel[0, st:].float()
+                                logger.info(
+                                    "DMD_DEBUG step=%d %s sigma=%.4f in(std=%.4f,mean=%.4f) "
+                                    "x0(std=%.4f,mean=%.4f) v(std=%.4f)", index, tag, s, xin.std(), xin.mean(),
+                                    x0dbg.std(), x0dbg.mean(), vel[0, st:].float().std())
+                    if stochastic_renoise:
+                        # x0 = sample + sigma * v (velocity = noise - clean),
+                        # then fresh-noise re-mix at the next sigma — the same
+                        # hop rule as the DMD2 training rollout.
+                        assert self.scheduler.sigmas is not None and self.audio_scheduler.sigmas is not None
+                        for latents, velocity, start, sigmas in (
+                            (batch.latents, video_velocity, video_start, self.scheduler.sigmas),
+                            (batch.audio_latents, audio_velocity, audio_start, self.audio_scheduler.sigmas),
+                        ):
+                            sigma = float(sigmas[index])
+                            sigma_next = float(sigmas[index + 1])
+                            sample = latents[start:].float()
+                            pred_x0 = sample + sigma * velocity[0, start:].float()
+                            if sigma_next > 0.0:
+                                noise = torch.randn_like(pred_x0)
+                                nxt = (1.0 - sigma_next) * pred_x0 + sigma_next * noise
+                            else:
+                                nxt = pred_x0
+                            latents[start:] = nxt.to(latents.dtype)
+                    else:
+                        batch.latents[video_start:] = self.scheduler.step(
+                            video_velocity[0, video_start:].float(),
+                            video_timestep,
+                            batch.latents[video_start:],
+                            return_dict=False,
+                        )[0]
+                        batch.audio_latents[audio_start:] = self.audio_scheduler.step(
+                            audio_velocity[0, audio_start:].float(),
+                            audio_timestep,
+                            batch.audio_latents[audio_start:],
+                            return_dict=False,
+                        )[0]
                     batch.step_index = index
                     batch.timestep = video_timestep
         finally:

fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py

@@ -313,21 +313,35 @@ def validate_preprocessed_training_data(
     print(f"Validated one Crush-Smol H3 training row in {expected_path}")
 
 
-def main() -> None:
-    """Encode the selected Crush-Smol record with each H3 component in sequence.
+def main(
+    video: Path | None = None,
+    prompt: str | None = None,
+    output_dir: Path = OUTPUT_DIR,
+    model_path: Path = MODEL_PATH,
+) -> None:
+    """Encode one audio-video-caption record with each H3 component in sequence.
 
+    With no arguments this encodes the pinned Crush-Smol record; ``--video`` +
+    ``--prompt`` preprocess an arbitrary mp4 (with soundtrack) instead.
     Releasing each component before loading the next component keeps video,
     audio, and text preprocessing within one GPU's memory.
     """
     _init_single_process_distributed()
-    resolved_model_path = MODEL_PATH.resolve()
+    resolved_model_path = model_path.resolve()
     if not resolved_model_path.is_dir():
         raise FileNotFoundError(f"Filtered MiniMax H3 model directory is missing at {resolved_model_path}")
     model_index = verify_model_config_and_directory(str(resolved_model_path))
-    video_path, caption = load_crush_smol_training_sample(
-        DATA_DIR / "videos2caption.json",
-        DATA_DIR / "videos",
-    )
+    if video is None:
+        video_path, caption = load_crush_smol_training_sample(
+            DATA_DIR / "videos2caption.json",
+            DATA_DIR / "videos",
+        )
+    else:
+        if not prompt or not prompt.strip():
+            raise ValueError("--prompt is required when --video is given")
+        video_path, caption = video, prompt.strip()
+        if not video_path.is_file():
+            raise FileNotFoundError(f"Training video is missing at {video_path}")
     frames, waveform = load_training_media(video_path)
     pipeline_config = MiniMaxH3PipelineConfig()
     fastvideo_args = FastVideoArgs(
@@ -345,21 +359,34 @@ def main() -> None:
     audio_latents = encode_audio_latents(waveform, resolved_model_path, model_index, fastvideo_args)
     text_embedding = encode_text_embedding(caption, resolved_model_path, model_index, fastvideo_args)
     record = build_parquet_record(
-        file_name=TRAINING_VIDEO_NAME,
+        file_name=video_path.name,
         caption=caption,
         video_latents=video_latents,
         audio_latents=audio_latents,
         text_embedding=text_embedding,
     )
-    output_path = write_parquet(record, OUTPUT_DIR)
-    print(f"Wrote one Crush-Smol MiniMax H3 T2VA record to {output_path}")
+    output_path = write_parquet(record, output_dir)
+    print(f"Wrote one MiniMax H3 T2VA record to {output_path}")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--validate-only", action="store_true")
+    parser.add_argument("--video",
+                        type=Path,
+                        default=None,
+                        help="mp4 with a soundtrack (>=124 frames at 24 fps); "
+                        "default: the pinned Crush-Smol record")
+    parser.add_argument("--prompt", type=str, default=None, help="caption for --video")
+    parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR)
+    parser.add_argument("--model-path", type=Path, default=MODEL_PATH)
     cli_args = parser.parse_args()
     if cli_args.validate_only:
         validate_preprocessed_training_data()
     else:
-        main()
+        main(
+            video=cli_args.video,
+            prompt=cli_args.prompt,
+            output_dir=cli_args.output_dir,
+            model_path=cli_args.model_path,
+        )

fastvideo/pipelines/preprocess/preprocess_minimax_h3_text_only.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Encode a prompt list into MiniMax H3 text-only conditioning rows (data-free DMD2).
+
+Each line of ``--prompts-file`` is one complete H3 prompt document (for the
+VidProM-H3 set: ``integrated_multimodal_description: ... overall_soundscape:
+... non_diegetic_music: ...`` fed to the model verbatim as one string). Every
+prompt is tokenized raw (no chat template, ``add_special_tokens=False``) and
+encoded through ``MiniMaxH3ConditioningStage`` — the same Qwen3-VL layer-50
+path that produced the validated t2va overfit rows — then written as
+``pyarrow_schema_text_only`` records for ``rollout_mode: simulate`` training
+(``training.data.preprocessed_data_type: text_only``).
+
+Embeddings are stored as float32 to match the training collate, which decodes
+``text_embedding_bytes`` with a hard-coded ``np.float32``
+(``fastvideo/dataset/utils.py``). At ~300 tokens x 5120 dims that is ~6 MB per
+prompt — budget ~380 GB for the full 63k VidProM set.
+
+Sharding: ``--num-shards N`` splits the prompt list round-robin by line index;
+run one process per GPU with distinct ``--shard-index``/``CUDA_VISIBLE_DEVICES``
+(and a distinct ``MASTER_PORT`` — each process initializes a one-rank process
+group). Each shard writes ``<output-dir>/shard_XX/``; the training dataloader
+walks the directory tree, so pointing ``training.data.data_path`` at
+``--output-dir`` picks up every shard. Restarting a shard resumes after the
+rows already on disk (delete the shard directory to re-encode from scratch).
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import pyarrow.parquet as pq
+import torch
+
+from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, records_to_table)
+from fastvideo.dataset.dataloader.record_schema import text_only_record_creator
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_text_only
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.loader.component_loader import PipelineComponentLoader
+from fastvideo.pipelines import ForwardBatch
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MiniMaxH3ConditioningStage
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import MINIMAX_H3_KEYFRAMES_KEY
+from fastvideo.utils import verify_model_config_and_directory
+
+
+def _init_single_process_distributed(shard_index: int) -> None:
+    """Initialize the one-rank process groups required by component loaders.
+
+    Concurrent shards on one host must not share a rendezvous port, so the
+    default port is offset by the shard index (an explicit ``MASTER_PORT``
+    still wins).
+    """
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", str(29531 + shard_index))
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    from fastvideo.distributed import maybe_init_distributed_environment_and_model_parallel
+
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+
+
+def _load_component(
+    name: str,
+    model_path: Path,
+    model_index: dict[str, Any],
+    fastvideo_args: FastVideoArgs,
+) -> Any:
+    """Load one checkpoint component through the inference component registry."""
+    transformers_or_diffusers, _ = model_index[name][:2]
+    return PipelineComponentLoader.load_module(
+        module_name=name,
+        component_model_path=str(model_path / name),
+        transformers_or_diffusers=transformers_or_diffusers,
+        fastvideo_args=fastvideo_args,
+    )
+
+
+def _load_shard_prompts(prompts_file: Path, shard_index: int, num_shards: int) -> list[tuple[int, str]]:
+    """Return this shard's ``(global_line_index, prompt)`` pairs, round-robin by line."""
+    entries: list[tuple[int, str]] = []
+    with prompts_file.open(encoding="utf-8") as handle:
+        for index, line in enumerate(handle):
+            prompt = line.strip()
+            if prompt:
+                entries.append((index, prompt))
+    return entries[shard_index::num_shards]
+
+
+def _count_existing_rows(shard_dir: Path) -> int:
+    """Count rows already written under a shard directory (resume offset)."""
+    total = 0
+    for parquet_path in sorted(shard_dir.rglob("*.parquet")):
+        total += pq.ParquetFile(parquet_path).metadata.num_rows
+    return total
+
+
+def main(args: argparse.Namespace) -> None:
+    _init_single_process_distributed(args.shard_index)
+
+    model_path = args.model_path.resolve()
+    if not model_path.is_dir():
+        raise FileNotFoundError(f"MiniMax H3 model directory is missing at {model_path}")
+    model_index = verify_model_config_and_directory(str(model_path))
+
+    shard = _load_shard_prompts(args.prompts_file, args.shard_index, args.num_shards)
+    shard_dir = args.output_dir / f"shard_{args.shard_index:02d}"
+    already_done = _count_existing_rows(shard_dir) if shard_dir.is_dir() else 0
+    if already_done >= len(shard):
+        print(f"Shard {args.shard_index}/{args.num_shards}: all {len(shard)} rows already encoded")
+        return
+    todo = shard[already_done:]
+    if args.limit is not None:
+        todo = todo[:args.limit]
+    print(f"Shard {args.shard_index}/{args.num_shards}: {len(shard)} prompts total, "
+          f"{already_done} already on disk, encoding {len(todo)} now -> {shard_dir}")
+
+    fastvideo_args = FastVideoArgs(
+        model_path=str(model_path),
+        pipeline_config=MiniMaxH3PipelineConfig(),
+        num_gpus=1,
+        tp_size=1,
+        sp_size=1,
+        hsdp_shard_dim=1,
+        use_fsdp_inference=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+    )
+    print("Loading MiniMax H3 tokenizer, processor, and Qwen3-VL encoder")
+    tokenizer = _load_component("tokenizer", model_path, model_index, fastvideo_args)
+    processor = _load_component("processor", model_path, model_index, fastvideo_args)
+    conditioner = _load_component("text_encoder", model_path, model_index, fastvideo_args)
+    stage = MiniMaxH3ConditioningStage(
+        conditioner=conditioner,
+        tokenizer=tokenizer,
+        processor=processor,
+    )
+
+    writer = ParquetDatasetWriter(out_dir=str(shard_dir), samples_per_file=args.samples_per_file)
+    records: list[dict[str, Any]] = []
+    started = time.monotonic()
+    with torch.inference_mode():
+        for done, (global_index, prompt) in enumerate(todo, start=1):
+            batch = ForwardBatch(data_type="video", prompt=prompt)
+            batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = []
+            batch = stage.forward(batch, fastvideo_args)
+            if not batch.prompt_embeds:
+                raise RuntimeError(f"MiniMax H3 conditioning returned no embedding for line {global_index}")
+            # float32 to match the training collate's np.frombuffer dtype.
+            text_embedding = batch.prompt_embeds[0].squeeze(0).float().cpu().contiguous().numpy()
+            records.append(
+                text_only_record_creator(
+                    text_name=f"vidprom_{global_index:05d}",
+                    text_embedding=text_embedding,
+                    caption=prompt,
+                ))
+
+            if len(records) >= args.flush_every or done == len(todo):
+                writer.append_table(records_to_table(records, pyarrow_schema_text_only))
+                records = []
+                written = writer.flush(write_remainder=done == len(todo))
+                rate = done / (time.monotonic() - started)
+                remaining = (len(todo) - done) / rate if rate > 0 else float("inf")
+                print(f"[shard {args.shard_index}] {done}/{len(todo)} encoded "
+                      f"({rate:.2f} prompts/s, ~{remaining / 60:.0f} min left, flushed {written} rows)")
+
+    print(f"Shard {args.shard_index} complete: {already_done + len(todo)}/{len(shard)} rows in {shard_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--prompts-file", type=Path, required=True, help="one H3 prompt document per line")
+    parser.add_argument("--model-path", type=Path, required=True, help="MiniMax-H3 checkpoint directory")
+    parser.add_argument("--output-dir", type=Path, required=True, help="dataset root; shards write shard_XX/ under it")
+    parser.add_argument("--shard-index", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    parser.add_argument("--samples-per-file", type=int, default=64)
+    parser.add_argument("--flush-every", type=int, default=256, help="rows buffered between parquet flushes")
+    parser.add_argument("--limit", type=int, default=None, help="encode at most N prompts this run (smoke tests)")
+    cli_args = parser.parse_args()
+    if not 0 <= cli_args.shard_index < cli_args.num_shards:
+        parser.error(f"--shard-index {cli_args.shard_index} must be in [0, {cli_args.num_shards})")
+    main(cli_args)

fastvideo/platforms/cuda.py

@@ -184,6 +184,19 @@ class CudaPlatformBase(Platform):
                 raise ImportError("The Video Sparse Attention backend is not installed. "
                                   "To install it, please follow the instructions at: "
                                   "https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation ") from e
+        elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3:
+            try:
+                from fastvideo_kernel.block_sparse_attn_256 import (  # noqa: F401
+                    block_sparse_attn_256_bshd)
+
+                logger.info("Using MiniMax-H3 Video Sparse Attention backend.")
+
+                return "fastvideo.attention.backends.video_sparse_attn_h3.MiniMaxH3VSABackend"
+            except ImportError as e:
+                logger.error("Failed to import H3 Video Sparse Attention backend: %s", str(e))
+                raise ImportError("VIDEO_SPARSE_ATTN_H3 selected but fastvideo_kernel's block-sparse "
+                                  "kernels are unavailable. Install fastvideo-kernel or pick a different "
+                                  "FASTVIDEO_ATTENTION_BACKEND.") from e
         elif selected_backend == AttentionBackendEnum.BSA_ATTN:
             try:
                 from fastvideo.attention.backends.bsa_attn import (  # noqa: F401

fastvideo/platforms/interface.py

@@ -18,6 +18,7 @@ class AttentionBackendEnum(enum.Enum):
     ATTN_QAT_INFER = enum.auto()
     ATTN_QAT_TRAIN = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()
+    VIDEO_SPARSE_ATTN_H3 = enum.auto()
     BSA_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()
     SLA_ATTN = enum.auto()

fastvideo/tests/api/test_extra_overrides_routing.py

@@ -36,6 +36,10 @@ def test_passthrough_keys_cover_ltx2_audio_conditioning() -> None:
         "ltx2_audio_denoise_mask",
         "audio_num_frames",
         "video_position_offset_sec",
+        # MiniMax-H3 VSA per-request knobs (consumed by MiniMaxH3DenoisingStage)
+        "vsa_mode",
+        "vsa_dense_first_n_steps",
+        "vsa_dense_layers",
     }
     assert set(_BATCH_EXTRA_PASSTHROUGH_KEYS) == expected
 

fastvideo/tests/attention/test_vsa_h3_metadata.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU checks for the VSA-H3 backend: tiling geometry, mask policy, and
+end-to-end equivalence against dense SDPA through a token-level mask
+reference. The same reference doubles as the GPU kernel parity oracle."""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, MiniMaxH3VSAImpl,
+                                                               MiniMaxH3VSAMetadataBuilder, _build_block_mask,
+                                                               _pool_tiles, token_tile_and_valid)
+
+_720P = dict(raw_latent_shape=(30, 44, 80), patch_size=(1, 2, 2), prefix_segments=(512, 1760, 400))
+_TINY = dict(raw_latent_shape=(8, 8, 12), patch_size=(1, 2, 2), prefix_segments=(7, 5, 3))
+
+_CPU = torch.device("cpu")
+
+
+def _build(spec, sparsity=0.0, device=_CPU):
+    return MiniMaxH3VSAMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=spec["raw_latent_shape"],
+        patch_size=spec["patch_size"],
+        VSA_sparsity=sparsity,
+        prefix_segments=spec["prefix_segments"],
+        device=device,
+    )
+
+
+def _impl():
+    return MiniMaxH3VSAImpl(num_heads=2, head_size=8, causal=False, softmax_scale=8**-0.5)
+
+
+def reference_sparse_attention(query, key, value, mask, meta):
+    """Token-level oracle: SDPA over the padded tile buffer with the block
+    mask expanded to tokens. query/key/value: tiled [B, S_pad, H, D]."""
+    token_tile, token_valid = token_tile_and_valid(meta.variable_block_sizes)
+    out = torch.empty_like(query)
+    for b in range(query.shape[0]):
+        for h in range(query.shape[2]):
+            allow = mask[b, h][token_tile][:, token_tile] & token_valid[None, :]
+            bias = torch.zeros(allow.shape, dtype=query.dtype, device=query.device)
+            bias.masked_fill_(~allow, float("-inf"))
+            out[b, :, h] = F.scaled_dot_product_attention(
+                query[b, :, h][None],
+                key[b, :, h][None],
+                value[b, :, h][None],
+                attn_mask=bias[None],
+            )[0]
+    return out
+
+
+def test_geometry_720p():
+    meta = _build(_720P)
+    seq = meta.total_seq_length
+    assert seq == 512 + 1760 + 400 + 26400
+    assert meta.num_prefix_tiles == 2 + 7 + 2
+    assert meta.num_video_tiles == 8 * 3 * 5
+    assert int(meta.variable_block_sizes.sum()) == seq
+    # (permutation coverage of [0, seq) is implied by the roundtrip below:
+    # untile_combined_index scatters seq distinct rows and recovers all of x)
+    # segment purity: no prefix tile straddles a segment boundary
+    boundaries = [512, 512 + 1760, 512 + 1760 + 400]
+    start = 0
+    for size in meta.variable_block_sizes[:meta.num_prefix_tiles].tolist():
+        end = start + size
+        assert all(not (start < b < end) for b in boundaries), (start, end)
+        start = end
+    # untile(tile(x)) == x
+    x = torch.randn(1, seq, 2, 4)
+    buf = _impl().tile(x, meta)
+    assert buf.shape[1] == meta.variable_block_sizes.numel() * _TILE_ELEMS
+    assert torch.equal(buf[:, meta.untile_combined_index], x)
+
+
+def test_mask_policy():
+    meta = _build(_720P, sparsity=0.9)
+    n = meta.num_prefix_tiles + meta.num_video_tiles
+    P, V = meta.num_prefix_tiles, meta.num_video_tiles
+    k_vid = math.ceil(0.1 * V)
+    scores = torch.randn(1, 2, n, n)
+
+    exempt = _build_block_mask(scores, P, V, 0.9, exempt=True)
+    assert exempt[:, :, :P].all(), "prefix queries must be dense"
+    assert exempt[..., :P].all(), "prefix keys must be visible to every query"
+    assert (exempt[:, :, P:, P:].sum(-1) == k_vid).all(), "video rows select exactly k_vid video tiles"
+
+    compete = _build_block_mask(scores, P, V, 0.9, exempt=False)
+    assert compete[:, :, :P].all()
+    assert (compete[:, :, P:].sum(-1) == min(k_vid + P, n)).all(), "budget-matched top-k"
+
+    dense = _build_block_mask(scores, P, V, 0.0, exempt=True)
+    assert dense.all(), "sparsity 0 must select everything"
+
+
+def test_sparsity_zero_matches_dense_sdpa():
+    torch.manual_seed(0)
+    meta = _build(_TINY)
+    seq = meta.total_seq_length
+    q, k, v = (torch.randn(1, seq, 2, 8) for _ in range(3))
+    impl = _impl()
+    tq, tk, tv = (impl.tile(t, meta).clone() for t in (q, k, v))
+
+    scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes),
+                          _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1))
+    mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.0, exempt=True)
+    sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta)
+
+    dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
+    assert torch.allclose(sparse_out, dense_out, atol=1e-5), (sparse_out - dense_out).abs().max()
+
+
+def test_prefix_queries_stay_dense_at_high_sparsity():
+    torch.manual_seed(1)
+    meta = _build(_TINY, sparsity=0.75)
+    seq = meta.total_seq_length
+    prefix_len = sum(_TINY["prefix_segments"])
+    q, k, v = (torch.randn(1, seq, 2, 8) for _ in range(3))
+    impl = _impl()
+    tq, tk, tv = (impl.tile(t, meta).clone() for t in (q, k, v))
+
+    scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes),
+                          _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1))
+    for exempt in (True, False):
+        mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.75, exempt=exempt)
+        sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta)
+        dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2),
+                                                   v.transpose(1, 2)).transpose(1, 2)
+        assert torch.allclose(sparse_out[:, :prefix_len], dense_out[:, :prefix_len], atol=1e-5)
+        assert not torch.allclose(sparse_out[:, prefix_len:], dense_out[:, prefix_len:], atol=1e-5), \
+            "video rows should actually be sparse at 75%"
+
+
+def test_tile_under_grad_does_not_reuse_shared_buffer():
+    """Training forwards need fresh tile buffers; in-place reuse of the shared
+    holder would trip autograd's saved-tensor version check at backward."""
+    meta = _build(_TINY)
+    impl = _impl()
+    seq = meta.total_seq_length
+
+    x1 = torch.randn(1, seq, 2, 8, requires_grad=True)
+    x2 = torch.randn(1, seq, 2, 8, requires_grad=True)
+    buf1 = impl.tile(x1, meta)
+    saved = (buf1 * buf1).sum()  # saves buf1 for backward, like the kernel
+    buf2 = impl.tile(x2, meta)
+    assert buf1 is not buf2
+    saved.backward()  # raises "modified by an inplace operation" on reuse
+    assert x1.grad is not None and x2.grad is None
+
+    # no-grad paths (inference / rollout) keep the single-buffer reuse
+    with torch.no_grad():
+        y1 = impl.tile(x1.detach(), meta)
+        y2 = impl.tile(x2.detach(), meta)
+    assert y1 is y2
+
+
+if __name__ == "__main__":
+    test_geometry_720p()
+    test_mask_policy()
+    test_sparsity_zero_matches_dense_sdpa()
+    test_prefix_queries_stay_dense_at_high_sparsity()
+    test_tile_under_grad_does_not_reuse_shared_buffer()
+    print("all VSA-H3 CPU checks passed")

fastvideo/tests/loader/test_shard_cache.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contract tests for the sharded base-weight cache.
+
+Runs on a single-process gloo group with a (1, 1) CPU device mesh: DTensor
+round-trip through write/load, the FQN reconciliation matrix (allowed
+zero-init params, disallowed extras, shape mismatches), and the
+never-fail-the-run contract.
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import Replicate, Shard, distribute_tensor
+
+from fastvideo.models.loader.shard_cache import (
+    ShardCacheContext,
+    try_load_from_shard_cache,
+    write_shard_cache,
+)
+
+
+@pytest.fixture(scope="module")
+def cpu_mesh():
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29581")
+    if not dist.is_initialized():
+        dist.init_process_group("gloo", rank=0, world_size=1)
+    return init_device_mesh("cpu", (1, 1), mesh_dim_names=("replicate", "shard"))
+
+
+def _make_model(cpu_mesh, *, extra_param: str | None = None, weight_rows: int = 8) -> nn.Module:
+    model = nn.Module()
+    placements = (Replicate(), Shard(0))
+    weight = distribute_tensor(torch.randn(weight_rows, 4), cpu_mesh, placements)
+    bias = distribute_tensor(torch.randn(weight_rows), cpu_mesh, placements)
+    model.register_parameter("weight", nn.Parameter(weight))
+    model.register_parameter("bias", nn.Parameter(bias))
+    model.register_buffer("scale", torch.full((1, ), 2.0))
+    if extra_param is not None:
+        extra = distribute_tensor(torch.randn(4, 4), cpu_mesh, placements)
+        model.register_parameter(extra_param.replace(".", "_"), nn.Parameter(extra))
+        # register under the dotted name via a child module for realism
+    model.reverse_param_names_mapping = {"weight": ("hf.weight", None, None)}
+    return model
+
+
+def _ctx(tmp_path) -> ShardCacheContext:
+    return ShardCacheContext(entry_dir=tmp_path / "entry", key="testkey", shard_index=0, num_shards=1, is_writer=True)
+
+
+def test_round_trip_restores_tensors_and_reverse_mapping(cpu_mesh, tmp_path):
+    src = _make_model(cpu_mesh)
+    ctx = _ctx(tmp_path)
+    write_shard_cache(src, ctx)
+
+    dst = _make_model(cpu_mesh)
+    with torch.no_grad():
+        dst.weight.mul_(0)
+        dst.bias.mul_(0)
+    dst.reverse_param_names_mapping = {}
+    assert try_load_from_shard_cache(dst, ctx, torch.device("cpu"))
+    assert torch.equal(dst.weight.to_local(), src.weight.to_local())
+    assert torch.equal(dst.bias.to_local(), src.bias.to_local())
+    assert torch.equal(dst.scale, src.scale)
+    assert dst.reverse_param_names_mapping == {"weight": ("hf.weight", None, None)}
+
+
+def test_allowed_new_param_zero_inits_on_hit(cpu_mesh, tmp_path):
+    src = _make_model(cpu_mesh)
+    ctx = _ctx(tmp_path)
+    write_shard_cache(src, ctx)
+
+    dst = _make_model(cpu_mesh)
+    gate = distribute_tensor(torch.randn(4, 4), cpu_mesh, (Replicate(), Shard(0)))
+    dst.register_parameter("to_gate_compress", nn.Parameter(gate))
+    assert try_load_from_shard_cache(dst, ctx, torch.device("cpu"))
+    assert torch.equal(dst.to_gate_compress.to_local(), torch.zeros(4, 4))
+
+
+def test_disallowed_missing_param_misses_without_mutation(cpu_mesh, tmp_path):
+    src = _make_model(cpu_mesh)
+    ctx = _ctx(tmp_path)
+    write_shard_cache(src, ctx)
+
+    dst = _make_model(cpu_mesh)
+    mystery = distribute_tensor(torch.randn(4, 4), cpu_mesh, (Replicate(), Shard(0)))
+    dst.register_parameter("mystery", nn.Parameter(mystery))
+    before = dst.weight.to_local().clone()
+    assert not try_load_from_shard_cache(dst, ctx, torch.device("cpu"))
+    assert torch.equal(dst.weight.to_local(), before)
+
+
+def test_shape_mismatch_misses(cpu_mesh, tmp_path):
+    src = _make_model(cpu_mesh)
+    ctx = _ctx(tmp_path)
+    write_shard_cache(src, ctx)
+
+    dst = _make_model(cpu_mesh, weight_rows=9)
+    assert not try_load_from_shard_cache(dst, ctx, torch.device("cpu"))
+
+
+def test_missing_entry_misses_cleanly(cpu_mesh, tmp_path):
+    dst = _make_model(cpu_mesh)
+    ctx = ShardCacheContext(entry_dir=tmp_path / "absent", key="k2", shard_index=0, num_shards=1, is_writer=True)
+    assert not try_load_from_shard_cache(dst, ctx, torch.device("cpu"))

fastvideo/tests/train/fixtures/minimax_h3_dmd2_min.yaml

@@ -0,0 +1,61 @@
+# Minimal H3 DMD2 trio and method configuration for CPU contract tests.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: false
+  critic:
+    _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel
+    init_from: MiniMaxAI/MiniMax-H3
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: data_latent
+  generator_update_interval: 1
+  real_score_guidance_scale: 3.5
+  dmd_denoising_steps: [1000, 757, 522]
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+  cfg_uncond:
+    text: zero
+  fake_score_learning_rate: 1.0e-3
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 1
+  data:
+    data_path: /tmp/minimax_h3_t2va
+    preprocessed_data_type: t2va
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    # Tiny geometry: [1, 24, 2, 4, 4] video latents, [1, 2, 32, 8] audio.
+    num_latent_t: 2
+    num_height: 64
+    num_width: 64
+    num_frames: 5
+  optimizer:
+    learning_rate: 1.0e-3
+    betas: [0.0, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+  loop:
+    max_train_steps: 4
+  model:
+    precondition_outputs: false
+  dit_precision: bf16
+
+callbacks: {}
+pipeline: {}

fastvideo/tests/train/methods/test_minimax_h3_dmd2.py

@@ -0,0 +1,509 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contract tests for MiniMax H3 DMD2 distillation.
+
+Covers the packed dual-modality adapter (MiniMaxH3DMDModel) and one full
+DMD2Method.single_train_step on a tiny CPU trio: student rollout, critic
+flow-matching loss, generator DMD loss, both backwards, both optimizers.
+"""
+
+import math
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import yaml
+
+from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata
+from fastvideo.forward_context import get_forward_context
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
+from fastvideo.train.models.minimax_h3 import MiniMaxH3DMDModel, MiniMaxH3Model
+from fastvideo.train.models.minimax_h3.minimax_h3 import shift_noise_amount
+from fastvideo.train.utils.config import load_run_config
+
+_FIXTURE = Path(__file__).resolve().parent.parent / "fixtures" / "minimax_h3_dmd2_min.yaml"
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+_EXPERIMENT_CONFIG = _REPO_ROOT / "examples/train/configs/distribution_matching/minimax_h3/dmd2_t2va.yaml"
+_VSA_OVERFIT_CONFIG = (_REPO_ROOT / "examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml")
+
+# Fixture geometry: video latents [1, 24, 2, 4, 4] and audio latents
+# [1, 2, 32, 8]; the packed adapter stores video-major [1, T, C, H, W].
+_VIDEO_SHAPE = (1, 2, 24, 4, 4)
+_AUDIO_SHAPE = (1, 2, 32, 8)
+_PACKED_NUMEL = math.prod(_VIDEO_SHAPE) + math.prod(_AUDIO_SHAPE)
+
+
+class _TinyJointTransformer(torch.nn.Module):
+    """Scale packed H3 rows with one trainable parameter."""
+
+    patch_size = (1, 2, 2)
+
+    def __init__(self, scale: float = 1.0) -> None:
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.tensor(scale))
+        self.last_encoder_hidden_states: torch.Tensor | None = None
+        self.last_attn_metadata = None
+
+    def forward(self, **kwargs):
+        self.last_encoder_hidden_states = kwargs["encoder_hidden_states"]
+        self.last_attn_metadata = get_forward_context().attn_metadata
+        return (
+            kwargs["hidden_states"] * self.scale,
+            kwargs["audio_hidden_states"] * self.scale,
+        )
+
+
+def _make_model(
+    monkeypatch: pytest.MonkeyPatch,
+    training_config,
+    *,
+    trainable: bool = True,
+    scale: float = 1.0,
+) -> MiniMaxH3DMDModel:
+    monkeypatch.setattr(MiniMaxH3Model, "device", property(lambda _self: torch.device("cpu")))
+    model = MiniMaxH3DMDModel.__new__(MiniMaxH3DMDModel)
+    model._trainable = trainable
+    model.transformer = _TinyJointTransformer(scale)
+    model.training_config = training_config
+    model.sp_group = None
+    model.attention_backend = None
+    return model
+
+
+def _tiny_training_config():
+    return SimpleNamespace(
+        data=SimpleNamespace(
+            num_latent_t=2,
+            num_frames=5,
+            num_height=64,
+            num_width=64,
+        ),
+        distributed=SimpleNamespace(sp_size=1),
+        vsa_sparsity=0.0,
+    )
+
+
+def _raw_batch(seed: int = 1) -> dict[str, torch.Tensor]:
+    generator = torch.Generator().manual_seed(seed)
+    return {
+        "vae_latent": torch.randn(1, 24, 2, 4, 4, generator=generator),
+        "audio_latent": torch.randn(1, 2, 32, 8, generator=generator),
+        "text_embedding": torch.randn(1, 4, 5120, generator=generator),
+        "text_attention_mask": torch.tensor([[1, 1, 0, 0]], dtype=torch.float32),
+    }
+
+
+def _build_method(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    rollout_mode: str,
+    generator_update_interval: int = 1,
+) -> DMD2Method:
+    config = load_run_config(str(_FIXTURE))
+    config.method["rollout_mode"] = rollout_mode
+    config.method["generator_update_interval"] = generator_update_interval
+    # Distinct role scales keep the critic-vs-teacher DMD gradient non-zero.
+    student = _make_model(monkeypatch, config.training, scale=1.0)
+    teacher = _make_model(monkeypatch, config.training, trainable=False, scale=0.5)
+    critic = _make_model(monkeypatch, config.training, scale=0.25)
+    student.init_preprocessors = lambda training_config: None
+    method = DMD2Method(
+        cfg=config,
+        role_models={
+            "student": student,
+            "teacher": teacher,
+            "critic": critic,
+        },
+    )
+    method.cuda_generator = torch.Generator(device="cpu").manual_seed(0)
+    return method
+
+
+# ----------------------------------------------------------------------
+# Core gate: one full DMD2 train step on CPU
+# ----------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("rollout_mode", ["data_latent", "simulate"])
+def test_dmd2_single_train_step_updates_student_and_critic(
+    monkeypatch: pytest.MonkeyPatch,
+    rollout_mode: str,
+) -> None:
+    """Run rollout, critic loss, generator loss, both backwards and steps."""
+    method = _build_method(monkeypatch, rollout_mode=rollout_mode)
+    student = method.student
+    teacher = method.teacher
+    critic = method.critic
+
+    loss_map, outputs, metrics = method.single_train_step(_raw_batch(), iteration=0)
+
+    assert metrics["update_student"] == 1.0
+    for key in ("total_loss", "generator_loss", "fake_score_loss"):
+        assert torch.isfinite(loss_map[key]), key
+    assert loss_map["generator_loss"].item() > 0.0
+    assert loss_map["fake_score_loss"].item() > 0.0
+
+    method.backward(loss_map, outputs)
+    assert student.transformer.scale.grad is not None
+    assert torch.isfinite(student.transformer.scale.grad)
+    assert critic.transformer.scale.grad is not None
+    assert torch.isfinite(critic.transformer.scale.grad)
+    assert teacher.transformer.scale.grad is None
+
+    student_before = student.transformer.scale.detach().clone()
+    critic_before = critic.transformer.scale.detach().clone()
+    method.optimizers_schedulers_step(0)
+    assert student.transformer.scale.detach() != student_before
+    assert critic.transformer.scale.detach() != critic_before
+
+
+def test_dmd2_generator_update_interval_gates_student(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Off-interval iterations train the critic only."""
+    method = _build_method(
+        monkeypatch,
+        rollout_mode="data_latent",
+        generator_update_interval=5,
+    )
+
+    loss_map, _outputs, metrics = method.single_train_step(_raw_batch(), iteration=1)
+
+    assert metrics["update_student"] == 0.0
+    assert loss_map["generator_loss"].item() == 0.0
+    assert loss_map["fake_score_loss"].item() > 0.0
+    assert method.get_optimizers(1) == [method._critic_optimizer]
+
+
+# ----------------------------------------------------------------------
+# Packed dual-modality adapter units
+# ----------------------------------------------------------------------
+
+
+def test_packed_adapter_roundtrip_and_prepare_batch(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify pack/unpack inversion and packed clean latents in the batch."""
+    model = _make_model(monkeypatch, _tiny_training_config())
+    video = torch.randn(_VIDEO_SHAPE)
+    audio = torch.randn(_AUDIO_SHAPE)
+
+    packed = model.pack_latents(video, audio)
+    assert packed.shape == (1, _PACKED_NUMEL)
+    video_out, audio_out = model.unpack_latents(packed)
+    torch.testing.assert_close(video_out, video)
+    torch.testing.assert_close(audio_out, audio)
+
+    raw_batch = _raw_batch()
+    batch = model.prepare_batch(
+        raw_batch,
+        generator=torch.Generator().manual_seed(7),
+    )
+    assert batch.latents.shape == (1, _PACKED_NUMEL)
+    video_clean, audio_clean = model.unpack_latents(batch.latents)
+    torch.testing.assert_close(
+        video_clean,
+        raw_batch["vae_latent"].permute(0, 2, 1, 3, 4).to(torch.bfloat16),
+    )
+    torch.testing.assert_close(audio_clean, batch.audio_latents)
+
+
+def test_packed_add_noise_applies_modality_shifts(monkeypatch: pytest.MonkeyPatch) -> None:
+    """One shared base timestep must map to two shifted noise amounts."""
+    model = _make_model(monkeypatch, _tiny_training_config())
+    clean = torch.ones(1, _PACKED_NUMEL)
+    noise = torch.zeros(1, _PACKED_NUMEL)
+
+    torch.testing.assert_close(
+        model.add_noise(clean, noise, torch.tensor([0])),
+        clean,
+    )
+    torch.testing.assert_close(
+        model.add_noise(clean, noise, torch.tensor([1000])),
+        noise,
+    )
+
+    mixed = model.add_noise(clean, noise, torch.tensor([500]))
+    video_mixed, audio_mixed = model.unpack_latents(mixed)
+    base = torch.tensor([0.5])
+    torch.testing.assert_close(
+        video_mixed,
+        torch.full(_VIDEO_SHAPE, float(1.0 - shift_noise_amount(base, 12.0))),
+    )
+    torch.testing.assert_close(
+        audio_mixed,
+        torch.full(_AUDIO_SHAPE, float(1.0 - shift_noise_amount(base, 3.0))),
+    )
+
+
+def test_packed_predict_noise_plumbs_timesteps_and_tolerates_vsa(monkeypatch: pytest.MonkeyPatch, ) -> None:
+    """Explicit method timesteps must rewrite both modality clean-times."""
+    model = _make_model(monkeypatch, _tiny_training_config())
+    batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7))
+    noisy = torch.randn(1, _PACKED_NUMEL).to(torch.bfloat16)
+    timestep = torch.tensor([757], dtype=torch.long)
+
+    # attn_kind="vsa" must silently mean dense (both metadata views are None).
+    prediction = model.predict_noise(
+        noisy,
+        timestep,
+        batch,
+        conditional=True,
+        attn_kind="vsa",
+    )
+
+    base = torch.tensor([0.757])
+    torch.testing.assert_close(
+        batch.timesteps,
+        1.0 - shift_noise_amount(base, 12.0),
+    )
+    torch.testing.assert_close(
+        batch.audio_timesteps,
+        1.0 - shift_noise_amount(base, 3.0),
+    )
+    # The unit-scale transformer echoes packed rows, and the H3 wrapper
+    # negates them into noise-minus-clean form.
+    torch.testing.assert_close(prediction, -noisy)
+
+    x0 = model.predict_x0(noisy, timestep, batch, conditional=True)
+    noisy_video, noisy_audio = model.unpack_latents(noisy)
+    sigma_video = shift_noise_amount(base, 12.0).to(torch.bfloat16)
+    sigma_audio = shift_noise_amount(base, 3.0).to(torch.bfloat16)
+    expected = model.pack_latents(
+        noisy_video + sigma_video * noisy_video,
+        noisy_audio + sigma_audio * noisy_audio,
+    )
+    torch.testing.assert_close(x0, expected)
+
+
+def test_uncond_forward_zeroes_text_and_guards_policies(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Teacher-CFG unconditional forwards zero text; other policies fail fast."""
+    model = _make_model(monkeypatch, _tiny_training_config())
+    batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7))
+    noisy = torch.randn(1, _PACKED_NUMEL).to(torch.bfloat16)
+    timestep = torch.tensor([500], dtype=torch.long)
+
+    model.predict_noise(
+        noisy,
+        timestep,
+        batch,
+        conditional=False,
+        cfg_uncond={"text": "zero"},
+    )
+    assert torch.all(model.transformer.last_encoder_hidden_states == 0)
+
+    model.predict_noise(
+        noisy,
+        timestep,
+        batch,
+        conditional=True,
+        cfg_uncond={"text": "zero"},
+    )
+    assert torch.any(model.transformer.last_encoder_hidden_states != 0)
+
+    with pytest.raises(ValueError, match="cfg_uncond"):
+        model.predict_noise(noisy, timestep, batch, conditional=False)
+    with pytest.raises(ValueError, match="negative-prompt"):
+        model.set_requires_negative_conditioning(True)
+    model.set_requires_negative_conditioning(False)
+
+
+# ----------------------------------------------------------------------
+# VSA-H3 wiring
+# ----------------------------------------------------------------------
+
+
+def test_prepare_batch_builds_vsa_h3_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The VSA-H3 role gets real packed-sequence metadata; dense view stays None."""
+    tc = _tiny_training_config()
+    tc.vsa_sparsity = 0.35
+    model = _make_model(monkeypatch, tc)
+    model.attention_backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
+
+    batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7))
+
+    meta = batch.attn_metadata_vsa
+    assert isinstance(meta, MiniMaxH3VSAMetadata)
+    assert batch.attn_metadata is None
+    assert meta.VSA_sparsity == pytest.approx(0.35)
+    # Packed layout: 2 text rows | 0 condition rows | 16 stereo audio rows |
+    # 8 video rows ([1, 24, 2, 4, 4] latents at patch (1, 2, 2)).
+    assert meta.total_seq_length == 26
+    assert meta.num_prefix_tiles == 2
+    assert meta.num_video_tiles == 1
+    assert meta.variable_block_sizes.tolist() == [2, 16, 8]
+    assert int(meta.variable_block_sizes.sum()) == meta.total_seq_length
+
+
+def test_predict_noise_routes_vsa_metadata_by_attn_kind(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Student "vsa" forwards see the VSA metadata; "dense" forwards see None."""
+    model = _make_model(monkeypatch, _tiny_training_config())
+    model.attention_backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
+    batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7))
+    noisy = torch.randn(1, _PACKED_NUMEL).to(torch.bfloat16)
+    timestep = torch.tensor([757], dtype=torch.long)
+
+    model.predict_noise(noisy, timestep, batch, conditional=True, attn_kind="vsa")
+    assert model.transformer.last_attn_metadata is batch.attn_metadata_vsa
+    assert isinstance(model.transformer.last_attn_metadata, MiniMaxH3VSAMetadata)
+
+    model.predict_noise(noisy, timestep, batch, conditional=True, attn_kind="dense")
+    assert model.transformer.last_attn_metadata is None
+
+
+def _ctor_training_config() -> SimpleNamespace:
+    """The minimum surface MiniMaxH3Model.__init__ reads from TrainingConfig."""
+    return SimpleNamespace(
+        pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(uniform_parameter_dtype=False)),
+        data=SimpleNamespace(
+            train_batch_size=1,
+            training_cfg_rate=0.0,
+            preprocessed_data_type="t2va",
+        ),
+        model=SimpleNamespace(enable_gradient_checkpointing_type=None),
+    )
+
+
+def test_per_role_attention_backend_override_resolves(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Each role's backend reaches the loader; unsupported backends fail fast."""
+    captured: dict[str, AttentionBackendEnum | None] = {}
+
+    def _fake_load(**kwargs):
+        captured[kwargs["model_path"]] = kwargs["attention_backend"]
+        return _TinyJointTransformer()
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.minimax_h3.minimax_h3.load_module_from_path",
+        _fake_load,
+    )
+
+    student = MiniMaxH3DMDModel(
+        init_from="role/student",
+        training_config=_ctor_training_config(),
+        trainable=True,
+        attention_backend="VIDEO_SPARSE_ATTN_H3",
+    )
+    teacher = MiniMaxH3DMDModel(
+        init_from="role/teacher",
+        training_config=_ctor_training_config(),
+        trainable=False,
+        attention_backend="FLASH_ATTN",
+    )
+
+    assert student.attention_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
+    assert teacher.attention_backend is AttentionBackendEnum.FLASH_ATTN
+    # load_module_from_path turns this request into the construction scope
+    # that binds the backend to the transformer's attention layers.
+    assert captured["role/student"] is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
+    assert captured["role/teacher"] is AttentionBackendEnum.FLASH_ATTN
+    assert not any(p.requires_grad for p in teacher.transformer.parameters())
+
+    with pytest.raises(ValueError, match="supports the attention backends"):
+        MiniMaxH3DMDModel(
+            init_from="role/bad",
+            training_config=_ctor_training_config(),
+            attention_backend="VIDEO_SPARSE_ATTN",
+        )
+
+
+# ----------------------------------------------------------------------
+# Config contracts
+# ----------------------------------------------------------------------
+
+
+def test_vsa0_overfit_config_pins_roles_and_experiment() -> None:
+    """The overfit config runs student VSA-H3 at sparsity 0, dense FA roles."""
+    config = yaml.safe_load(_VSA_OVERFIT_CONFIG.read_text())
+    models, method, training = config["models"], config["method"], config["training"]
+
+    assert models["student"]["attention_backend"] == "VIDEO_SPARSE_ATTN_H3"
+    assert models["teacher"]["attention_backend"] == "FLASH_ATTN"
+    assert models["critic"]["attention_backend"] == "FLASH_ATTN"
+    assert models["teacher"]["trainable"] is False
+    assert training["vsa"]["sparsity"] == 0.0
+    assert method["rollout_mode"] == "data_latent"
+    assert method["generator_update_interval"] == 5
+    assert method["dmd_denoising_steps"] == [1000, 757, 522]
+    assert training["data"]["train_batch_size"] == 1
+    assert training["data"]["data_path"] == "/mnt/h3-dmd2-overfit/data"
+    assert training["data"]["num_height"] == 768
+    assert training["data"]["num_width"] == 1344
+    assert training["data"]["num_frames"] == 124
+    assert training["loop"]["max_train_steps"] == 2000
+    assert training["tracker"]["project_name"] == "h3-dmd2-vsa"
+
+
+def test_h3_dmd2_fixture_resolves_trio_contract() -> None:
+    """The fixture must wire the H3 DMD trio through the modular builder path."""
+    config = load_run_config(str(_FIXTURE))
+
+    for role in ("student", "teacher", "critic"):
+        assert config.models[role]["_target_"] == ("fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel")
+    assert config.models["teacher"]["trainable"] is False
+    assert config.method["_target_"] == ("fastvideo.train.methods.distribution_matching.dmd2.DMD2Method")
+    assert config.training.data.preprocessed_data_type == "t2va"
+
+
+def test_h3_dmd2_experiment_config_mirrors_wan_recipe() -> None:
+    """The example config keeps the studio DMD2 defaults and the H3 contract."""
+    config = yaml.safe_load(_EXPERIMENT_CONFIG.read_text())
+    method = config["method"]
+
+    for role in ("student", "teacher", "critic"):
+        assert config["models"][role]["_target_"] == ("fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel")
+    assert config["models"]["teacher"]["trainable"] is False
+    assert config["models"]["critic"]["trainable"] is True
+    assert method["_target_"] == ("fastvideo.train.methods.distribution_matching.dmd2.DMD2Method")
+    assert method["rollout_mode"] == "data_latent"
+    assert method["generator_update_interval"] == 5
+    assert method["dmd_denoising_steps"] == [1000, 757, 522]
+    assert method["cfg_uncond"] == {"text": "zero"}
+    assert method["fake_score_learning_rate"] == 8.0e-6
+    assert method["fake_score_betas"] == [0.0, 0.999]
+    assert method["fake_score_lr_scheduler"] == "constant"
+    assert config["training"]["data"]["preprocessed_data_type"] == "t2va"
+    assert config["training"]["data"]["train_batch_size"] == 1
+    assert config["training"]["data"]["training_cfg_rate"] == 0.0
+
+
+def test_validation_dmd_sigmas_match_training_noise_amounts() -> None:
+    """``pipeline_config.dmd_denoising_steps`` replays the trained jump points.
+
+    The H3 denoising stage normalizes the method's integer steps to base time
+    and lets each scheduler apply its own shift; the resulting clean-times must
+    match ``1 - shift_noise_amount(base)`` — the exact noising the packed DMD
+    adapter applies during training rollouts — with one forward per step.
+    """
+    from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+
+    steps = [1000, 757, 522]
+    base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32)
+    video = MiniMaxH3Scheduler(shift=12.0)
+    audio = MiniMaxH3Scheduler(shift=3.0)
+    video.set_timesteps(sigmas=video.shift_sigmas(base))
+    audio.set_timesteps(sigmas=audio.shift_sigmas(base))
+
+    assert video.num_inference_steps == len(steps)
+    assert audio.num_inference_steps == len(steps)
+    for index, step in enumerate(steps):
+        base_step = torch.tensor([step / 1000.0])
+        assert video.timesteps[index].item() == pytest.approx(1.0 - shift_noise_amount(base_step, 12.0).item())
+        assert audio.timesteps[index].item() == pytest.approx(1.0 - shift_noise_amount(base_step, 3.0).item())
+
+
+def test_validation_callback_injects_method_denoising_steps() -> None:
+    """The callback copies the trained step list onto the validation config."""
+    from fastvideo.train.callbacks.validation import ValidationCallback
+
+    callback = ValidationCallback.__new__(ValidationCallback)
+    callback.method = SimpleNamespace(method_config={"dmd_denoising_steps": [1000, 757, 522]})
+
+    config = SimpleNamespace(dmd_denoising_steps=None)
+    callback._inject_method_denoising_steps(config)
+    assert config.dmd_denoising_steps == [1000, 757, 522]
+
+    explicit = SimpleNamespace(dmd_denoising_steps=[1000, 500])
+    callback._inject_method_denoising_steps(explicit)
+    assert explicit.dmd_denoising_steps == [1000, 500]
+
+    callback.method = SimpleNamespace(method_config={"dmd_denoising_steps": [1000, 757], "warp_denoising_step": True})
+    warped = SimpleNamespace(dmd_denoising_steps=None)
+    callback._inject_method_denoising_steps(warped)
+    assert warped.dmd_denoising_steps is None

fastvideo/tests/train/utils/test_tracking.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU tests for the modular trainer's wandb tracking guard.
+
+The Trainer logs each step's loss dict through ``build_tracker``'s tracker on
+rank 0; these tests pin the graceful no-op path (wandb missing / no
+credentials) and the logged loss-dict shape with a monkeypatched wandb.
+"""
+
+import builtins
+import sys
+from types import SimpleNamespace
+
+from fastvideo.train.utils import tracking
+from fastvideo.train.utils.training_config import (
+    CheckpointConfig,
+    TrackerConfig,
+)
+from fastvideo.training.trackers import DummyTracker, WandbTracker
+
+
+class _FakeRun:
+
+    def __init__(self, **init_kwargs):
+        self.init_kwargs = init_kwargs
+        self.logged = []
+
+    def log(self, metrics, step=None):
+        self.logged.append((metrics, step))
+
+    def finish(self):
+        pass
+
+
+def _fake_wandb():
+    runs = []
+
+    def _init(**kwargs):
+        run = _FakeRun(**kwargs)
+        runs.append(run)
+        return run
+
+    return SimpleNamespace(init=_init, runs=runs, api=SimpleNamespace(api_key="key"))
+
+
+def test_wandb_usable_false_when_not_importable(monkeypatch) -> None:
+    real_import = builtins.__import__
+
+    def _blocked(name, *args, **kwargs):
+        if name == "wandb":
+            raise ImportError("no wandb")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", _blocked)
+    monkeypatch.delitem(sys.modules, "wandb", raising=False)
+    assert tracking._wandb_usable() is False
+
+
+def test_wandb_usable_false_without_credentials(monkeypatch) -> None:
+    monkeypatch.delenv("WANDB_API_KEY", raising=False)
+    monkeypatch.delenv("WANDB_MODE", raising=False)
+    monkeypatch.setitem(
+        sys.modules,
+        "wandb",
+        SimpleNamespace(api=SimpleNamespace(api_key=None)),
+    )
+    assert tracking._wandb_usable() is False
+
+
+def test_wandb_usable_with_env_key(monkeypatch) -> None:
+    monkeypatch.setenv("WANDB_API_KEY", "key")
+    monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace())
+    assert tracking._wandb_usable() is True
+
+
+def test_build_tracker_noops_cleanly_without_wandb(monkeypatch, tmp_path) -> None:
+    """A wandb-project config must degrade to the dummy tracker, not crash."""
+    monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=0))
+    monkeypatch.setattr(tracking, "_wandb_usable", lambda: False)
+
+    tracker = tracking.build_tracker(
+        TrackerConfig(project_name="h3-dmd2-vsa"),
+        CheckpointConfig(output_dir=str(tmp_path)),
+        config={"method": {}},
+    )
+
+    assert isinstance(tracker, DummyTracker)
+    tracker.log({"total_loss": 1.0}, 1)  # must not raise
+    tracker.finish()
+
+
+def test_build_tracker_logs_loss_dict_with_monkeypatched_wandb(monkeypatch, tmp_path) -> None:
+    fake = _fake_wandb()
+    monkeypatch.setitem(sys.modules, "wandb", fake)
+    monkeypatch.setenv("WANDB_API_KEY", "key")
+    monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=0))
+
+    run_config = {"method": {"dmd_denoising_steps": [1000, 757, 522]}}
+    tracker = tracking.build_tracker(
+        TrackerConfig(project_name="h3-dmd2-vsa", run_name="dmd2_vsa0_overfit"),
+        CheckpointConfig(output_dir=str(tmp_path)),
+        config=run_config,
+    )
+
+    assert isinstance(tracker, WandbTracker)
+    (run, ) = fake.runs
+    assert run.init_kwargs["project"] == "h3-dmd2-vsa"
+    assert run.init_kwargs["name"] == "dmd2_vsa0_overfit"
+    assert run.init_kwargs["config"] == run_config
+
+    # The per-step dict the Trainer logs on rank 0 (DMD2 loss map + metrics).
+    metrics = {
+        "total_loss": 0.5,
+        "generator_loss": 0.25,
+        "fake_score_loss": 0.25,
+        "update_student": 1.0,
+        "step_time_sec": 0.1,
+        "vsa_sparsity": 0.0,
+    }
+    tracker.log(metrics, 7)
+    assert run.logged == [(metrics, 7)]
+
+
+def test_build_tracker_nonzero_rank_never_inits_wandb(monkeypatch, tmp_path) -> None:
+    fake = _fake_wandb()
+    monkeypatch.setitem(sys.modules, "wandb", fake)
+    monkeypatch.setenv("WANDB_API_KEY", "key")
+    monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=1))
+
+    tracker = tracking.build_tracker(
+        TrackerConfig(project_name="h3-dmd2-vsa"),
+        CheckpointConfig(output_dir=str(tmp_path)),
+        config=None,
+    )
+
+    assert isinstance(tracker, DummyTracker)
+    assert fake.runs == []

fastvideo/train/callbacks/callback.py

@@ -24,6 +24,7 @@ _BUILTIN_CALLBACKS: dict[str, str] = {
     "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
     "validation": "fastvideo.train.callbacks.validation.ValidationCallback",
     "ema": "fastvideo.train.callbacks.ema.EMACallback",
+    "latent_vis": "fastvideo.train.callbacks.latent_vis.LatentVisCallback",
 }
 
 

fastvideo/train/callbacks/latent_vis.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Intermediate-latent visualization callback for distillation methods.
+
+Port of the legacy ``fastvideo/training/distillation_pipeline.py`` latent
+logging to the modular trainer: every ``every_steps`` iterations, rank 0
+decodes the method's latest latent snapshots (``method.latent_vis`` — the
+student's rollout prediction plus the real- and fake-score predictions on
+generator-update steps) through the model's ``decode_vis_latents`` hook and
+logs them to the tracker as videos.
+
+Both hooks are optional: methods that never populate ``latent_vis`` or models
+without ``decode_vis_latents`` turn this callback into a no-op.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+from fastvideo.distributed import get_world_group
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.training.trackers import DummyTracker
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+_DEFAULT_KEYS = (
+    "generator_pred_video",
+    "real_score_pred_video",
+    "faker_score_pred_video",
+)
+
+
+class LatentVisCallback(Callback):
+    """Decode and log intermediate training latents as tracker videos."""
+
+    def __init__(
+        self,
+        every_steps: int = 100,
+        keys: list[str] | None = None,
+        fps: int = 24,
+    ) -> None:
+        self.every_steps = int(every_steps)
+        self.keys = tuple(keys) if keys else _DEFAULT_KEYS
+        self.fps = int(fps)
+        self.tracker: Any = DummyTracker()
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        tracker = getattr(method, "tracker", None)
+        if tracker is not None:
+            self.tracker = tracker
+
+    def on_training_step_end(
+        self,
+        method: TrainingMethod,
+        loss_dict: dict[str, Any],
+        iteration: int = 0,
+    ) -> None:
+        if self.every_steps <= 0 or iteration % self.every_steps != 0:
+            return
+        if get_world_group().rank != 0:
+            return
+        vis = getattr(method, "latent_vis", None)
+        if not vis:
+            return
+        decode = getattr(getattr(method, "student", None), "decode_vis_latents", None)
+        if decode is None:
+            return
+
+        artifacts: dict[str, Any] = {}
+        for key in self.keys:
+            latent = vis.get(key)
+            if not isinstance(latent, torch.Tensor):
+                continue
+            try:
+                clip = decode(latent)
+            except Exception as exc:
+                logger.warning("Latent visualization decode failed for %r: %s", key, exc)
+                continue
+            art = self.tracker.video(clip, fps=self.fps, format="mp4")
+            if art is not None:
+                artifacts[f"latent_vis/{key}"] = art
+        if artifacts:
+            self.tracker.log_artifacts(artifacts, iteration)

fastvideo/train/callbacks/validation.py

@@ -382,6 +382,17 @@ class ValidationCallback(Callback):
             self._restore_inactive_role_modules(module_records)
             self._restore_optimizer_states(optimizer_tensor_records)
             self._empty_cuda_cache()
+            off_cuda = [
+                f"{name} on {getattr(p, '_local_tensor', p).device}"
+                for name, p in validation_transformer.named_parameters()
+                if getattr(p, "_local_tensor", p).device.type != "cuda"
+            ]
+            if off_cuda:
+                logger.warning(
+                    "Post-validation: %d validation-transformer params off-CUDA; first: %s",
+                    len(off_cuda),
+                    off_cuda[:5],
+                )
 
     def _offload_optimizer_states_to_cpu(
         self,
@@ -466,6 +477,20 @@ class ValidationCallback(Callback):
             device = self._first_cuda_tensor_device(module)
             if device is None:
                 continue
+            if self._is_fsdp_managed(module):
+                # `.to()` round-trips on fully_shard modules replace DTensor
+                # local storage behind FSDP2's bookkeeping; the corruption
+                # surfaces as device-mismatch errors on the first backward
+                # through the module a few steps after restore. Keep sharded
+                # roles resident — the optimizer-state offload above already
+                # returns the bulk of the memory.
+                logger.info(
+                    "Keeping role %r transformer on %s during validation "
+                    "(FSDP-managed modules do not survive .to() round-trips).",
+                    role,
+                    device,
+                )
+                continue
             try:
                 module.to("cpu")
             except Exception as exc:
@@ -494,6 +519,17 @@ class ValidationCallback(Callback):
                 device,
             )
 
+    @staticmethod
+    def _is_fsdp_managed(module: torch.nn.Module) -> bool:
+        try:
+            from torch.distributed.fsdp import FSDPModule
+            from torch.distributed.tensor import DTensor
+        except ImportError:
+            return False
+        if isinstance(module, FSDPModule):
+            return True
+        return any(isinstance(p, DTensor) for p in module.parameters(recurse=True))
+
     @staticmethod
     def _first_cuda_tensor_device(module: torch.nn.Module) -> torch.device | None:
         for tensor in list(module.parameters(recurse=True)) + list(module.buffers(recurse=True)):
@@ -1237,6 +1273,25 @@ class ValidationCallback(Callback):
                 getattr(transformer, name),
             )
 
+    def _inject_method_denoising_steps(self, validation_config: Any) -> None:
+        """Sample validation at the method's trained DMD jump points.
+
+        Distillation students only ever denoise from ``dmd_denoising_steps``;
+        stages that honor ``pipeline_config.dmd_denoising_steps`` should visit
+        the same points instead of the scheduler's native grid. An explicit
+        value already on the pipeline config wins, and warped lists stay out
+        (their raw entries are grid indices, not timesteps).
+        """
+        method_config = getattr(self.method, "method_config", None)
+        if not isinstance(method_config, dict):
+            return
+        steps = method_config.get("dmd_denoising_steps")
+        if not steps or bool(method_config.get("warp_denoising_step", False)):
+            return
+        if getattr(validation_config, "dmd_denoising_steps", "unset") is not None:
+            return
+        validation_config.dmd_denoising_steps = [int(step) for step in steps]
+
     def _get_pipeline(
         self,
         *,
@@ -1288,6 +1343,7 @@ class ValidationCallback(Callback):
                 validation_config,
                 loaded_config,
             )
+            self._inject_method_denoising_steps(validation_config)
             self._pipeline.fastvideo_args.pipeline_config = validation_config
             arch_config = self._pipeline.fastvideo_args.pipeline_config.dit_config.arch_config
             logger.info(

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -62,6 +62,7 @@ def _save_role_pretrained(
     output_dir: str,
     module_names: list[str] | None = None,
     overwrite: bool = False,
+    link_base: bool = False,
     model: Any,
 ) -> str:
     """Export a role's modules into a diffusers-style model dir.
@@ -117,16 +118,51 @@ def _save_role_pretrained(
 
         logger.info(
             "Creating pretrained export dir at %s "
-            "(base=%s)",
+            "(base=%s, link_base=%s)",
             dst,
             local_base,
+            link_base,
         )
-        shutil.copytree(
-            local_base,
-            dst,
-            symlinks=False,
-            copy_function=_copy_or_link,
-        )
+        if link_base:
+            # Space-lean export: symlink every base component except the
+            # module dirs we are about to rewrite (those get a real dir with
+            # the base's non-weight files, e.g. config.json — weights and
+            # index are produced fresh below). Cross-user hardlinks are
+            # blocked by fs.protected_hardlinks, symlinks are not.
+            rewritten = set(module_names or ["transformer"])
+            dst.mkdir(parents=True, exist_ok=True)
+            for entry in sorted(local_base.iterdir()):
+                if entry.name == ".cache" or entry.name.startswith(".git"):
+                    continue
+                target = dst / entry.name
+                if entry.is_dir() and entry.name in rewritten:
+                    target.mkdir()
+                    for f in sorted(entry.iterdir()):
+                        if f.name.endswith(".safetensors") or f.name.endswith(".safetensors.index.json"):
+                            continue
+                        shutil.copy2(os.path.realpath(f), target / f.name)
+                else:
+                    os.symlink(os.path.realpath(entry), target)
+        else:
+            try:
+                shutil.copytree(
+                    local_base,
+                    dst,
+                    symlinks=False,
+                    copy_function=_copy_or_link,
+                    # HF hub bookkeeping under the base checkpoint (.cache/)
+                    # may be unreadable when the base belongs to another
+                    # user; it is not part of the model.
+                    ignore=shutil.ignore_patterns(".cache", ".git*"),
+                )
+            except shutil.Error as exc:
+                # copytree collects per-file failures and raises at the end;
+                # tolerate stragglers as long as the component manifest made
+                # it.
+                logger.warning("copytree finished with %d skipped entries (first: %s)", len(exc.args[0]),
+                               exc.args[0][0] if exc.args[0] else "?")
+        if not ((dst / "modular_model_index.json").is_file() or (dst / "model_index.json").is_file()):
+            raise FileNotFoundError(f"Export dir {dst} is missing its model index after copy.")
 
     _barrier()
 
@@ -160,6 +196,10 @@ def _save_role_pretrained(
         if _rank() == 0:
             for path in module_dir.glob("*.safetensors"):
                 path.unlink(missing_ok=True)
+            # A leftover shard index from the base would point at the shards
+            # deleted above and shadow the fresh single-file export.
+            for path in module_dir.glob("*.safetensors.index.json"):
+                path.unlink(missing_ok=True)
 
             # Convert internal parameter names back to HF format.
             # load_model_from_full_model_state_dict builds reverse_param_names_mapping
@@ -233,6 +273,8 @@ def convert(
     role: str = "student",
     overwrite: bool = False,
     verify: bool = False,
+    weights_only: bool = False,
+    link_base: bool = False,
 ) -> str:
     """Load a DCP checkpoint and export as a diffusers model.
 
@@ -296,6 +338,11 @@ def convert(
 
     # -- Load DCP weights into the model --
     states = method.checkpoint_state()
+    if weights_only:
+        # Export needs only module weights; optimizer states (a full fp32
+        # second-moment tensor per trainable role) can double GPU memory and
+        # force structure-matching between the shim and checkpoint optimizers.
+        states = {k: v for k, v in states.items() if k.startswith("roles.")}
     logger.info(
         "Loading DCP checkpoint from %s",
         resolved,
@@ -321,6 +368,7 @@ def convert(
         base_model_path=base_model_path,
         output_dir=output_dir,
         overwrite=overwrite,
+        link_base=link_base,
         model=model,
     )
     logger.info("Export complete: %s", result)
@@ -443,6 +491,22 @@ def main() -> None:
               "the exported directory to catch key-mapping bugs "
               "immediately."),
     )
+    parser.add_argument(
+        "--weights-only",
+        action="store_true",
+        help=("Load only roles.* module weights from the checkpoint, "
+              "skipping optimizer/scheduler states (halves GPU memory "
+              "and allows exporting via a shim config whose optimizer "
+              "differs from the checkpoint's)."),
+    )
+    parser.add_argument(
+        "--link-base",
+        action="store_true",
+        help=("Symlink base-model components into the export dir instead "
+              "of copying them (only the exported module dirs are real). "
+              "Saves hundreds of GB per export; the export then depends on "
+              "the base model dir staying in place."),
+    )
     args = parser.parse_args(sys.argv[1:])
 
     convert(
@@ -452,6 +516,8 @@ def main() -> None:
         role=args.role,
         overwrite=args.overwrite,
         verify=args.verify,
+        weights_only=args.weights_only,
+        link_base=args.link_base,
     )
 
 

fastvideo/train/entrypoint/train.py

@@ -97,6 +97,7 @@ def run_training_from_config(
     ckpt_config = CheckpointConfig(
         save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0),
         keep_last=int(tc.checkpoint.checkpoints_total_limit or 0),
+        start_step=int(tc.checkpoint.checkpointing_start_step or 0),
     )
 
     checkpoint_manager = CheckpointManager(

fastvideo/train/methods/base.py

@@ -109,6 +109,17 @@ class TrainingMethod(torch.nn.Module, ABC):
     def _lr_scheduler_dict(self) -> dict[str, Any]:
         ...
 
+    def apply_configured_lrs(self) -> None:
+        """Re-apply the config's learning rates to live optimizers/schedulers.
+
+        Called after a checkpoint resume when
+        ``training.checkpoint.reset_lr_on_resume`` is set: DCP restores
+        param-group ``lr``/``initial_lr`` and scheduler ``base_lrs`` from the
+        checkpoint, which would otherwise silently override an LR change in
+        the YAML. Methods with per-role LRs override this.
+        """
+        return None
+
     def checkpoint_state(self) -> dict[str, Any]:
         """Return DCP-ready checkpoint state for all trainable roles.
 
@@ -275,6 +286,6 @@ class TrainingMethod(torch.nn.Module, ABC):
     def _infer_attn_kind(self) -> Literal["dense", "vsa"]:
         """Derive metadata mode from the student's resolved backend."""
         backend = (self.student.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND)
-        if backend == "VIDEO_SPARSE_ATTN":
+        if backend in ("VIDEO_SPARSE_ATTN", "VIDEO_SPARSE_ATTN_H3"):
             return "vsa"
         return "dense"

fastvideo/train/methods/distribution_matching/dmd2.py

@@ -62,6 +62,7 @@ class DMD2Method(TrainingMethod):
             self._score_min_timestep,
             self._score_max_timestep,
         ) = self._parse_score_timestep_bounds()
+        self._score_timestep_shift = self._parse_score_timestep_shift()
 
         # Initialize preprocessors on student.
         self.student.init_preprocessors(self.training_config)
@@ -110,19 +111,29 @@ class DMD2Method(TrainingMethod):
             dtype=training_batch.latents.dtype,
         )
         student_ctx = None
+        generator_metrics: dict[str, LogScalar] = {}
+        rollout_for_critic: torch.Tensor | None = None
         if update_student:
             generator_pred_x0 = self._student_rollout(training_batch, with_grad=True)
             student_ctx = (
                 training_batch.timesteps,
                 training_batch.attn_metadata_vsa,
             )
-            generator_loss = self._dmd_loss(generator_pred_x0, training_batch)
+            generator_loss, generator_metrics = self._dmd_loss(generator_pred_x0, training_batch)
+            # Reference DMD2 trains the critic on the same generated sample
+            # the generator step used; a second independent rollout here
+            # would only add ~len(dmd_denoising_steps) no-grad forwards.
+            rollout_for_critic = generator_pred_x0.detach()
 
         (
             fake_score_loss,
             critic_ctx,
             critic_outputs,
-        ) = self._critic_flow_matching_loss(training_batch)
+            critic_metrics,
+        ) = self._critic_flow_matching_loss(
+            training_batch,
+            generator_pred_x0=rollout_for_critic,
+        )
 
         total_loss = generator_loss + fake_score_loss
         loss_map = {
@@ -137,7 +148,18 @@ class DMD2Method(TrainingMethod):
             "student_ctx": student_ctx,
             "critic_ctx": critic_ctx,
         }
-        metrics: dict[str, LogScalar] = {"update_student": float(update_student)}
+        metrics: dict[str, LogScalar] = {
+            "update_student": float(update_student),
+            **generator_metrics,
+            **critic_metrics,
+        }
+        # Rank-local latent snapshots for LatentVisCallback. The critic dict
+        # always carries generator_pred_video; the DMD dict adds the real- and
+        # fake-score predictions on generator-update iterations.
+        self.latent_vis = {
+            **(training_batch.fake_score_latent_vis_dict or {}),
+            **(training_batch.dmd_latent_vis_dict or {}),
+        }
         return loss_map, outputs, metrics
 
     # TrainingMethod override: backward
@@ -370,6 +392,43 @@ class DMD2Method(TrainingMethod):
             scheduler_name=critic_sched,
         )
 
+    def _modality_slices(self) -> tuple[tuple[str, slice], ...] | None:
+        """Named packed-latent slices exposed by multi-modality models.
+
+        Packed adapters (e.g. MiniMax-H3's video+audio flattening) weight
+        modalities by element count under a single global mean, which mutes
+        the smaller stream — H3 audio is <1% of packed elements. When the
+        student exposes ``modality_slices()``, every DMD2 loss and normalizer
+        is computed per modality and combined via ``_modality_weight``.
+        """
+        getter = getattr(self.student, "modality_slices", None)
+        if getter is None:
+            return None
+        slices = tuple(getter())
+        return slices or None
+
+    def _modality_weight(self, name: str) -> float:
+        raw = self.method_config.get("modality_loss_weights", None)
+        if not isinstance(raw, dict):
+            return 1.0
+        value = raw.get(name, 1.0)
+        return 1.0 if value is None else float(value)
+
+    def apply_configured_lrs(self) -> None:
+        """Force student/critic LRs back to the configured values (post-resume)."""
+        student_lr = float(self.training_config.optimizer.learning_rate)
+        critic_lr = float(self.method_config.get("fake_score_learning_rate"))
+        for optimizer, scheduler, lr in (
+            (self._student_optimizer, self._student_lr_scheduler, student_lr),
+            (self._critic_optimizer, self._critic_lr_scheduler, critic_lr),
+        ):
+            for group in optimizer.param_groups:
+                group["lr"] = lr
+                if "initial_lr" in group:
+                    group["initial_lr"] = lr
+            if hasattr(scheduler, "base_lrs"):
+                scheduler.base_lrs = [lr] * len(scheduler.base_lrs)
+
     def _should_update_student(
         self,
         iteration: int,
@@ -461,15 +520,54 @@ class DMD2Method(TrainingMethod):
             int(max_ratio * num_timesteps),
         )
 
-    def _sample_score_timestep(self, device: torch.device) -> torch.Tensor:
-        timestep = torch.randint(
-            0,
-            int(self.student.num_train_timesteps),
-            [1],
-            device=device,
-            dtype=torch.long,
-            generator=self.cuda_generator,
+    def _parse_score_timestep_shift(self) -> float:
+        """Resolve the sampling-density shift for the score timestep.
+
+        ``1.0`` (default) keeps the legacy uniform draw over base timesteps.
+        For rectified-flow models with a large timestep shift ``s``, uniform
+        base-t concentrates ~s× more supervision near sigma=1 than near the
+        floor; setting this to ``s`` samples uniformly in that modality's
+        *shifted sigma* instead (draw u ~ U over sigma-space bounds, invert
+        t = u / (s - (s-1)u)), equalizing score supervision across its noise
+        axis. Only the sampling density changes — every timestep still maps
+        to per-modality sigmas exactly as before.
+        """
+        shift = get_optional_float(
+            self.method_config,
+            "score_timestep_shift",
+            where="method.score_timestep_shift",
         )
+        shift = 1.0 if shift is None else float(shift)
+        if shift <= 0.0:
+            raise ValueError("method.score_timestep_shift must be > 0, "
+                             f"got {shift}")
+        return shift
+
+    def _sample_score_timestep(self, device: torch.device) -> torch.Tensor:
+        shift = self._score_timestep_shift
+        if shift == 1.0:
+            timestep = torch.randint(
+                0,
+                int(self.student.num_train_timesteps),
+                [1],
+                device=device,
+                dtype=torch.long,
+                generator=self.cuda_generator,
+            )
+        else:
+            num_timesteps = float(self.student.num_train_timesteps)
+            t_lo = self._score_min_timestep / num_timesteps
+            t_hi = self._score_max_timestep / num_timesteps
+            sigma_lo = shift * t_lo / (1.0 + (shift - 1.0) * t_lo)
+            sigma_hi = shift * t_hi / (1.0 + (shift - 1.0) * t_hi)
+            u = torch.rand(
+                [1],
+                device=device,
+                dtype=torch.float32,
+                generator=self.cuda_generator,
+            ) * (sigma_hi - sigma_lo) + sigma_lo
+            t = u / (shift - (shift - 1.0) * u)
+            timestep = (t * num_timesteps).round().to(torch.long)
         timestep = self.student.shift_and_clamp_timestep(timestep)
         return timestep.clamp(
             self._score_min_timestep,
@@ -601,9 +699,12 @@ class DMD2Method(TrainingMethod):
     def _critic_flow_matching_loss(
         self,
         batch: Any,
-    ) -> tuple[torch.Tensor, Any, dict[str, Any]]:
-        with torch.no_grad():
-            generator_pred_x0 = self._student_rollout(batch, with_grad=False)
+        *,
+        generator_pred_x0: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, Any, dict[str, Any], dict[str, LogScalar]]:
+        if generator_pred_x0 is None:
+            with torch.no_grad():
+                generator_pred_x0 = self._student_rollout(batch, with_grad=False)
 
         device = generator_pred_x0.device
         fake_score_timestep = self._sample_score_timestep(device)
@@ -625,7 +726,17 @@ class DMD2Method(TrainingMethod):
             attn_kind="dense",
         )
         target = noise - generator_pred_x0
-        flow_matching_loss = torch.mean((pred_noise - target)**2)
+        slices = self._modality_slices()
+        emit_modality_metrics = slices is not None
+        if slices is None:
+            slices = (("packed", slice(None)), )
+        flow_matching_loss = torch.zeros((), device=device, dtype=torch.float32)
+        metrics: dict[str, LogScalar] = {}
+        for name, modality in slices:
+            loss_m = torch.mean((pred_noise[:, modality].float() - target[:, modality].float())**2)
+            flow_matching_loss = flow_matching_loss + self._modality_weight(name) * loss_m
+            if emit_modality_metrics:
+                metrics[f"fake_score_loss_{name}"] = loss_m.detach()
 
         batch.fake_score_latent_vis_dict = {
             "generator_pred_video": generator_pred_x0,
@@ -636,13 +747,14 @@ class DMD2Method(TrainingMethod):
             flow_matching_loss,
             (batch.timesteps, batch.attn_metadata),
             outputs,
+            metrics,
         )
 
     def _dmd_loss(
         self,
         generator_pred_x0: torch.Tensor,
         batch: Any,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, dict[str, LogScalar]]:
         guidance_scale = get_optional_float(
             self.method_config,
             "real_score_guidance_scale",
@@ -679,22 +791,48 @@ class DMD2Method(TrainingMethod):
                 cfg_uncond=self._cfg_uncond,
                 attn_kind="dense",
             )
-            real_uncond_x0 = self.teacher.predict_x0(
-                noisy_latents,
-                timestep,
-                batch,
-                conditional=False,
-                cfg_uncond=self._cfg_uncond,
-                attn_kind="dense",
-            )
-            real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale
+            if float(guidance_scale) == 1.0:
+                # At scale 1.0 the CFG combination is the identity, so the
+                # unconditional forward is pure waste. This is also the
+                # correct setting for guidance-distilled teachers (e.g.
+                # MiniMax-H3, whose own inference stack rejects CFG): their
+                # conditional prediction already bakes in guidance, and they
+                # define no trained unconditional branch to extrapolate from.
+                real_cfg_x0 = real_cond_x0
+            else:
+                real_uncond_x0 = self.teacher.predict_x0(
+                    noisy_latents,
+                    timestep,
+                    batch,
+                    conditional=False,
+                    cfg_uncond=self._cfg_uncond,
+                    attn_kind="dense",
+                )
+                real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale
 
-            denom = torch.abs(generator_pred_x0 - real_cfg_x0).mean()
-            grad = (faker_x0 - real_cfg_x0) / denom
-            grad = torch.nan_to_num(grad)
+            # Intermediate-latent visualization state (legacy training/
+            # distillation_pipeline.py key names); LatentVisCallback decodes
+            # these on rank 0 every N steps.
+            batch.dmd_latent_vis_dict.update({
+                "real_score_pred_video": real_cfg_x0.detach(),
+                "faker_score_pred_video": faker_x0.detach(),
+                "dmd_timestep": timestep.detach(),
+            })
 
-        loss = 0.5 * F.mse_loss(
-            generator_pred_x0.float(),
-            (generator_pred_x0.float() - grad.float()).detach(),
-        )
-        return loss
+        slices = self._modality_slices()
+        emit_modality_metrics = slices is not None
+        if slices is None:
+            slices = (("packed", slice(None)), )
+        loss = torch.zeros((), device=device, dtype=torch.float32)
+        metrics: dict[str, LogScalar] = {}
+        for name, modality in slices:
+            with torch.no_grad():
+                real_m = real_cfg_x0[:, modality]
+                denom = torch.abs(generator_pred_x0[:, modality] - real_m).mean()
+                grad = torch.nan_to_num((faker_x0[:, modality] - real_m) / denom)
+            gen_m = generator_pred_x0[:, modality].float()
+            loss_m = 0.5 * F.mse_loss(gen_m, (gen_m - grad.float()).detach())
+            loss = loss + self._modality_weight(name) * loss_m
+            if emit_modality_metrics:
+                metrics[f"generator_loss_{name}"] = loss_m.detach()
+        return loss, metrics

fastvideo/train/models/minimax_h3/__init__.py

@@ -3,3 +3,5 @@
 
 from fastvideo.train.models.minimax_h3.minimax_h3 import (
     MiniMaxH3Model as MiniMaxH3Model, )
+from fastvideo.train.models.minimax_h3.minimax_h3_dmd import (
+    MiniMaxH3DMDModel as MiniMaxH3DMDModel, )

fastvideo/train/models/minimax_h3/minimax_h3.py

@@ -7,6 +7,9 @@ from typing import Any, Literal, TYPE_CHECKING
 
 import torch
 
+import fastvideo.envs as envs
+from fastvideo.attention.backends.video_sparse_attn_h3 import (
+    MiniMaxH3VSAMetadataBuilder, )
 from fastvideo.distributed import get_sp_group
 from fastvideo.forward_context import set_forward_context
 from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
@@ -22,6 +25,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import (
     unpack_audio_tokens,
     unpatchify_video_tokens,
 )
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (
+    _h3_vsa_prefix_segments, )
 from fastvideo.platforms import AttentionBackendEnum
 
 from fastvideo.train.models.base import ModelBase, NoisePrediction
@@ -39,6 +44,15 @@ _AUDIO_SCHEDULER_SHIFT = 3.0
 _VIDEO_LATENT_CHANNELS = 24
 _AUDIO_LATENT_CHANNELS = 32
 
+# Dense TORCH_SDPA is the default; per-role overrides allow FLASH_ATTN
+# (teacher/critic, FA4 via FASTVIDEO_FA4=1) and the packed-sequence VSA-H3
+# backend (distillation student).
+_ALLOWED_ATTENTION_BACKENDS = (
+    AttentionBackendEnum.TORCH_SDPA,
+    AttentionBackendEnum.FLASH_ATTN,
+    AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,
+)
+
 
 def shift_noise_amount(base_noise_amount: torch.Tensor, shift: float) -> torch.Tensor:
     """Apply the MiniMax H3 rational shift to a unit noise amount."""
@@ -68,10 +82,12 @@ class MiniMaxH3Model(ModelBase):
             trainable=trainable,
             attention_backend=attention_backend,
         )
-        # PyTorch scaled dot product attention (SDPA) provides dense attention
-        # without adding another attention-kernel dependency to H3 training.
-        if self.attention_backend != AttentionBackendEnum.TORCH_SDPA:
-            raise ValueError("MiniMaxH3Model requires the TORCH_SDPA attention backend")
+        # Attention layers bind their backend during construction, so this is
+        # the per-role selection point (student/teacher/critic can differ).
+        if self.attention_backend not in _ALLOWED_ATTENTION_BACKENDS:
+            allowed = ", ".join(b.name for b in _ALLOWED_ATTENTION_BACKENDS)
+            raise ValueError("MiniMaxH3Model supports the attention backends "
+                             f"{{{allowed}}}, got {self.attention_backend}")
         if training_config.pipeline_config is None:
             raise ValueError("MiniMaxH3Model requires a resolved MiniMax H3 pipeline config")
         # Packed row indices describe one text-video-audio document without a
@@ -83,9 +99,12 @@ class MiniMaxH3Model(ModelBase):
         if float(training_config.data.training_cfg_rate) != 0.0:
             raise ValueError("MiniMaxH3Model requires training.data.training_cfg_rate=0.0")
         # Joint supervision requires paired video and stereo-audio latents from
-        # every parquet row.
-        if str(training_config.data.preprocessed_data_type) != "t2va":
-            raise ValueError("MiniMaxH3Model requires training.data.preprocessed_data_type='t2va'")
+        # every parquet row ('t2va'). 'text_only' rows carry prompt conditioning
+        # alone and are valid only for data-free methods that synthesize latent
+        # shapes from config (DMD2 enforces rollout_mode='simulate' for them).
+        if str(training_config.data.preprocessed_data_type) not in ("t2va", "text_only"):
+            raise ValueError("MiniMaxH3Model requires training.data.preprocessed_data_type "
+                             "'t2va' or 'text_only'")
 
         # FastVideo's Fully Sharded Data Parallel loading path requires one BF16
         # parameter dtype, including modules that H3 inference keeps in FP32.
@@ -135,15 +154,17 @@ class MiniMaxH3Model(ModelBase):
 
     def init_preprocessors(self, training_config: TrainingConfig) -> None:
         """Load precomputed text embeddings and paired video-audio latents."""
-        from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2va
+        from fastvideo.dataset.dataloader.schema import (pyarrow_schema_t2va, pyarrow_schema_text_only)
         from fastvideo.train.utils.dataloader import build_parquet_t2v_train_dataloader
 
         self.sp_group = get_sp_group()
         text_config = training_config.pipeline_config.text_encoder_configs[0]  # type: ignore[union-attr]
+        parquet_schema = (pyarrow_schema_text_only
+                          if str(training_config.data.preprocessed_data_type) == "text_only" else pyarrow_schema_t2va)
         self.dataloader = build_parquet_t2v_train_dataloader(
             training_config.data,
             text_len=int(text_config.arch_config.text_len),
-            parquet_schema=pyarrow_schema_t2va,
+            parquet_schema=parquet_schema,
         )
         self.start_step = 0
 
@@ -294,8 +315,40 @@ class MiniMaxH3Model(ModelBase):
         training_batch.minimax_h3_layout = layout
         training_batch.attn_metadata = None
         training_batch.attn_metadata_vsa = None
+        self._maybe_build_vsa_metadata(training_batch)
         return training_batch
 
+    def _maybe_build_vsa_metadata(self, batch: TrainingBatch) -> None:
+        """Populate the VSA view when this model runs the VSA-H3 backend.
+
+        Methods route it through ``predict_noise(attn_kind="vsa")``; dense
+        models keep ``attn_metadata_vsa`` as ``None`` so their forwards stay
+        dense. One builder per model keeps the padded tile buffer reused
+        across steps (see ``MiniMaxH3VSAMetadata.tile_buf_holder``).
+        """
+        backend = (getattr(self, "attention_backend_name", None) or envs.FASTVIDEO_ATTENTION_BACKEND)
+        if backend != "VIDEO_SPARSE_ATTN_H3":
+            return
+        layout = batch.minimax_h3_layout
+        patch_size = tuple(self.transformer.patch_size)
+        builder = getattr(self, "_vsa_metadata_builder", None)
+        if builder is None:
+            builder = self._vsa_metadata_builder = MiniMaxH3VSAMetadataBuilder()
+        batch.attn_metadata_vsa = builder.build(
+            # Training builds one metadata per batch and reuses it across the
+            # step's forwards; the step index only feeds probe bookkeeping.
+            current_timestep=0,
+            raw_latent_shape=(
+                layout.num_video_latent_frames,
+                layout.latent_height,
+                layout.latent_width,
+            ),
+            patch_size=patch_size,
+            VSA_sparsity=float(self.training_config.vsa_sparsity),
+            prefix_segments=_h3_vsa_prefix_segments(layout, patch_size),
+            device=self.device,
+        )
+
     def add_noise(
         self,
         clean_latents: torch.Tensor,
@@ -320,10 +373,12 @@ class MiniMaxH3Model(ModelBase):
     ) -> NoisePrediction:
         """Pack modality timesteps and convert H3 outputs to noise-minus-clean."""
         del timestep
-        if not conditional or cfg_uncond is not None:
-            raise ValueError("MiniMaxH3Model predicts one conditional T2VA sample")
-        if attn_kind != "dense":
-            raise ValueError("MiniMaxH3Model supports dense attention for training")
+        # Under dense backends both metadata views are None, so "vsa" silently
+        # means dense (mirrors WanModel). MiniMaxH3DMDModel.prepare_batch
+        # populates attn_metadata_vsa when the role runs VIDEO_SPARSE_ATTN_H3.
+        if attn_kind not in ("dense", "vsa"):
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+        attn_metadata = (batch.attn_metadata_vsa if attn_kind == "vsa" else batch.attn_metadata)
         layout = batch.minimax_h3_layout
         if not isinstance(layout, MiniMaxH3PackedLayout):
             raise RuntimeError("prepare_batch() must set TrainingBatch.minimax_h3_layout")
@@ -332,6 +387,16 @@ class MiniMaxH3Model(ModelBase):
         if batch.timesteps is None or batch.audio_timesteps is None:
             raise RuntimeError("prepare_batch() must set video and audio timesteps")
 
+        encoder_hidden_states = batch.encoder_hidden_states
+        if not conditional:
+            # H3 has no negative-prompt encoder at training time, so the only
+            # supported unconditional branch (teacher CFG in distillation)
+            # zeroes the text embeddings.
+            if (cfg_uncond or {}).get("text") != "zero":
+                raise ValueError("MiniMaxH3Model unconditional forwards require "
+                                 "method.cfg_uncond={'text': 'zero'}")
+            encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
+
         dtype = torch.bfloat16
         device = self.device
         video_bcthw = noisy_latents.permute(0, 2, 1, 3, 4).to(dtype)
@@ -354,12 +419,12 @@ class MiniMaxH3Model(ModelBase):
 
         with torch.autocast(device.type, dtype=dtype), set_forward_context(
                 current_timestep=unique_timesteps,
-                attn_metadata=None,
+                attn_metadata=attn_metadata,
         ):
             video_velocity, audio_velocity = self.transformer(
                 hidden_states=video_rows[None],
                 audio_hidden_states=audio_rows[None],
-                encoder_hidden_states=batch.encoder_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
                 timestep=unique_timesteps,
                 timestep_indices=timestep_indices,
                 token_tags=layout.token_tags.to(device),
@@ -390,6 +455,19 @@ class MiniMaxH3Model(ModelBase):
     ) -> None:
         """Restore the forward context and average accumulated microbatch gradients."""
         timesteps, attn_metadata = ctx
+        expected_device_type = self.device.type
+        offenders = []
+        for name, param in self.transformer.named_parameters():
+            local = getattr(param, "_local_tensor", param)
+            if local.device.type != expected_device_type:
+                offenders.append(f"param {name} on {local.device}")
+            if param.grad is not None:
+                grad_local = getattr(param.grad, "_local_tensor", param.grad)
+                if grad_local.device.type != expected_device_type:
+                    offenders.append(f"grad {name} on {grad_local.device}")
+        if offenders:
+            raise RuntimeError(f"{len(offenders)} training tensors off-CUDA before backward; "
+                               f"first offenders: {offenders[:8]}")
         with set_forward_context(
                 current_timestep=timesteps,
                 attn_metadata=attn_metadata,

fastvideo/train/models/minimax_h3/minimax_h3_dmd.py

@@ -0,0 +1,309 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 distribution-matching adapter (packed dual-modality latents)."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Literal
+
+import torch
+
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.pipelines.basic.minimax_h3.packing import (
+    MINIMAX_H3_AUDIO_CHANNELS,
+    audio_latent_num_frames,
+)
+from fastvideo.train.models.minimax_h3.minimax_h3 import (
+    _AUDIO_LATENT_CHANNELS,
+    _AUDIO_SCHEDULER_SHIFT,
+    _VIDEO_LATENT_CHANNELS,
+    _VIDEO_SCHEDULER_SHIFT,
+    MiniMaxH3Model,
+    shift_noise_amount,
+)
+
+# DMD2 samples integer score timesteps on the legacy [0, 1000] scale. H3 maps
+# them to its shared base noise amount before applying the modality shifts.
+_DMD_TIMESTEP_SCALE = 1000
+
+
+class MiniMaxH3DMDModel(MiniMaxH3Model):
+    """Present H3's dual (video, audio) streams to DMD2 as one packed tensor.
+
+    ``DMD2Method``'s rollout and loss math assume one latent tensor per
+    sample. This adapter flattens both modality latents into one ``[1, N]``
+    tensor (video's ``[1, T, 24, H, W]`` elements first, stereo audio's
+    ``[1, 2, 32, Ta]`` elements after) so ``dmd2.py`` stays model-agnostic.
+    Integer method timesteps become one shared base noise amount that is
+    shifted per modality (video 12.0, audio 3.0), exactly as H3's paired
+    schedulers synchronize the two streams during fine-tuning and inference.
+
+    ``modality_slices()`` exposes the packed video/audio column ranges so
+    DMD2 computes losses and normalizers per modality instead of one packed
+    mean that video's element count would dominate.
+    """
+
+    @property
+    def num_train_timesteps(self) -> int:
+        return _DMD_TIMESTEP_SCALE
+
+    # ------------------------------------------------------------------
+    # Packed dual-modality helpers
+    # ------------------------------------------------------------------
+
+    def _modality_shapes(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
+        """Return the ``[1, T, C, H, W]`` video and ``[1, 2, 32, Ta]`` audio shapes."""
+        data = self.training_config.data
+        video_shape = (
+            1,
+            int(data.num_latent_t),
+            _VIDEO_LATENT_CHANNELS,
+            int(data.num_height) // 16,
+            int(data.num_width) // 16,
+        )
+        audio_shape = (
+            1,
+            MINIMAX_H3_AUDIO_CHANNELS,
+            _AUDIO_LATENT_CHANNELS,
+            audio_latent_num_frames(int(data.num_frames)),
+        )
+        return video_shape, audio_shape
+
+    def pack_latents(
+        self,
+        video_latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+    ) -> torch.Tensor:
+        """Flatten both modality latents into one ``[1, N]`` tensor."""
+        return torch.cat(
+            (video_latents.reshape(1, -1), audio_latents.reshape(1, -1)),
+            dim=1,
+        )
+
+    def modality_slices(self) -> tuple[tuple[str, slice], ...]:
+        """Named packed-latent column slices for per-modality DMD2 losses.
+
+        Video is ~3.6M packed elements against audio's ~15-30k, so a single
+        global mean would give audio <1% of the distillation signal; DMD2
+        consumes these slices to normalize and weight each stream separately.
+        """
+        video_shape, audio_shape = self._modality_shapes()
+        split = math.prod(video_shape)
+        return (
+            ("video", slice(0, split)),
+            ("audio", slice(split, split + math.prod(audio_shape))),
+        )
+
+    def unpack_latents(self, packed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split one packed ``[1, N]`` tensor back into (video, audio) latents."""
+        video_shape, audio_shape = self._modality_shapes()
+        split = math.prod(video_shape)
+        if packed.shape != (1, split + math.prod(audio_shape)):
+            raise ValueError("Packed latents must have shape "
+                             f"[1, {split + math.prod(audio_shape)}], got {tuple(packed.shape)}")
+        return (
+            packed[:, :split].reshape(video_shape),
+            packed[:, split:].reshape(audio_shape),
+        )
+
+    def _noise_amounts(self, timestep: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Map one integer method timestep to both modality noise amounts."""
+        base = (timestep.reshape(-1)[:1].to(torch.float32) / _DMD_TIMESTEP_SCALE)
+        base = base.clamp(0.0, 1.0)
+        return (
+            shift_noise_amount(base, _VIDEO_SCHEDULER_SHIFT),
+            shift_noise_amount(base, _AUDIO_SCHEDULER_SHIFT),
+        )
+
+    # ------------------------------------------------------------------
+    # ModelBase overrides (packed convention)
+    # ------------------------------------------------------------------
+
+    def set_requires_negative_conditioning(self, requires: bool) -> None:
+        """Fail fast: H3 cannot encode negative prompts at training time."""
+        if requires:
+            raise ValueError("MiniMaxH3DMDModel has no negative-prompt encoder; set "
+                             "method.cfg_uncond={'text': 'zero'} for unconditional forwards")
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        """Prepare the T2VA batch, then expose clean latents in packed form."""
+        batch = super().prepare_batch(
+            raw_batch,
+            generator=generator,
+            latents_source=latents_source,
+        )
+        # DMD2 draws its own noise and timesteps per forward; only the packed
+        # clean latents matter here (the base prepare_batch already built the
+        # VSA metadata view for VSA-H3 roles). The fine-tuning noisy fields
+        # are refreshed by predict_noise on every call.
+        batch.latents = self.pack_latents(batch.latents, batch.audio_latents)
+        return batch
+
+    def add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Noise packed latents at one shared timestep, shifted per modality."""
+        sigma_video, sigma_audio = self._noise_amounts(timestep)
+        clean_video, clean_audio = self.unpack_latents(clean_latents)
+        noise_video, noise_audio = self.unpack_latents(noise)
+        return self.pack_latents(
+            self._mix(clean_video, noise_video, sigma_video),
+            self._mix(clean_audio, noise_audio, sigma_audio),
+        )
+
+    @staticmethod
+    def _mix(
+        clean: torch.Tensor,
+        noise: torch.Tensor,
+        sigma: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = sigma.to(device=clean.device, dtype=clean.dtype)
+        return (1.0 - sigma) * clean + sigma * noise
+
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        """Run one packed joint forward at an explicit method timestep.
+
+        Both modality clean-time fields on ``batch`` are rewritten from
+        ``timestep`` so the packed-row timestep plan and the backward
+        forward-context stay coherent with this call.
+        """
+        sigma_video, sigma_audio = self._noise_amounts(timestep)
+        noisy_video, noisy_audio = self.unpack_latents(noisy_latents)
+        batch.timesteps = (1.0 - sigma_video).to(noisy_latents.device)
+        batch.audio_timesteps = (1.0 - sigma_audio).to(noisy_latents.device)
+        batch.audio_noisy_model_input = noisy_audio
+        video_pred, audio_pred = super().predict_noise(
+            noisy_video,
+            timestep,
+            batch,
+            conditional=conditional,
+            cfg_uncond=cfg_uncond,
+            attn_kind=attn_kind,
+        )
+        return self.pack_latents(video_pred, audio_pred)
+
+    def predict_x0(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        """Convert packed noise-minus-clean predictions to packed clean latents."""
+        pred_noise = self.predict_noise(
+            noisy_latents,
+            timestep,
+            batch,
+            conditional=conditional,
+            cfg_uncond=cfg_uncond,
+            attn_kind=attn_kind,
+        )
+        sigma_video, sigma_audio = self._noise_amounts(timestep)
+        noisy_video, noisy_audio = self.unpack_latents(noisy_latents)
+        pred_video, pred_audio = self.unpack_latents(pred_noise)
+        return self.pack_latents(
+            self._to_x0(noisy_video, pred_video, sigma_video),
+            self._to_x0(noisy_audio, pred_audio, sigma_audio),
+        )
+
+    @staticmethod
+    def _to_x0(
+        noisy: torch.Tensor,
+        pred_noise: torch.Tensor,
+        sigma: torch.Tensor,
+    ) -> torch.Tensor:
+        # noisy = (1 - sigma) * clean + sigma * noise and pred approximates
+        # noise - clean, so clean = noisy - sigma * pred.
+        sigma = sigma.to(device=noisy.device, dtype=noisy.dtype)
+        return noisy - sigma * pred_noise
+
+    # ------------------------------------------------------------------
+    # Intermediate-latent visualization (LatentVisCallback)
+    # ------------------------------------------------------------------
+
+    def _load_vis_vae(self) -> Any:
+        """Lazily load the H3 video VAE for visualization decodes.
+
+        The module stays CPU-resident between decodes; ``decode_vis_latents``
+        moves it to the GPU per call. Loading mirrors the H3 preprocess
+        scripts: the inference component registry keeps precision policy and
+        normalization identical to the published decode recipe.
+        """
+        vae = getattr(self, "_vis_vae_module", None)
+        if vae is not None:
+            return vae
+        import os
+
+        from fastvideo.fastvideo_args import FastVideoArgs
+        from fastvideo.models.loader.component_loader import PipelineComponentLoader
+        from fastvideo.utils import verify_model_config_and_directory
+
+        model_index = verify_model_config_and_directory(self._init_from)
+        transformers_or_diffusers, _ = model_index["vae"][:2]
+        args = FastVideoArgs(
+            model_path=self._init_from,
+            pipeline_config=self.training_config.pipeline_config,
+            num_gpus=1,
+            tp_size=1,
+            sp_size=1,
+            hsdp_shard_dim=1,
+            use_fsdp_inference=False,
+            vae_cpu_offload=True,
+            text_encoder_cpu_offload=True,
+        )
+        vae = PipelineComponentLoader.load_module(
+            module_name="vae",
+            component_model_path=os.path.join(self._init_from, "vae"),
+            transformers_or_diffusers=transformers_or_diffusers,
+            fastvideo_args=args,
+        )
+        vae.to("cpu")
+        self._vis_vae_module = vae
+        return vae
+
+    @torch.no_grad()
+    def decode_vis_latents(self, packed: torch.Tensor) -> Any:
+        """Decode the packed video stream into a uint8 ``[B, T, C, H, W]`` clip.
+
+        Follows ``MiniMaxH3VideoDecodingStage``: denormalize latents, decode
+        under the published FP16-autocast-over-FP32 recipe, denormalize
+        pixels. The audio stream is dropped — the tracker artifact is a
+        silent video.
+        """
+        video_latents, _ = self.unpack_latents(packed.detach())
+        latents = video_latents.permute(0, 2, 1, 3, 4).to(device=self.device, dtype=torch.float32)
+        vae = self._load_vis_vae()
+        vae.to(self.device)
+        try:
+            latents = vae.denormalize_latents(latents)
+            with torch.autocast(self.device.type, dtype=torch.float16, enabled=self.device.type == "cuda"):
+                video = vae.decode(latents).sample
+            video = vae.denormalize_pixels(video.float()).clamp_(0.0, 1.0).cpu()
+        finally:
+            vae.to("cpu")
+        video = video.permute(0, 2, 1, 3, 4)
+        return (video * 255.0).to(torch.uint8).numpy()
+
+
+__all__ = ["MiniMaxH3DMDModel"]

fastvideo/train/trainer.py

@@ -11,6 +11,7 @@ import torch
 from tqdm.auto import tqdm
 
 from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.logger import init_logger
 from fastvideo.train.callbacks.callback import CallbackDict
 from fastvideo.train.methods.base import LogScalar, TrainingMethod
 from fastvideo.train.utils.tracking import build_tracker
@@ -19,6 +20,41 @@ if TYPE_CHECKING:
     from fastvideo.train.utils.training_config import (
         TrainingConfig, )
 
+logger = init_logger(__name__)
+
+
+def _verify_master_weight_precision(method: TrainingMethod, tc: TrainingConfig) -> None:
+    """Refuse to train on low-precision master weights unless opted in.
+
+    Optimizer steps applied in-place to bf16/fp16 parameters round away
+    updates below ~half an ulp of each weight's magnitude; O(1)-magnitude
+    parameters (norm gains) freeze entirely at typical distillation learning
+    rates, and ``zeros_like``-allocated optimizer state inherits the same
+    starved dtype. fp32 sharded masters (``training.dit_precision: fp32``)
+    keep compute in bf16 via FSDP's param_dtype while fixing both.
+    """
+    if bool(getattr(tc.model, "allow_low_precision_master_weights", False)):
+        return
+    offenders: list[str] = []
+    for role, model in getattr(method, "_role_models", {}).items():
+        if not getattr(model, "_trainable", False):
+            continue
+        transformer = getattr(model, "transformer", None)
+        if transformer is None:
+            continue
+        for name, param in transformer.named_parameters():
+            if param.requires_grad and param.dtype != torch.float32:
+                offenders.append(f"{role}:{name} ({param.dtype})")
+                break
+    if offenders:
+        raise RuntimeError("Trainable master weights are not fp32: "
+                           f"{offenders}. bf16/fp16 parameter storage silently rounds away "
+                           "optimizer updates below ~half an ulp per weight (norm-scale "
+                           "parameters freeze completely). Set training.dit_precision: fp32 "
+                           "(fp32 sharded masters; compute stays bf16 via FSDP param_dtype), "
+                           "or acknowledge the effect explicitly with "
+                           "training.model.allow_low_precision_master_weights: true.")
+
 
 def _coerce_log_scalar(
     value: Any,
@@ -114,6 +150,7 @@ class Trainer:
         )
 
         method.set_tracker(self.tracker)
+        _verify_master_weight_precision(method, tc)
         method.on_train_start()
         self.callbacks.on_train_start(
             method,
@@ -127,6 +164,11 @@ class Trainer:
             resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint)))
             if resumed_step is not None:
                 start_step = int(resumed_step)
+                if bool(getattr(tc.checkpoint, "reset_lr_on_resume", False)):
+                    # The DCP load above restored the checkpoint's optimizer
+                    # LRs and scheduler base_lrs; re-apply the YAML's values.
+                    method.apply_configured_lrs()
+                    logger.info("reset_lr_on_resume: re-applied configured learning rates at step %s", start_step)
         self.callbacks.on_validation_begin(
             method,
             iteration=start_step,

fastvideo/train/utils/checkpoint.py

@@ -62,7 +62,10 @@ def _find_latest_checkpoint(output_dir: Path) -> Path | None:
             continue
         if not _CHECKPOINT_DIR_RE.match(child.name):
             continue
-        if not (child / "dcp").is_dir():
+        # dcp.save writes .metadata last — its presence is the completion
+        # marker. A checkpoint dir from a crashed save (e.g. ENOSPC mid-write)
+        # has dcp/ but no .metadata; resuming from it would fail at boot.
+        if not (child / "dcp" / ".metadata").is_file():
             continue
         try:
             step = _parse_step_from_dir(child)
@@ -184,6 +187,9 @@ class _CallbackStateWrapper:
 class CheckpointConfig:
     save_steps: int
     keep_last: int
+    # Suppress periodic saves before this step (early checkpoints of a long
+    # run are rarely useful and cost ~100 GiB each). 0 disables the gate.
+    start_step: int = 0
 
 
 class CheckpointManager:
@@ -234,6 +240,8 @@ class CheckpointManager:
         save_steps = int(self.config.save_steps or 0)
         if save_steps <= 0:
             return
+        if step < int(self.config.start_step or 0):
+            return
         if step % save_steps != 0:
             return
         if self._last_saved_step == step:

fastvideo/train/utils/config.py

@@ -422,6 +422,8 @@ def _build_training_config(
             resume_from_checkpoint=str(ck.get("resume_from_checkpoint", "") or ""),
             training_state_checkpointing_steps=int(ck.get("training_state_checkpointing_steps", 0) or 0),
             checkpoints_total_limit=int(ck.get("checkpoints_total_limit", 0) or 0),
+            checkpointing_start_step=int(ck.get("checkpointing_start_step", 0) or 0),
+            reset_lr_on_resume=bool(ck.get("reset_lr_on_resume", False)),
         ),
         tracker=TrackerConfig(
             trackers=list(tr.get("trackers", []) or []),
@@ -439,6 +441,7 @@ def _build_training_config(
             precondition_outputs=bool(m.get("precondition_outputs", False)),
             moba_config=dict(m.get("moba_config", {}) or {}),
             enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
+            allow_low_precision_master_weights=bool(m.get("allow_low_precision_master_weights", False)),
         ),
         pipeline_config=pipeline_config,
         model_path=model_path,

fastvideo/train/utils/moduleloader.py

@@ -73,7 +73,13 @@ def make_inference_args(
     args = _make_training_args(tc, model_path=model_path)
     args.inference_mode = True
     args.mode = ExecutionMode.INFERENCE
-    args.dit_cpu_offload = True
+    # Never CPU-offload the DiT here: validation samples with the LIVE
+    # training transformer, and the denoising stages' post-sampling
+    # ``transformer.to("cpu")`` strands the FSDP-sharded training params on
+    # CPU — the next training backward then dies assigning CUDA grads to
+    # CPU tensors (no-grad forwards keep working off gathered buffers,
+    # which is why it surfaces steps later).
+    args.dit_cpu_offload = False
     args.VSA_sparsity = tc.vsa_sparsity
     return args
 

fastvideo/train/utils/optimizer.py

@@ -18,6 +18,75 @@ if TYPE_CHECKING:
     )
 
 
+class AdamWBeta1Zero(torch.optim.Optimizer):
+    """AdamW specialized for ``beta1 == 0``: identical update, no ``exp_avg``.
+
+    With ``beta1 = 0`` Adam's first moment reduces to the raw gradient
+    (``m_t = g_t``, bias correction 1), so the buffer only doubles optimizer
+    state for nothing — one full parameter-sized tensor per model. The op
+    sequence below mirrors ``torch.optim.AdamW``'s single-tensor path
+    exactly, so the parameter trajectory is bitwise-equivalent to
+    ``AdamW(betas=(0.0, beta2))``.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float,
+        beta2: float,
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+    ) -> None:
+        if not 0.0 <= beta2 < 1.0:
+            raise ValueError(f"Invalid beta2: {beta2}")
+        defaults = dict(lr=float(lr), beta2=float(beta2), eps=float(eps), weight_decay=float(weight_decay))
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            lr = group["lr"]
+            beta2 = group["beta2"]
+            eps = group["eps"]
+            weight_decay = group["weight_decay"]
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                state = self.state[p]
+                if not state:
+                    state["step"] = 0
+                    # Low-precision params must not starve the second moment:
+                    # bf16 v loses g^2 increments below ~0.4% of its running
+                    # magnitude. Keep v in fp32 whenever p is not fp32 (the
+                    # fp32-master path is unchanged and stays bitwise-equal
+                    # to torch.optim.AdamW).
+                    if p.dtype == torch.float32:
+                        state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(p, dtype=torch.float32)
+                state["step"] += 1
+                exp_avg_sq = state["exp_avg_sq"]
+                p.mul_(1 - lr * weight_decay)
+                if p.dtype == torch.float32:
+                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+                    bias_correction2_sqrt = (1 - beta2**state["step"])**0.5
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+                    p.addcdiv_(grad, denom, value=-lr)
+                else:
+                    grad_f = grad.float()
+                    exp_avg_sq.mul_(beta2).addcmul_(grad_f, grad_f, value=1 - beta2)
+                    bias_correction2_sqrt = (1 - beta2**state["step"])**0.5
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+                    update = grad_f.div_(denom).mul_(-lr)
+                    p.add_(update.to(p.dtype))
+        return loss
+
+
 def build_optimizer_and_scheduler(
     *,
     params: list[torch.nn.Parameter],
@@ -36,13 +105,24 @@ def build_optimizer_and_scheduler(
         raise ValueError("No trainable parameters passed to "
                          "build_optimizer_and_scheduler")
 
-    optimizer = torch.optim.AdamW(
-        params,
-        lr=float(learning_rate),
-        betas=betas,
-        weight_decay=float(optimizer_config.weight_decay),
-        eps=1e-8,
-    )
+    if float(betas[0]) == 0.0:
+        # beta1 == 0 makes Adam's first-moment buffer redundant; skip it to
+        # save one parameter-sized state tensor per trainable model.
+        optimizer: torch.optim.Optimizer = AdamWBeta1Zero(
+            params,
+            lr=float(learning_rate),
+            beta2=float(betas[1]),
+            weight_decay=float(optimizer_config.weight_decay),
+            eps=1e-8,
+        )
+    else:
+        optimizer = torch.optim.AdamW(
+            params,
+            lr=float(learning_rate),
+            betas=betas,
+            weight_decay=float(optimizer_config.weight_decay),
+            eps=1e-8,
+        )
 
     scheduler = get_scheduler(
         str(scheduler_name),

fastvideo/train/utils/tracking.py

@@ -6,6 +6,7 @@ import os
 from typing import Any, TYPE_CHECKING
 
 from fastvideo.distributed import get_world_group
+from fastvideo.logger import init_logger
 from fastvideo.training.trackers import (
     initialize_trackers,
     Trackers,
@@ -17,6 +18,25 @@ if TYPE_CHECKING:
         TrackerConfig,
     )
 
+logger = init_logger(__name__)
+
+
+def _wandb_usable() -> bool:
+    """True when wandb can actually start a run (importable + credentials)."""
+    try:
+        import wandb
+    except Exception:
+        return False
+    if os.environ.get("WANDB_API_KEY"):
+        return True
+    if os.environ.get("WANDB_MODE") in ("offline", "disabled"):
+        return True
+    try:
+        # Covers ~/.netrc logins from a prior `wandb login`.
+        return wandb.api.api_key is not None
+    except Exception:
+        return False
+
 
 def build_tracker(
     tracker_config: TrackerConfig,
@@ -33,6 +53,11 @@ def build_tracker(
         trackers.append(Trackers.WANDB.value)
     if world_group.rank != 0:
         trackers = []
+    if Trackers.WANDB.value in trackers and not _wandb_usable():
+        logger.warning("wandb tracking requested but wandb is not importable "
+                       "or no credentials are configured (WANDB_API_KEY); "
+                       "continuing without the wandb tracker.")
+        trackers = [t for t in trackers if t != Trackers.WANDB.value]
 
     tracker_log_dir = (checkpoint_config.output_dir or os.getcwd())
     if trackers:
@@ -43,11 +68,40 @@ def build_tracker(
     tracker_run_name = tracker_config.run_name or None
     project = (tracker_config.project_name or "fastvideo")
 
-    return initialize_trackers(
-        trackers,
-        experiment_name=project,
-        config=tracker_config_dict,
-        log_dir=tracker_log_dir,
-        entity=tracker_entity,
-        run_name=tracker_run_name,
-    )
+    try:
+        return initialize_trackers(
+            trackers,
+            experiment_name=project,
+            config=tracker_config_dict,
+            log_dir=tracker_log_dir,
+            entity=tracker_entity,
+            run_name=tracker_run_name,
+        )
+    except Exception as exc:
+        if Trackers.WANDB.value not in trackers:
+            raise
+        # A revoked API key or unreachable api.wandb.ai passes _wandb_usable()
+        # (it only proves a key exists) and then throws inside wandb.init —
+        # which must not kill a multi-node run at boot. Offline init never
+        # contacts the API; the run stays syncable later via `wandb sync`.
+        logger.warning("Tracker init failed (%s); retrying wandb in offline mode.", exc)
+        os.environ["WANDB_MODE"] = "offline"
+        try:
+            return initialize_trackers(
+                trackers,
+                experiment_name=project,
+                config=tracker_config_dict,
+                log_dir=tracker_log_dir,
+                entity=tracker_entity,
+                run_name=tracker_run_name,
+            )
+        except Exception as offline_exc:
+            logger.warning("Offline tracker init also failed (%s); continuing without trackers.", offline_exc)
+            return initialize_trackers(
+                [],
+                experiment_name=project,
+                config=None,
+                log_dir=tracker_log_dir,
+                entity=None,
+                run_name=None,
+            )

fastvideo/train/utils/training_config.py

@@ -58,6 +58,11 @@ class CheckpointConfig:
     resume_from_checkpoint: str = ""
     training_state_checkpointing_steps: int = 0
     checkpoints_total_limit: int = 0
+    checkpointing_start_step: int = 0
+    # DCP checkpoints restore optimizer param-group LRs and scheduler
+    # base_lrs, silently overriding the YAML on resume. Set true to re-apply
+    # the configured learning rates after loading (LR-change experiments).
+    reset_lr_on_resume: bool = False
 
 
 @dataclass(slots=True)
@@ -77,6 +82,13 @@ class ModelTrainingConfig:
     precondition_outputs: bool = False
     moba_config: dict = field(default_factory=dict)
     enable_gradient_checkpointing_type: str | None = None
+    # Optimizer steps applied to bf16/fp16 parameter storage round away
+    # updates smaller than ~half an ulp of each weight's magnitude —
+    # O(1)-magnitude parameters (norm gains) freeze entirely at typical
+    # distillation LRs. The trainer refuses to start unless master weights
+    # are fp32 (training.dit_precision: fp32) or this explicit opt-in
+    # acknowledges the effect (memory-constrained topologies).
+    allow_low_precision_master_weights: bool = False
 
 
 @dataclass(slots=True)

fastvideo/training/distillation_pipeline.py

@@ -33,7 +33,8 @@ from fastvideo.training.activation_checkpoint import (apply_activation_checkpoin
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
                                                get_scheduler, load_distillation_checkpoint,
-                                               save_distillation_checkpoint, shift_timestep)
+                                               save_distillation_checkpoint, shift_timestep,
+                                               verify_master_weight_precision)
 from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory)
 
 try:
@@ -1282,6 +1283,12 @@ class DistillationPipeline(TrainingPipeline):
         """Main training loop with distillation-specific logging."""
         assert self.training_args.seed is not None, "seed must be set"
         seed = self.training_args.seed
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+            "fake_score_transformer": getattr(self, "fake_score_transformer", None),
+            "fake_score_transformer_2": getattr(self, "fake_score_transformer_2", None),
+        })
 
         # Set the same seed within each SP group to ensure reproducibility
         if self.sp_world_size > 1:

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -21,7 +21,7 @@ from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfF
 from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines import TrainingBatch
 from fastvideo.training.distillation_pipeline import DistillationPipeline
-from fastvideo.training.training_utils import (EMA_FSDP, save_distillation_checkpoint)
+from fastvideo.training.training_utils import (EMA_FSDP, save_distillation_checkpoint, verify_master_weight_precision)
 from fastvideo.utils import is_vsa_available, set_random_seed
 from fastvideo.profiler import profile_region
 
@@ -797,6 +797,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         """Main training loop with self-forcing specific logging."""
         assert self.training_args.seed is not None, "seed must be set"
         seed = self.training_args.seed
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+            "fake_score_transformer": getattr(self, "fake_score_transformer", None),
+            "fake_score_transformer_2": getattr(self, "fake_score_transformer_2", None),
+        })
 
         # Set the same seed within each SP group to ensure reproducibility
         if self.sp_world_size > 1:

fastvideo/training/training_pipeline.py

@@ -42,7 +42,8 @@ from fastvideo.training.activation_checkpoint import (apply_activation_checkpoin
 from fastvideo.training.trackers import (DummyTracker, TrackerType, initialize_trackers, Trackers)
 from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases,
                                                compute_density_for_timestep_sampling, count_trainable, get_scheduler,
-                                               get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint)
+                                               get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint,
+                                               verify_master_weight_precision)
 from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict)
 
 try:
@@ -519,6 +520,10 @@ class TrainingPipeline(LoRAPipeline, ABC):
         logger.info('rank: %s: start training', self.global_rank, local_main_process_only=False)
         if not self.post_init_called:
             self.post_init()
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+        })
         num_trainable_params = count_trainable(self.transformer)
         logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3))
 

fastvideo/training/training_utils.py

@@ -23,6 +23,36 @@ logger = init_logger(__name__)
 _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
 
 
+def verify_master_weight_precision(modules: dict[str, torch.nn.Module | None]) -> None:
+    """Refuse to train on low-precision master weights unless opted in.
+
+    Optimizer steps applied in-place to bf16/fp16 parameters round away
+    updates below ~half an ulp of each weight's magnitude; O(1)-magnitude
+    parameters (norm gains) freeze entirely at typical fine-tune/distillation
+    learning rates, and ``zeros_like``-allocated optimizer state inherits the
+    same starved dtype. ``--dit_precision fp32`` gives fp32 sharded masters
+    while compute stays bf16 via FSDP's ``param_dtype``.
+    """
+    if os.environ.get("FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS", "0") == "1":
+        return
+    offenders: list[str] = []
+    for module_name, module in modules.items():
+        if module is None:
+            continue
+        for name, param in module.named_parameters():
+            if param.requires_grad and param.dtype != torch.float32:
+                offenders.append(f"{module_name}:{name} ({param.dtype})")
+                break
+    if offenders:
+        raise RuntimeError("Trainable master weights are not fp32: "
+                           f"{offenders}. bf16/fp16 parameter storage silently rounds away "
+                           "optimizer updates below ~half an ulp per weight (norm-scale "
+                           "parameters freeze completely). Pass --dit_precision fp32 "
+                           "(fp32 sharded masters; compute stays bf16 via FSDP param_dtype), "
+                           "or acknowledge the effect explicitly with "
+                           "FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS=1.")
+
+
 def gather_state_dict_on_cpu_rank0(
     model,
     device: torch.device | None = None,

pyproject.toml

@@ -208,7 +208,8 @@ dreamverse-mock-server = "dreamverse.mock_server:cli"
 
 [tool.setuptools.packages.find]
 where = [".", "apps/dreamverse"]
-exclude = ["assets*", "docker*", "docs", "scripts*", "apps*", "dreamverse.tests*"]
+include = ["fastvideo", "fastvideo.*", "dreamverse", "dreamverse.*"]
+exclude = ["dreamverse.tests*"]
 
 [tool.setuptools.package-data]
 dreamverse = ["prompts/*.md"]