diff --git a/docs-site/navigation.json b/docs-site/navigation.json index 6e2faa661..88c2df91c 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -314,6 +314,10 @@ "label": "Attn-QAT Training", "href": "/training/attn_qat" }, + { + "label": "FastH3 NVFP4 QAD", + "href": "/training/minimax_h3_fp4_qad" + }, { "label": "HunyuanVideo 1.5 T2V", "href": "/training/hunyuan15_overfit" diff --git a/docs/training/minimax_h3_fp4_qad.md b/docs/training/minimax_h3_fp4_qad.md new file mode 100644 index 000000000..9544793c7 --- /dev/null +++ b/docs/training/minimax_h3_fp4_qad.md @@ -0,0 +1,253 @@ +# FastH3 NVFP4 QAD + +This page describes the shared quantization-aware distillation (QAD) recipe for +the FastH3 (MiniMax-H3) family. One recipe covers two model profiles: + +| Profile | Model plugin | Task | Sampling | VSA tile | +|---|---|---|---|---| +| OmniRef | `omniref_pdd.MiniMaxH3OmniRefPDDModel` | Ref2VA | 8-forward PDD ladder | 128 | +| V2 T2AV QAD | `v2_t2v.MiniMaxH3V2T2VModel` | T2VA | 8-step DMD | 64 | + +Both profiles use the same shared recipe: NVFP4 W4A4 linears in all 50 blocks, +the FP4 sparse-attention emulator, and the same objective and guards. QAD +trains the student against a frozen bf16 teacher so that quantization error is +absorbed by the weights instead of appearing as drift, ghosting or audio loss +at inference. The published OmniRef QAD checkpoint includes FP4 attention; the +published FastH3-QAD-v2 checkpoint predates it and keeps BF16 VSA attention (see +[Released checkpoints](#released-checkpoints)). + +The runnable configs are: + +- `examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml` +- `examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml` + +## Released checkpoints + +| Role | Checkpoint | +|---|---| +| OmniRef QAD result | `FastVideo/FastH3-OmniRef-NVFP4-QAD` | +| FastH3-QAD-v2 result | `FastVideo/FastH3-QAD-v2` | +| V2 bf16 teacher / starting weights | `FastVideo/FastVideo-FastH3-8-Step-V2` | +| NVFP4 consumer baseline | `FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer` | + +The OmniRef config starts from the **calibrated OmniRef checkpoint used as QAD +step 0**, not from the released OmniRef QAD result; point `init_from` at that +checkpoint. The V2 config starts from the bf16 V2 snapshot. + +The published FastH3-QAD-v2 checkpoint is the **linears-only variant** of this +schedule: NVFP4 W4A4 linears with BF16 VSA attention, 300 steps. The shared +FP4 recipe on this page additionally runs the FP4 attention emulator from step +0; it is the same training code and objective, but it does not reproduce the +published FastH3-QAD-v2 checkpoint bit for bit. Set +`models.student.qad.fp4_attention.enabled: false` to get the published +linears-only variant. + +## Target numerics + +The student starts from the bf16 masters, and the deployed numerics are active +from the first forward, so every update sees the quantization error it is +trained to absorb. + +- **Linears:** NVFP4 W4A4 for `attn.to_q`, `attn.to_k`, `attn.to_v`, + `attn.to_out`, `ff.fc_in` and `ff.fc_out` in every block. + - FFN inputs use static activation scales from the calibrated `amax` table. + - Attention and gate linears use the unit activation scale; their `amax` + entries are optional. + - `skip_blocks` is empty: no block falls back to bf16 linears. +- **Attention:** the FP4 sparse-attention emulator for VSA tile 128 (OmniRef) + or tile 64 (V2 T2AV). Two-level P quantization and K smoothing are on. The + emulator reproduces the deployed kernel's block order and rounding on any + CUDA GPU and provides straight-through gradients. +- **Kept frozen:** the VSA compression gates and every norm stay at their + released values. Training the gate weights drifts the sparse-attention tile + selection and produces ghosting; the gate linears still execute the deployed + NVFP4 format, they are just not updated. `proj_in`, the time embedder, + `proj_out` and the audio projections stay fp32 and frozen, matching + inference. + +The calibrated `amax` tables for both profiles are vendored under +`examples/train/calibration/fasth3/`: + +| Table | sha256 | +|---|---| +| `amax-omniref-mse-bf16-top8.json` | `14b93ec3…df083d` | +| `amax-consumer.json` | `0800516c…d39e52` | + +`qad.linears.amax_sha256` pins the table; the run copies it into the output +directory and refuses to start if the bytes differ. + +## Training objective + +For each row the frozen bf16 teacher walks the profile's sampling ladder. At +the same inputs and timesteps, the student minimizes: + +- **Output regression:** per-forward x0 error, normalized per forward and + modality by its step-0 value. +- **Layer-wise reconstruction:** each attention and FFN sub-layer receives the + teacher's input (stop-gradient) and matches the teacher sub-layer's output. + The loss is weighted per row group (target video and audio 1.0, + conditioning 0.25) because audio is a small fraction of the packed tokens. + `recon_normalize: step0` divides it by its step-0 value so `recon_lambda` + balances the two terms. + +From `rollout_start_step`, `rollout_fraction` of the rows continue from the +student's own trajectory instead of the teacher's, which reduces exposure bias +in the sampled output. + +## Evaluation and guards + +- **Teacher-forced eval** every `eval_every` steps on held-out rows. +- **Full eval** every `full_eval_every` steps: the student's own sampled output + against the teacher at the same seed, keyframe alignment, and saved audio and + video latents. +- **Noise floor:** `noise_floor_eval` re-runs the step-0 eval at seed + 1 so a + metric move can be compared with teacher seed noise. +- **Best checkpoint:** selected by the teacher-forced `score`. A checkpoint is + eligible only if the keys in `eligibility_keys` (audio error and endpoint) + stay within their step-0 guardrails. +- **Auto-stop:** after `patience` consecutive bad evals. +- **Step-0 cache:** `step0_cache` fingerprints the teacher, student, + calibration, quantization plan, eval set and loss weights, so reruns skip the + baseline evaluation. +- **Resume:** checkpoints carry the step-0 baseline and guard state; do not + delete `method.qad_baseline` if you edit a checkpoint by hand. + +## Data preparation + +V2 T2AV is data-free: it needs prompts and their Qwen3-VL text embeddings, not +VAE latents. + +1. Build the split manifests from prompt sources (parquet shards with a + `prompt` column, JSONL with a `prompt` field): + + ```bash + python scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py \ + --parquet 'data/prompts/*.parquet' --out-dir data/fasth3-v2 --eval-count 30 + ``` + + This writes `train.jsonl` and `eval.jsonl` plus an `embeddings/` directory. +2. Fill the embeddings (safe to rerun and shard across GPUs): + + ```bash + python scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py \ + --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml \ + --model-path data/models/FastH3-V2 \ + --manifests data/fasth3-v2/train.jsonl data/fasth3-v2/eval.jsonl \ + --shard 0 --num-shards 8 + ``` + + The V2 config's `train_manifest` and `eval_manifest` point at these files. + +OmniRef training rows are preprocessed Ref2VA parquet rows that carry the +reference latents, text embeddings and geometry; the encoded dataset itself is +an input to this recipe and is not produced by the training code. Provide a +`MANIFEST.jsonl` with `case`, `id`, `parquet` and `resolution` fields and point +`models.student.data.manifests` at it. + +`scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py` builds the +held-out eval plan and runs the PDD ladder once per row: use `--eval-dir` to +write the held-out `manifest.json` consumed by `models.student.data.eval_manifest` +(and the raw latents for decoding). + +## Running + +Fill in the placeholder paths in the config (starting checkpoint, data +manifests, output directory), then launch with `torchrun`. + +OmniRef, 4 nodes x 4 GPUs (sequence parallel 4): + +```bash +torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \ + --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \ + -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml +``` + +V2 T2AV, 4 nodes x 4 GPUs (16-way HSDP, lr 5e-6, 300 steps): + +```bash +torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \ + --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \ + -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml +``` + +To reproduce the published FastH3-QAD-v2 linears-only checkpoint, disable FP4 attention +for the same schedule: + +```bash + --models.student.qad.fp4_attention.enabled false +``` + +## FP4-attention adaptation + +The internal follow-up that this recipe tracks started from a linears-only +300-step DCP, enabled FP4 attention, and trained 300 more steps at lr 1.5e-6 +with no warmup and a constant LR on 16 GPUs. It is an **adaptation of the +linears-only checkpoint to FP4 attention**, not a second epoch of the FP4 +stage 1 above: the source checkpoint's attention numerics were BF16, and the +continuation needs a DCP of that linears-only model (the packed HF release is +not one). It is therefore not part of the public reproduction path, and the +published FastH3-QAD-v2 checkpoint remains the linears-only result. The exact overrides +are listed at the bottom of the V2 config. + +## Checkpoints + +Runs write regular `checkpoint-/` DCP checkpoints plus a `best/` +directory: + +- `best/`: the best eligible state (weights only), refreshed at every eval and + never pruned. +- `checkpoint-/`: weights, optimizer, RNG state and the method baseline; + the last step is always saved. + +The saved state is the bf16 master plus the quantization plan, not packed NVFP4 +weights; inference applies the same converter used for the released model. Use +`fastvideo/train/entrypoint/dcp_to_diffusers.py` to export a checkpoint if your +deployment path expects Diffusers weights. It detects the weights-only marker +on a tagged checkpoint, so `best/` can be exported directly: + +```bash +python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint outputs/v2_1_nvfp4_qad/best \ + --output-dir outputs/v2_1_nvfp4_qad/export +``` + +## Code map + +| Path | Purpose | +|---|---| +| `fastvideo/train/methods/knowledge_distillation/pdd_qad.py` | Method: ladder, losses, evaluation, guards | +| `fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py` | Rung aggregation and score | +| `fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py` | Layer-wise reconstruction hooks | +| `fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py` | Step-0 baseline cache | +| `fastvideo/train/models/minimax_h3/omniref_pdd.py` | OmniRef Ref2VA profile | +| `fastvideo/train/models/minimax_h3/v2_t2v.py` | V2 T2AV profile | +| `fastvideo/layers/quantization/nvfp4_qad.py` | NVFP4 linear emulation and STE | +| `fastvideo/attention/backends/fp4_vsa_qat.py` | VSA fine-attention override | +| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py` | Tile 64/128 FP4 block-sparse Triton kernel | +| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py` | PyTorch numeric reference | +| `scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py` | V2 prompt split manifests | +| `scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py` | V2 Qwen3-VL text embeddings | +| `scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py` | OmniRef held-out eval plan and latents | +| `examples/train/calibration/fasth3/` | Pinned FFN activation `amax` tables | + +## Expected behavior + +Teacher-forced video and audio error both fall over the first tens of steps, +and the student's own video samples improve beyond the teacher's seed noise. +Audio samples improve more slowly and can remain within noise even when +teacher-forced audio error clearly drops. Gains slow after roughly 60-100 +steps; the remaining error is dominated by FP4 attention. Evaluate a checkpoint +with the full-eval samples, not the teacher-forced score alone, before +promoting it. + +Released reference points (lower is better): + +| Run | Metric | Baseline | Selected | +|---|---|---|---| +| FastH3-QAD-v2 (linears-only) | mean per-forward x0 rel-L2 | 0.431 | 0.351 (video 0.382, audio 0.319) | +| OmniRef QAD | teacher-forced score | 0.362 | 0.303 at step 120 (video 0.198, audio 0.409) | + +The FastH3-QAD-v2 baseline is the NVFP4 Consumer model before QAD; the OmniRef baseline +is the calibrated checkpoint at QAD step 0. diff --git a/examples/train/calibration/fasth3/amax-consumer.json b/examples/train/calibration/fasth3/amax-consumer.json new file mode 100644 index 000000000..0058cf347 --- /dev/null +++ b/examples/train/calibration/fasth3/amax-consumer.json @@ -0,0 +1,102 @@ +{ + "b0.ff.fc_in": 747.9899299279987, + "b0.ff.fc_out": 2462.017121214983, + "b1.ff.fc_in": 671.2121750026851, + "b1.ff.fc_out": 1285.509862734237, + "b2.ff.fc_in": 1153.55816110724, + "b2.ff.fc_out": 3796.948980859055, + "b3.ff.fc_in": 1153.55816110724, + "b3.ff.fc_out": 4715.270819259384, + "b4.ff.fc_in": 1285.509862734237, + "b4.ff.fc_out": 4715.270819259384, + "b5.ff.fc_in": 2209.302022692089, + "b5.ff.fc_out": 5254.635977347013, + "b6.ff.fc_in": 2209.302022692089, + "b6.ff.fc_out": 5254.635977347013, + "b7.ff.fc_in": 2462.017121214983, + "b7.ff.fc_out": 4715.270819259384, + "b8.ff.fc_in": 2462.017121214983, + "b8.ff.fc_out": 5855.696365467314, + "b9.ff.fc_in": 2743.6392141259485, + "b9.ff.fc_out": 5855.696365467314, + "b10.ff.fc_in": 2743.6392141259485, + "b10.ff.fc_out": 5855.696365467314, + "b11.ff.fc_in": 2743.6392141259485, + "b11.ff.fc_out": 6525.510437495776, + "b12.ff.fc_in": 2743.6392141259485, + "b12.ff.fc_out": 6525.510437495776, + "b13.ff.fc_in": 3057.4751817363253, + "b13.ff.fc_out": 6525.510437495776, + "b14.ff.fc_in": 2743.6392141259485, + "b14.ff.fc_out": 6525.510437495776, + "b15.ff.fc_in": 2743.6392141259485, + "b15.ff.fc_out": 6525.510437495776, + "b16.ff.fc_in": 3407.2097180654623, + "b16.ff.fc_out": 6525.510437495776, + "b17.ff.fc_in": 3057.4751817363253, + "b17.ff.fc_out": 9030.717912847476, + "b18.ff.fc_in": 2743.6392141259485, + "b18.ff.fc_out": 9030.717912847476, + "b19.ff.fc_in": 2743.6392141259485, + "b19.ff.fc_out": 9030.717912847476, + "b20.ff.fc_in": 3057.4751817363253, + "b20.ff.fc_out": 10063.712585454898, + "b21.ff.fc_in": 3057.4751817363253, + "b21.ff.fc_out": 10063.712585454898, + "b22.ff.fc_in": 3407.2097180654623, + "b22.ff.fc_out": 10063.712585454898, + "b23.ff.fc_in": 3407.2097180654623, + "b23.ff.fc_out": 12497.700799896355, + "b24.ff.fc_in": 3796.948980859055, + "b24.ff.fc_out": 13927.271620769527, + "b25.ff.fc_in": 3796.948980859055, + "b25.ff.fc_out": 19274.08888735978, + "b26.ff.fc_in": 4231.26976101764, + "b26.ff.fc_out": 21478.789600085922, + "b27.ff.fc_in": 4715.270819259384, + "b27.ff.fc_out": 29724.70802687751, + "b28.ff.fc_in": 4715.270819259384, + "b28.ff.fc_out": 45841.764386937015, + "b29.ff.fc_in": 4715.270819259384, + "b29.ff.fc_out": 56928.958163414085, + "b30.ff.fc_in": 5254.635977347013, + "b30.ff.fc_out": 78784.54787887946, + "b31.ff.fc_in": 5254.635977347013, + "b31.ff.fc_out": 109030.7109780948, + "b32.ff.fc_in": 5855.696365467314, + "b32.ff.fc_out": 135400.63235256448, + "b33.ff.fc_in": 5855.696365467314, + "b33.ff.fc_out": 168148.35127510442, + "b34.ff.fc_in": 6525.510437495776, + "b34.ff.fc_out": 168148.35127510442, + "b35.ff.fc_in": 6525.510437495776, + "b35.ff.fc_out": 232702.1493904285, + "b36.ff.fc_in": 6525.510437495776, + "b36.ff.fc_out": 232702.1493904285, + "b37.ff.fc_in": 6525.510437495776, + "b37.ff.fc_out": 208816.33399986484, + "b38.ff.fc_in": 6525.510437495776, + "b38.ff.fc_out": 288982.97321111924, + "b39.ff.fc_in": 7271.942168450891, + "b39.ff.fc_out": 358875.77201571054, + "b40.ff.fc_in": 7271.942168450891, + "b40.ff.fc_out": 445672.69186462485, + "b41.ff.fc_in": 8103.755565259112, + "b41.ff.fc_out": 496651.7297372453, + "b42.ff.fc_in": 9030.717912847476, + "b42.ff.fc_out": 616770.844395513, + "b43.ff.fc_in": 9030.717912847476, + "b43.ff.fc_out": 616770.844395513, + "b44.ff.fc_in": 10063.712585454898, + "b44.ff.fc_out": 687321.2672027495, + "b45.ff.fc_in": 11214.867875490954, + "b45.ff.fc_out": 853555.2832070139, + "b46.ff.fc_in": 12497.700799896355, + "b46.ff.fc_out": 1059994.2350776498, + "b47.ff.fc_in": 13927.271620769527, + "b47.ff.fc_out": 1316362.0994398587, + "b48.ff.fc_in": 17295.691014283653, + "b48.ff.fc_out": 2030107.8425012752, + "b49.ff.fc_in": 19274.08888735978, + "b49.ff.fc_out": 2809486.555264971 +} \ No newline at end of file diff --git a/examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json b/examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json new file mode 100644 index 000000000..1f63628e6 --- /dev/null +++ b/examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json @@ -0,0 +1,102 @@ +{ + "transformer_blocks.0.ff.fc_in": 7.34375, + "transformer_blocks.0.ff.fc_out": 6432.0, + "transformer_blocks.1.ff.fc_in": 6.3125, + "transformer_blocks.1.ff.fc_out": 8960.0, + "transformer_blocks.10.ff.fc_in": 15.25, + "transformer_blocks.10.ff.fc_out": 6368.0, + "transformer_blocks.11.ff.fc_in": 17.75, + "transformer_blocks.11.ff.fc_out": 16128.0, + "transformer_blocks.12.ff.fc_in": 18.549999999999997, + "transformer_blocks.12.ff.fc_out": 4064.0, + "transformer_blocks.13.ff.fc_in": 20.25, + "transformer_blocks.13.ff.fc_out": 78848.0, + "transformer_blocks.14.ff.fc_in": 12.0, + "transformer_blocks.14.ff.fc_out": 1784.0, + "transformer_blocks.15.ff.fc_in": 13.1875, + "transformer_blocks.15.ff.fc_out": 1400.0, + "transformer_blocks.16.ff.fc_in": 12.000625, + "transformer_blocks.16.ff.fc_out": 2768.16, + "transformer_blocks.17.ff.fc_in": 22.049999999999997, + "transformer_blocks.17.ff.fc_out": 4960.0, + "transformer_blocks.18.ff.fc_in": 8.875, + "transformer_blocks.18.ff.fc_out": 15616.0, + "transformer_blocks.19.ff.fc_in": 12.4375, + "transformer_blocks.19.ff.fc_out": 3203.2, + "transformer_blocks.2.ff.fc_in": 4.2875, + "transformer_blocks.2.ff.fc_out": 7648.0, + "transformer_blocks.20.ff.fc_in": 8.96875, + "transformer_blocks.20.ff.fc_out": 5696.0, + "transformer_blocks.21.ff.fc_in": 10.0625, + "transformer_blocks.21.ff.fc_out": 6464.0, + "transformer_blocks.22.ff.fc_in": 14.000625, + "transformer_blocks.22.ff.fc_out": 11712.0, + "transformer_blocks.23.ff.fc_in": 13.125, + "transformer_blocks.23.ff.fc_out": 13312.0, + "transformer_blocks.24.ff.fc_in": 18.125, + "transformer_blocks.24.ff.fc_out": 9497.599999999999, + "transformer_blocks.25.ff.fc_in": 10.7625, + "transformer_blocks.25.ff.fc_out": 12864.0, + "transformer_blocks.26.ff.fc_in": 10.443750000000001, + "transformer_blocks.26.ff.fc_out": 11776.0, + "transformer_blocks.27.ff.fc_in": 13.9375, + "transformer_blocks.27.ff.fc_out": 13504.0, + "transformer_blocks.28.ff.fc_in": 13.9375, + "transformer_blocks.28.ff.fc_out": 8448.0, + "transformer_blocks.29.ff.fc_in": 11.4375, + "transformer_blocks.29.ff.fc_out": 13058.559999999998, + "transformer_blocks.3.ff.fc_in": 6.28125, + "transformer_blocks.3.ff.fc_out": 10880.0, + "transformer_blocks.30.ff.fc_in": 12.000625, + "transformer_blocks.30.ff.fc_out": 18816.0, + "transformer_blocks.31.ff.fc_in": 7.75, + "transformer_blocks.31.ff.fc_out": 8705.279999999999, + "transformer_blocks.32.ff.fc_in": 7.9375, + "transformer_blocks.32.ff.fc_out": 13312.0, + "transformer_blocks.33.ff.fc_in": 10.0, + "transformer_blocks.33.ff.fc_out": 8035.200000000001, + "transformer_blocks.34.ff.fc_in": 6.95625, + "transformer_blocks.34.ff.fc_out": 7424.0, + "transformer_blocks.35.ff.fc_in": 9.25, + "transformer_blocks.35.ff.fc_out": 20480.0, + "transformer_blocks.36.ff.fc_in": 9.40625, + "transformer_blocks.36.ff.fc_out": 29440.0, + "transformer_blocks.37.ff.fc_in": 16.0, + "transformer_blocks.37.ff.fc_out": 14272.0, + "transformer_blocks.38.ff.fc_in": 15.751249999999999, + "transformer_blocks.38.ff.fc_out": 28160.0, + "transformer_blocks.39.ff.fc_in": 15.064374999999998, + "transformer_blocks.39.ff.fc_out": 149504.0, + "transformer_blocks.4.ff.fc_in": 6.75, + "transformer_blocks.4.ff.fc_out": 14400.0, + "transformer_blocks.40.ff.fc_in": 11.438125, + "transformer_blocks.40.ff.fc_out": 31232.0, + "transformer_blocks.41.ff.fc_in": 17.25, + "transformer_blocks.41.ff.fc_out": 37383.67999999999, + "transformer_blocks.42.ff.fc_in": 15.0625, + "transformer_blocks.42.ff.fc_out": 15296.0, + "transformer_blocks.43.ff.fc_in": 17.625, + "transformer_blocks.43.ff.fc_out": 24064.0, + "transformer_blocks.44.ff.fc_in": 24.375, + "transformer_blocks.44.ff.fc_out": 99840.0, + "transformer_blocks.45.ff.fc_in": 40.0025, + "transformer_blocks.45.ff.fc_out": 350208.0, + "transformer_blocks.46.ff.fc_in": 42.75749999999999, + "transformer_blocks.46.ff.fc_out": 20070.399999999998, + "transformer_blocks.47.ff.fc_in": 57.2525, + "transformer_blocks.47.ff.fc_out": 49920.0, + "transformer_blocks.48.ff.fc_in": 145.0, + "transformer_blocks.48.ff.fc_out": 96256.0, + "transformer_blocks.49.ff.fc_in": 146.0, + "transformer_blocks.49.ff.fc_out": 200704.0, + "transformer_blocks.5.ff.fc_in": 13.1875, + "transformer_blocks.5.ff.fc_out": 15168.0, + "transformer_blocks.6.ff.fc_in": 19.375, + "transformer_blocks.6.ff.fc_out": 11904.0, + "transformer_blocks.7.ff.fc_in": 33.5, + "transformer_blocks.7.ff.fc_out": 15872.0, + "transformer_blocks.8.ff.fc_in": 16.5, + "transformer_blocks.8.ff.fc_out": 9216.0, + "transformer_blocks.9.ff.fc_in": 19.375, + "transformer_blocks.9.ff.fc_out": 4096.0 +} \ No newline at end of file diff --git a/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml new file mode 100644 index 000000000..013f10673 --- /dev/null +++ b/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml @@ -0,0 +1,135 @@ +# FastH3 OmniRef NVFP4 QAD (the shared FP4 QAD recipe, Ref2VA profile). +# +# Student: OmniRef Ref2VA (8-forward PDD ladder) with the deployed NVFP4 +# numerics active from step 0: NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in +# all 50 blocks, and the FP4 sparse-attention emulator at VSA tile 128. The +# VSA compression gates and every norm stay frozen at the calibrated bf16 +# values; training them drifts the tile selection and causes ghosting. The +# teacher is the same checkpoint in bf16. +# +# See docs/training/minimax_h3_fp4_qad.md for the recipe shared with the V2 +# (T2AV QAD) profile, calibration, and the checkpoint contract. +# +# Released result: FastVideo/FastH3-OmniRef-NVFP4-QAD (step 120 of 200). +# `init_from` is the calibrated OmniRef checkpoint used as QAD step 0, +# not that released result. The calibrated amax table is vendored under +# examples/train/calibration/fasth3/ and pinned by sha256 below. +models: + student: + _target_: fastvideo.train.models.minimax_h3.omniref_pdd.MiniMaxH3OmniRefPDDModel + init_from: data/models/omniref-nvfp4-ptq + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json + # Attention and gate linears use the unit activation scale; only the + # FFN linears read their calibrated amax from this table. + attention: true + gate: true + ffn: true + skip_blocks: [] + amax_sha256: 14b93ec308faf0d89bd13ded507d287c030bb3005c96b8063213eb8a09df083d + fp4_attention: + enabled: true + two_level_p: true + smooth_k: true + first_block_max_floor: null + data: + manifests: + - data/omniref/train-manifest.jsonl + cases: + - first_frame + - first_last_frame + - storyboard + - continue_scene + - continue_shot + - full_scene_audio_image_reference + per_group: + 480p: 400 + 768p: 50 + seed: 20261007 + max_frames: 243 + eval_manifest: data/omniref/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + # Only the projections and FFN linears train: the same set the deployed + # NVFP4 export quantizes. Gates, norms and modulation stay frozen. + trainable_patterns: transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\. + teacher: + _target_: fastvideo.train.models.minimax_h3.omniref_pdd.MiniMaxH3OmniRefPDDModel + init_from: data/models/omniref-nvfp4-ptq + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + rungs_per_row: 4 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/omniref_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/omniref_nvfp4_qad/step0_cache.json + +training: + distributed: + num_gpus: 16 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 16 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 5.0e-06 + betas: [0.9, 0.999] + eps: 1.0e-08 + weight_decay: 0.0 + lr_scheduler: cosine_with_min_lr + lr_warmup_steps: 10 + min_lr_ratio: 0.3 + loop: + max_train_steps: 200 + gradient_accumulation_steps: 1 + checkpoint: + output_dir: outputs/omniref_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: fasth3-omniref-nvfp4 + run_name: omniref-nvfp4-qad + model: + enable_gradient_checkpointing_type: full + dit_precision: fp32 + vsa_sparsity: 0.0 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: {} diff --git a/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml new file mode 100644 index 000000000..8d3c12c0a --- /dev/null +++ b/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml @@ -0,0 +1,149 @@ +# FastH3 V2 NVFP4 QAD (the shared FP4 QAD recipe, T2VA profile). +# +# Student: V2 bf16 weights with the deployed NVFP4 numerics active from step 0: +# NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks, and the FP4 +# sparse-attention emulator at VSA tile 64. The VSA compression gates and every +# norm stay frozen at the V2 bf16 values. The teacher is the same weights in +# bf16. Middle-block token compression is disabled: pure QAD. +# +# This is one recipe with two profiles; see docs/training/minimax_h3_fp4_qad.md +# for the shared objective, evaluation and checkpoint contract, and for the +# OmniRef (Ref2VA, tile 128) config. +# +# Paths are placeholders. Fill in the bf16 V2 model, the calibration table +# shipped with it, the latent manifests and the output dir. +# Starting weights: FastVideo/FastVideo-FastH3-8-Step-V2 (bf16). The released +# QAD result is FastVideo/FastH3-QAD-v2; the +# NVFP4 baseline it replaces is +# FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer. +models: + student: + _target_: fastvideo.train.models.minimax_h3.v2_t2v.MiniMaxH3V2T2VModel + init_from: data/models/FastH3-V2 + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/fasth3/amax-consumer.json + # Attention and gate linears use the unit activation scale; only the + # FFN linears read their calibrated amax from this table. + attention: true + gate: true + ffn: true + skip_blocks: [] + amax_sha256: 0800516caec04ec829d21c45ca8b18eede6b81008c040997c0b84d48e0d39e52 + fp4_attention: + enabled: true + two_level_p: true + smooth_k: true + first_block_max_floor: null + compression: + enabled: false + start: 3 + end: 41 + mode: subsample + data: + train_manifest: data/fasth3-v2/train.jsonl + eval_manifest: data/fasth3-v2/eval.jsonl + width: 832 + height: 480 + num_frames: 124 + # Only the projections and FFN linears train: the same set the deployed + # NVFP4 export quantizes. Gates, norms and modulation stay frozen. + trainable_patterns: transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\. + teacher: + _target_: fastvideo.train.models.minimax_h3.v2_t2v.MiniMaxH3V2T2VModel + init_from: data/models/FastH3-V2 + trainable: false + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + rungs_per_row: 4 + eval_every: 25 + full_eval_every: 50 + patience: 3 + fidelity_tolerance: 0.05 + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/v2_1_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 8 + noise_floor_eval: true + # A checkpoint is eligible for "best" only if audio stays inside its + # step-0 guardrail; the ELBO-style score alone can trade audio for video. + eligibility_keys: + - endpoint/audio + - x0_rel_l2/audio/mean + step0_cache: outputs/v2_1_nvfp4_qad/step0_cache.json + +training: + distributed: + num_gpus: 16 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 16 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 5.0e-06 + betas: [0.9, 0.999] + eps: 1.0e-08 + weight_decay: 0.0 + lr_scheduler: cosine_with_min_lr + lr_warmup_steps: 10 + min_lr_ratio: 0.3 + loop: + max_train_steps: 300 + gradient_accumulation_steps: 1 + checkpoint: + output_dir: outputs/v2_1_nvfp4_qad + training_state_checkpointing_steps: 50 + checkpoints_total_limit: 1 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: fasth3-v2-nvfp4-qad + run_name: v2-nvfp4-qad + model: + enable_gradient_checkpointing_type: full + dit_precision: fp32 + vsa_sparsity: 0.0 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: {} + +# The published FastVideo/FastH3-QAD-v2 +# checkpoint is the linears-only variant of this schedule: NVFP4 linears with +# BF16 VSA attention, i.e. the same config with fp4_attention.enabled=false. +# The shared FP4 recipe above adds FP4 attention from step 0 and does not +# reproduce that checkpoint bit for bit. +# +# FP4-attention adaptation (provenance note, not part of the public +# reproduction). The internal follow-up started from a linears-only 300-step +# DCP, enabled FP4 attention, and trained 300 more steps at lr 1.5e-6 with no +# warmup and a constant LR. It is an adaptation of the linears-only checkpoint +# to FP4 attention, not a second epoch of the FP4 stage 1 above, and it needs +# a DCP of the linears-only model (the packed HF release is not one): +# +# method.init_student_dcp= +# training.optimizer.learning_rate=1.5e-06 +# training.optimizer.lr_warmup_steps=0 +# training.optimizer.min_lr_ratio=1.0 +# training.loop.max_train_steps=300 +# training.checkpoint.output_dir=outputs/v2_1_nvfp4_qad_fp4_adapt diff --git a/fastvideo/attention/backends/fp4_vsa_qat.py b/fastvideo/attention/backends/fp4_vsa_qat.py index db9cb6bb5..6d9b65baf 100644 --- a/fastvideo/attention/backends/fp4_vsa_qat.py +++ b/fastvideo/attention/backends/fp4_vsa_qat.py @@ -105,4 +105,23 @@ class FP4VSAFineAttention: return out.transpose(1, 2) -__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "load_qat_vsa_modules"] +def vsa_impls(transformer: torch.nn.Module) -> list[Any]: + """Every VSA-H3 attention implementation in a MiniMax-H3 transformer.""" + impls = [] + for module in transformer.modules(): + impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None) + if impl is not None and hasattr(impl, "fine_attention_override"): + impls.append(impl) + return impls + + +def install_fp4_vsa_attention(transformer: torch.nn.Module, numerics: FP4AttentionNumerics | None) -> int: + """Route every VSA-H3 layer's fine attention through the emulator (``None`` restores the kernels).""" + override = FP4VSAFineAttention(numerics) if numerics is not None else None + impls = vsa_impls(transformer) + for impl in impls: + impl.fine_attention_override = override + return len(impls) + + +__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "install_fp4_vsa_attention", "load_qat_vsa_modules"] diff --git a/fastvideo/layers/quantization/nvfp4_qad.py b/fastvideo/layers/quantization/nvfp4_qad.py new file mode 100644 index 000000000..01f13ad26 --- /dev/null +++ b/fastvideo/layers/quantization/nvfp4_qad.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 quantization-aware distillation (QAD) for the MiniMax-H3 DiT block linears. + +``NVFP4QADLinearMethod`` replaces a ``ReplicatedLinear``'s ``quant_method`` +in place, so parameter names, FSDP sharding and checkpoints are unchanged. +Its forward *is* the deployed NVFP4 GEMM, computed from the layer's current +(BF16-cast) master weight exactly as the packed export and +``NVFP4QuantizeMethod.apply`` compute it: + +* weight: RTN NVFP4 with global scale ``448 * 6 / amax(W)`` (the converter's + ``quantize_dense_linear``; ModelOpt max calibration gives the same codes); +* activation: per-16 E4M3 block scales under a global scale that is + ``static`` (``448 * 6 / calibrated amax``, the FFN linears) or ``unit`` + (1.0, attention projections and VSA gates); +* GEMM: FlashInfer ``mm_fp4`` with ``alpha = (1 / weight_sf) / input_sf``. + +With gradients enabled the forward runs inside a straight-through estimator +whose backward is full precision. ``enabled = +False`` delegates to the original method, so a disabled student is the +teacher bit for bit. +""" +from __future__ import annotations + +import contextlib +import json +import re +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +FP4_RANGE = 448.0 * 6.0 +H3_BLOCK_LINEAR = re.compile(r"transformer_blocks\.(\d+)\.(attn\.to_q|attn\.to_k|attn\.to_v|attn\.to_out" + r"|attn\.to_gate_compress|ff\.fc_in|ff\.fc_out)$") +STATIC_SUFFIXES = ("ff.fc_in", "ff.fc_out") + + +def weight_global_sf(weight: torch.Tensor) -> torch.Tensor: + """``448 * 6 / amax(|W|)`` in fp32, as the packed-export converter computes it. + + The max of a BF16 tensor is exact in BF16, so one BF16 reduction gives the converter's value + (``weight.float().abs().nan_to_num().max()``) for finite weights without fp32 temporaries. + """ + return FP4_RANGE / torch.nan_to_num(weight.abs().amax().float()).clamp(min=1e-12) + + +def nvfp4_gemm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, + input_global_sf: torch.Tensor) -> torch.Tensor: + """The deployed NVFP4 linear on a BF16 weight: quantize both operands, FP4 GEMM.""" + from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend, _nvfp4_quantize + from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer + + sf_layout, _, _ = _require_flashinfer() + weight = weight.to(torch.bfloat16).contiguous() + w_sf = weight_global_sf(weight) + w_fp4, w_scale = _nvfp4_quantize(weight, w_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + alpha = (1.0 / w_sf) / input_global_sf + shape = x.shape + x2d = x.to(torch.bfloat16).reshape(-1, shape[-1]).contiguous() + x_fp4, x_scale = _nvfp4_quantize(x2d, input_global_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + out = _mm_fp4(x_fp4, w_fp4.T, x_scale, w_scale.T, alpha, torch.bfloat16, None, backend=_mm_fp4_backend()) + if bias is not None: + out = out + bias.to(out.dtype) + return out.view(*shape[:-1], weight.shape[0]) + + +class _NVFP4STE(torch.autograd.Function): + """Forward = ``nvfp4_gemm``; backward = full-precision linear gradients (straight-through).""" + + @staticmethod + def forward(ctx, x, weight, bias, input_global_sf): # type: ignore[override] + ctx.save_for_backward(x, weight) + ctx.has_bias = bias is not None + return nvfp4_gemm(x, weight, bias, input_global_sf) + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + x, weight = ctx.saved_tensors + grad_2d = grad_out.reshape(-1, grad_out.shape[-1]) + grad_x = (grad_2d @ weight.to(grad_2d.dtype)).reshape(x.shape).to(x.dtype) + grad_w = (grad_2d.t() @ x.reshape(-1, x.shape[-1]).to(grad_2d.dtype)).to(weight.dtype) + grad_b = grad_2d.sum(dim=0) if ctx.has_bias else None + return grad_x, grad_w, grad_b, None + + +class NVFP4QADLinearMethod: + """Drop-in ``quant_method`` for one H3 block linear (see module docstring).""" + + def __init__(self, + original: Any, + prefix: str, + act_scale: str, + input_amax: float | None = None, + allow_uncalibrated: bool = False) -> None: + if act_scale not in ("static", "unit"): + raise ValueError(f"act_scale must be 'static' or 'unit', got {act_scale!r}") + if act_scale == "static" and not allow_uncalibrated and (input_amax is None or not input_amax > 0): + raise ValueError(f"{prefix}: a static activation scale needs a positive calibrated amax") + self.original = original + self.prefix = prefix + self.act_scale = act_scale + self.input_amax = float(input_amax) if input_amax is not None else None + self.enabled = True + # Calibration: run the original (dense) path and fold each input's abs-max into ``observed_amax``. + self.calibrating = False + self.observed_amax: torch.Tensor | None = None + self._sf: dict[torch.device, torch.Tensor] = {} + + def set_input_amax(self, amax: float) -> None: + if not amax > 0: + raise ValueError(f"{self.prefix}: calibrated amax must be positive, got {amax}") + self.input_amax = float(amax) + self._sf.clear() + + def input_global_sf(self, device: torch.device) -> torch.Tensor: + if self.act_scale == "static" and self.input_amax is None: + raise RuntimeError(f"{self.prefix}: static activation scale used before calibration") + if device not in self._sf: + value = 1.0 if self.act_scale == "unit" else FP4_RANGE / max(self.input_amax or 0.0, 1e-12) + self._sf[device] = torch.tensor(value, dtype=torch.float32, device=device) + return self._sf[device] + + def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + if self.calibrating: + amax = x.detach().abs().amax().float() + self.observed_amax = amax if self.observed_amax is None else torch.maximum(self.observed_amax, amax) + return self.original.apply(layer, x, bias) + if not self.enabled: + return self.original.apply(layer, x, bias) + input_sf = self.input_global_sf(x.device) + if torch.is_grad_enabled() and (layer.weight.requires_grad or x.requires_grad): + return _NVFP4STE.apply(x, layer.weight, bias, input_sf) + return nvfp4_gemm(x, layer.weight, bias, input_sf) + + def __getattr__(self, name: str) -> Any: + # Callers probe optional quant-method hooks (e.g. wants_prequantized_input); fall back to the original. + # Before __init__ (copy.deepcopy, unpickling) there is no original: AttributeError keeps hasattr working. + try: + original = self.__dict__["original"] + except KeyError: + raise AttributeError(name) from None + return getattr(original, name) + + +@dataclass(frozen=True) +class NVFP4QADPlan: + """Which block linears are quantized and how (``h3_dit_vsa``: attention, gates and FFN).""" + quantize_attention: bool = True + quantize_gate: bool = True + quantize_ffn: bool = True + skip_blocks: tuple[int, ...] = () + + +def load_amax_table(path: str | Path) -> dict[str, float]: + """``FASTVIDEO_NVFP4_ACT_AMAX``-format JSON (``b.`` or full prefix -> amax or {"all": amax}).""" + raw = json.loads(Path(path).read_text()) + return {key: float(value["all"] if isinstance(value, dict) else value) for key, value in raw.items()} + + +def _amax_for(table: dict[str, float], name: str, block: str, sub: str) -> float | None: + for key in (name, f"b{block}.{sub}", f"transformer_blocks.{block}.{sub}"): + if key in table: + return table[key] + return None + + +def install_nvfp4_qad(transformer: torch.nn.Module, + amax_table: dict[str, float] | None, + plan: NVFP4QADPlan, + allow_uncalibrated: bool = False) -> dict[str, NVFP4QADLinearMethod]: + """Install QAD methods on the selected block linears; returns them by module name. + + FFN linears take ``static`` scales from ``amax_table`` (required unless + ``allow_uncalibrated``, which leaves them for ``calibrate``); attention + projections and gates use the unit scale, as the deployed export does. + """ + installed: dict[str, NVFP4QADLinearMethod] = {} + for raw_name, module in transformer.named_modules(): + # Activation checkpointing wraps blocks; key everything by the checkpoint (unwrapped) name. + name = raw_name.replace("_checkpoint_wrapped_module.", "") + match = H3_BLOCK_LINEAR.search(name) + if match is None or not hasattr(module, "quant_method") or int(match.group(1)) in plan.skip_blocks: + continue + block, sub = match.group(1), match.group(2) + if sub.startswith("ff.") and not plan.quantize_ffn: + continue + if sub == "attn.to_gate_compress" and not plan.quantize_gate: + continue + if sub.startswith("attn.to_") and sub != "attn.to_gate_compress" and not plan.quantize_attention: + continue + original = module.quant_method + if isinstance(original, NVFP4QADLinearMethod): + original = original.original + if sub in STATIC_SUFFIXES: + amax = _amax_for(amax_table or {}, name, block, sub) + if amax is None and not allow_uncalibrated: + raise KeyError(f"no calibrated input amax for {name} (key b{block}.{sub}); run Stage A calibration") + method = NVFP4QADLinearMethod(original, name, "static", amax, allow_uncalibrated=allow_uncalibrated) + else: + method = NVFP4QADLinearMethod(original, name, "unit") + module.quant_method = method + installed[name] = method + if not installed: + raise ValueError("no MiniMax-H3 block linears matched; is this a MiniMax-H3 transformer?") + return installed + + +@contextlib.contextmanager +def calibrating(methods: dict[str, NVFP4QADLinearMethod]) -> Iterator[None]: + """Max-calibrate the static linears that have no amax yet (dense forwards while active). + + On exit the observed maxima are all-reduced (MAX) over the default process + group, so every rank, including each sequence-parallel shard, agrees. + """ + pending = {name: m for name, m in methods.items() if m.act_scale == "static" and m.input_amax is None} + for method in pending.values(): + method.calibrating, method.observed_amax = True, None + try: + yield + finally: + for method in pending.values(): + method.calibrating = False + if not pending: + return + names = sorted(pending) + device = next((m.observed_amax.device for m in pending.values() if m.observed_amax is not None), + torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else None) + values = torch.stack([(pending[n].observed_amax if pending[n].observed_amax is not None else torch.zeros( + (), device=device)).float() for n in names]) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.MAX) + for name, value in zip(names, values.tolist(), strict=True): + if not value > 0: + raise RuntimeError(f"{name} saw no activations during calibration") + pending[name].set_input_amax(value) + + +def set_nvfp4_qad_enabled(methods: dict[str, NVFP4QADLinearMethod], enabled: bool) -> None: + for method in methods.values(): + method.enabled = enabled + + +def nvfp4_qad_export_scales(methods: dict[str, NVFP4QADLinearMethod]) -> dict[str, float]: + """The static FFN amax values the export must carry (``--act-amax`` input of the packed converter).""" + table = {} + for name, method in methods.items(): + match = H3_BLOCK_LINEAR.search(name) + if method.act_scale == "static" and match is not None: + table[f"b{match.group(1)}.{match.group(2)}"] = float(method.input_amax or 0.0) + return table + + +__all__ = [ + "NVFP4QADLinearMethod", "NVFP4QADPlan", "calibrating", "install_nvfp4_qad", "load_amax_table", "nvfp4_gemm", + "nvfp4_qad_export_scales", "set_nvfp4_qad_enabled", "weight_global_sf" +] diff --git a/fastvideo/models/dits/minimax_h3_token_compression.py b/fastvideo/models/dits/minimax_h3_token_compression.py new file mode 100644 index 000000000..f0b4d5439 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_token_compression.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Middle-block spatial token compression for FastH3 (QAD training and inference; hook based). + +Blocks ``[start, end)`` run on fewer video tokens. Text and audio rows are untouched; a kept video token keeps its +original RoPE coordinates and AdaLN row; after the last compressed block every video token receives the update of +its nearest kept token: ``H_out = H + expand(F(P H) - P H)`` where ``P`` is the compression operator. With +compression off nothing changes (no hook acts), so the student equals the teacher bit for bit. + +Rather than replacing the block list, this installs forward hooks, so module names, FSDP sharding, activation +checkpointing and QAD linears are untouched: + +* block ``start`` pre-hook compresses the hidden states and swaps AdaLN rows / RoPE / sequence length; +* every block in the range gets the compressed AdaLN rows / RoPE; +* the inner attention module of each compressed block swaps the VSA metadata for the compressed video grid (a hook + on the *inner* module, so activation-checkpoint recompute during backward swaps it too); +* block ``end - 1`` post-hook expands and restores. + +Single sequence-parallel rank only (the compression indexes the whole packed sequence). +""" +from __future__ import annotations + +import contextlib +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn + +from fastvideo.forward_context import get_forward_context + +MODES = ("subsample", "pair_x", "pool2x2") + + +@dataclass(frozen=True) +class CompressionSpec: + start: int + end: int + mode: str = "subsample" + # Compress only the first ``first_steps`` of the sampling forwards (None = all); the low-noise steps that set + # fine detail then run on the full token grid. + first_steps: int | None = None + + def __post_init__(self) -> None: + if self.mode not in MODES: + raise ValueError(f"compression mode must be one of {MODES}, got {self.mode!r}") + if not 0 <= self.start < self.end: + raise ValueError(f"need 0 <= start < end, got [{self.start}, {self.end})") + + @classmethod + def parse(cls, text: str) -> CompressionSpec | None: + """``"off"``/empty -> None; ``"mode:start:end"`` (e.g. ``subsample:3:41``).""" + text = text.strip().lower() + if text in ("", "0", "off", "none"): + return None + parts = text.split(":") + if len(parts) not in (3, 4): + raise ValueError(f"token compression must be 'off' or 'mode:start:end[:first_steps]', got {text!r}") + return cls(int(parts[1]), int(parts[2]), parts[0], int(parts[3]) if len(parts) == 4 else None) + + @classmethod + def from_config(cls, cfg: dict[str, Any] | None) -> CompressionSpec | None: + if not cfg or not cfg.get("enabled", True): + return None + first = cfg.get("first_steps") + return cls(int(cfg["start"]), int(cfg["end"]), str(cfg.get("mode", "subsample")), + None if first is None else int(first)) + + +def compressed_grid(mode: str, grid: tuple[int, int, int]) -> tuple[int, int, int]: + frames, height, width = grid + if mode == "pair_x": + return frames, height, (width + 1) // 2 + return frames, (height + 1) // 2, (width + 1) // 2 + + +def token_reduction(mode: str, grid: tuple[int, int, int]) -> float: + sub = compressed_grid(mode, grid) + return (grid[0] * grid[1] * grid[2]) / (sub[0] * sub[1] * sub[2]) + + +def _pad_to_even(x: torch.Tensor, dims: tuple[int, ...]) -> torch.Tensor: + for dim in dims: + if x.shape[dim] % 2: + x = torch.cat((x, x.narrow(dim, x.shape[dim] - 1, 1)), dim=dim) + return x + + +def grid_ops(mode: str, grid: tuple[int, int, int]) -> tuple[Callable, Callable, torch.Tensor]: + """``(compress, expand, keep_index)`` for a video token grid ``(T, H, W)``. + + ``compress``: rows ``[B, T*H*W, C]`` -> ``[B, T*h*w, C]``; ``expand``: delta ``[B, T*h*w, C]`` -> ``[B, T*H*W, C]``; + ``keep_index``: for every compressed token the video row whose RoPE / AdaLN row it keeps (CPU long). + """ + frames, height, width = grid + sub = compressed_grid(mode, grid) + index = torch.arange(frames * height * width).reshape(frames, height, width) + + def as_grid(rows: torch.Tensor) -> torch.Tensor: + return rows.reshape(rows.shape[0], frames, height, width, rows.shape[-1]) + + if mode == "subsample": + keep = index[:, ::2, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + return as_grid(rows)[:, :, ::2, ::2].reshape(rows.shape[0], -1, rows.shape[-1]) + + elif mode == "pair_x": + keep = index[:, :, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + g = _pad_to_even(as_grid(rows), (3, )) + pair = g.float().unflatten(3, (-1, 2)).mean(dim=4) + return pair.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1]) + + else: # pool2x2 + keep = index[:, ::2, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + g = _pad_to_even(as_grid(rows), (2, 3)) + pooled = g.float().unflatten(3, (-1, 2)).unflatten(2, (-1, 2)).mean(dim=(3, 5)) + return pooled.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1]) + + repeat_h = 1 if mode == "pair_x" else 2 + + def expand(delta: torch.Tensor) -> torch.Tensor: + d = delta.reshape(delta.shape[0], *sub, delta.shape[-1]) + if repeat_h == 2: + d = d.repeat_interleave(2, dim=2)[:, :, :height] + d = d.repeat_interleave(2, dim=3)[:, :, :, :width] + return d.reshape(delta.shape[0], -1, delta.shape[-1]) + + return compress, expand, keep + + +def _unwrap(module: nn.Module) -> nn.Module: + return getattr(module, "_checkpoint_wrapped_module", module) + + +@dataclass +class _Forward: + """Per-forward state (one transformer call).""" + full: torch.Tensor | None = None + compressed_in: torch.Tensor | None = None + adaln: torch.Tensor | None = None + rotary: tuple[torch.Tensor, torch.Tensor] | None = None + seq_len: int = 0 + num_prefix: int = 0 + + +class TokenCompression: + """Installs the hooks; ``scope(grid, compressed_metadata)`` arms them for one forward (and its backward).""" + + def __init__(self, transformer: nn.Module, spec: CompressionSpec) -> None: + self.spec = spec + self.blocks = list(transformer.transformer_blocks) + if spec.end > len(self.blocks): + raise ValueError(f"compression range [{spec.start}, {spec.end}) exceeds {len(self.blocks)} blocks") + self._handles: list[Any] = [] + self._grid: tuple[int, int, int] | None = None + self._ops: tuple[Callable, Callable, torch.Tensor] | None = None + self._metadata: Any = None + self._saved_metadata: Any = None + self._state = _Forward() + for index in range(spec.start, spec.end): + block = self.blocks[index] + self._handles.append(block.register_forward_pre_hook(self._pre(index))) + if index == spec.end - 1: + self._handles.append(block.register_forward_hook(self._post)) + attn = _unwrap(block).attn + self._handles.append(attn.register_forward_pre_hook(self._attn_pre)) + self._handles.append(attn.register_forward_hook(self._attn_post)) + + @property + def active(self) -> bool: + return self._grid is not None + + @contextlib.contextmanager + def scope(self, grid: tuple[int, int, int], compressed_metadata: Any) -> Iterator[None]: + """Arm the hooks for forwards (and their backward recompute) inside the block.""" + self._grid, self._metadata = grid, compressed_metadata + self._ops = grid_ops(self.spec.mode, grid) + try: + yield + finally: + self._grid = self._ops = self._metadata = None + self._state = _Forward() + + def remove(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + # -- hooks + def _pre(self, index: int): + + def hook(module: nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...] | None: + if not self.active: + return None + hidden, temb, adaln, rotary, _ = args + if index == self.spec.start: + self._enter(hidden, adaln, rotary) + hidden = self._state.compressed_in + state = self._state + return hidden, temb, state.adaln, state.rotary, state.seq_len + + return hook + + def _enter(self, hidden: torch.Tensor, adaln: torch.Tensor, rotary: tuple[torch.Tensor, torch.Tensor]) -> None: + assert self._grid is not None and self._ops is not None + compress, _, keep = self._ops + num_video = self._grid[0] * self._grid[1] * self._grid[2] + num_prefix = hidden.shape[1] - num_video + if num_prefix < 0 or adaln.shape[0] != hidden.shape[1]: + raise ValueError("token compression needs the whole packed sequence on one rank " + f"(hidden {tuple(hidden.shape)}, adaln {tuple(adaln.shape)}, video {num_video})") + rows = torch.cat((torch.arange(num_prefix), keep + num_prefix)).to(hidden.device) + self._state = _Forward(full=hidden, + compressed_in=torch.cat((hidden[:, :num_prefix], compress(hidden[:, num_prefix:])), dim=1), + adaln=adaln.index_select(0, rows), + rotary=(rotary[0].index_select(0, rows), rotary[1].index_select(0, rows)), + seq_len=int(rows.numel()), + num_prefix=num_prefix) + + def _post(self, module: nn.Module, args: tuple[Any, ...], output: torch.Tensor) -> torch.Tensor | None: + if not self.active: + return None + assert self._ops is not None + state = self._state + n = state.num_prefix + delta = self._ops[1](output[:, n:] - state.compressed_in[:, n:]) + restored = torch.cat((output[:, :n], state.full[:, n:] + delta.to(state.full.dtype)), dim=1) + self._state = _Forward() + return restored + + def _attn_pre(self, module: nn.Module, args: tuple[Any, ...], kwargs: Any = None) -> None: + if not self.active: + return + context = get_forward_context() + self._saved_metadata = context.attn_metadata + context.attn_metadata = self._metadata + + def _attn_post(self, module: nn.Module, args: tuple[Any, ...], output: Any) -> None: + if not self.active: + return + get_forward_context().attn_metadata = self._saved_metadata + + +def block_flops_share(spec: CompressionSpec | None, grid: tuple[int, int, int], num_blocks: int, num_other_tokens: int, + attn_keep: float = 0.2) -> dict[str, float]: + """Relative DiT cost of the compressed model (1.0 = uncompressed), by FLOPs. + + Per block, per token: linear layers ~ ``12 d^2`` MACs (qkv+out+gate ~ 5 d^2, FFN SwiGLU 3 * d * 2.67 d), attention + ~ ``2 * attended_keys * d``; VSA keeps ``attn_keep`` of the video keys. ``d`` cancels in the ratio when expressed in + ``d`` units, so only the token counts matter; ``d = 5376`` is used for the attention/linear balance. + """ + d = 5376.0 + video = grid[0] * grid[1] * grid[2] + + def cost(tokens_video: float) -> float: + n = tokens_video + num_other_tokens + linear = 13.4 * d * d * n + attn = 2.0 * 2.0 * d * n * (attn_keep * tokens_video + num_other_tokens) + return linear + attn + + full = cost(video) + if spec is None: + return {"relative_flops": 1.0, "speedup_flops": 1.0} + reduced = cost(video / token_reduction(spec.mode, grid)) + inside = spec.end - spec.start + total = (num_blocks - inside) * full + inside * reduced + return {"relative_flops": total / (num_blocks * full), "speedup_flops": num_blocks * full / total} + + +__all__ = ["CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops", "token_reduction"] diff --git a/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py new file mode 100644 index 000000000..02e133046 --- /dev/null +++ b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py @@ -0,0 +1,614 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniRef PDD QAD: student initialization, deployed NVFP4 numerics, rows and metrics. + +CPU tests: a disabled QAD student is bit-identical to the untouched model, +static scales come from the calibration table, the row plan is the latent +generator's plan, and the eval statistics are consistent. GPU tests (GB200 or +any FlashInfer FP4 GPU): the QAD linear forward is bit-identical to the +deployed ``NVFP4QuantizeMethod`` on a converter-packed weight, and the STE +backward is the dense gradient. +""" +from __future__ import annotations + +import copy +import math + +import pytest +import torch + +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config +from fastvideo.forward_context import set_forward_context +from fastvideo.layers.quantization.nvfp4_qad import (NVFP4QADLinearMethod, NVFP4QADPlan, install_nvfp4_qad, + nvfp4_qad_export_scales, set_nvfp4_qad_enabled) +from fastvideo.models.dits import minimax_h3 +from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_packed_sequence, build_row_timesteps +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.tests.loader.test_minimax_h3_pdd_loading import TINY_ARCH, VIDEO_PATCH_DIM, cpu_loader # noqa: F401 +from fastvideo.tests.loader.test_minimax_h3_pdd_loading import single_process_group # noqa: F401 +from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import RungStats, summarize +from fastvideo.train.models.minimax_h3.omniref_data import RowSpec, RowStream, select_rows + +needs_fp4 = pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="NVFP4 GEMMs need a Blackwell GPU with FlashInfer") + + +def _tiny_model() -> torch.nn.Module: + arch = MiniMaxH3ArchConfig(**{**TINY_ARCH, "patch_size": tuple(TINY_ARCH["patch_size"])}) + model = minimax_h3.MiniMaxH3Transformer3DModel(MiniMaxH3Config(arch_config=arch), hf_config={}) + generator = torch.Generator().manual_seed(0) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1) + return model.eval() + + +def _amax_table(num_blocks: int = 2) -> dict[str, float]: + return {f"b{b}.ff.{sub}": 3.0 + b for b in range(num_blocks) for sub in ("fc_in", "fc_out")} + + +def _forward(model: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor]: + layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2)) + unique, inverse = build_row_timesteps(layout, 0.3, 0.4, 0.999, 1.0) + generator = torch.Generator().manual_seed(1) + inputs = dict(hidden_states=torch.randn(1, int(layout.video_indices.numel()), VIDEO_PATCH_DIM, generator=generator), + audio_hidden_states=torch.randn(1, int(layout.audio_indices.numel()), TINY_ARCH["audio_in_channels"], + generator=generator), + encoder_hidden_states=torch.randn(1, 3, TINY_ARCH["text_dim"], generator=generator), + timestep=unique, timestep_indices=inverse, token_tags=layout.token_tags, + position_ids=layout.position_ids, video_indices=layout.video_indices, + audio_indices=layout.audio_indices, text_indices=layout.text_indices) + with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None, + forward_batch=ForwardBatch(data_type="dummy")): + return model(**inputs) + + +def test_disabled_qad_student_is_bit_identical_to_teacher(cpu_loader, single_process_group): # noqa: F811 + teacher = _tiny_model() + student = copy.deepcopy(teacher) + methods = install_nvfp4_qad(student, _amax_table(), NVFP4QADPlan()) + # Two blocks x (q, k, v, out, fc_in, fc_out); SDPA builds no VSA gates. + assert len(methods) == 12 + assert all(isinstance(m, NVFP4QADLinearMethod) for m in methods.values()) + set_nvfp4_qad_enabled(methods, False) + for expected, got in zip(_forward(teacher), _forward(student), strict=True): + assert torch.equal(expected, got) + assert all(torch.equal(a, b) for a, b in zip(teacher.state_dict().values(), student.state_dict().values())) + + +def test_static_scales_require_calibration_and_round_trip(cpu_loader, single_process_group): # noqa: F811 + model = _tiny_model() + table = _amax_table() + with pytest.raises(KeyError, match="calibrated input amax"): + install_nvfp4_qad(copy.deepcopy(model), {k: v for k, v in table.items() if k != "b1.ff.fc_out"}, + NVFP4QADPlan()) + methods = install_nvfp4_qad(model, table, NVFP4QADPlan()) + assert nvfp4_qad_export_scales(methods) == table + unit = [name for name, m in methods.items() if m.act_scale == "unit"] + assert len(unit) == 8 and all(".attn." in name for name in unit) + assert float(methods["transformer_blocks.1.ff.fc_out"].input_global_sf(torch.device("cpu"))) == pytest.approx( + 2688.0 / 4.0) + # Re-installing wraps the original method, never a QAD method. + again = install_nvfp4_qad(model, table, NVFP4QADPlan(quantize_attention=False)) + assert len(again) == 4 and not isinstance(again["transformer_blocks.0.ff.fc_in"].original, NVFP4QADLinearMethod) + + +def _reference_plan(groups: dict[tuple[str, str], list[dict]], seed: int, per_group: int) -> list[dict]: + """The latent generator's selection: per-group seeded shuffle, then interleave.""" + import random + picked = [] + for key in sorted(groups): + entries = sorted(groups[key], key=lambda entry: entry["id"]) + random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries) + picked.append(entries[:per_group]) + plan = [] + for index in range(max(len(entries) for entries in picked)): + for entries in picked: + if index < len(entries): + entry = entries[index] + source = entry["id"].split(":", 1)[-1] + plan.append({ + **entry, + "id": f"omniref-{entry['case']}-{source}-{entry['resolution']}", + "source": source, + }) + return plan + + +def test_qad_student_can_be_deep_copied_and_pickled(cpu_loader, single_process_group): # noqa: F811 + """EMA copies and teacher snapshots deep-copy the student after install_nvfp4_qad.""" + import pickle + student = _tiny_model() + install_nvfp4_qad(student, _amax_table(), NVFP4QADPlan()) + clone = copy.deepcopy(student) + method = clone.transformer_blocks[0].ff.fc_in.quant_method + assert isinstance(method, NVFP4QADLinearMethod) and method.input_amax == 3.0 + assert method.original is not student.transformer_blocks[0].ff.fc_in.quant_method.original + restored = pickle.loads(pickle.dumps(student.transformer_blocks[0].ff.fc_out.quant_method)) + assert restored.prefix.endswith("transformer_blocks.0.ff.fc_out") and restored.input_amax == 3.0 + with pytest.raises(AttributeError): + object.__new__(NVFP4QADLinearMethod).anything + + +def test_row_plan_is_the_latent_generators_plan(): + groups = {(case, res): [{"id": f"{case}:src{i}", "case": case, "resolution": res, "parquet": f"/p/{case}{i}"} + for i in range(12)] for case in ("first_frame", "storyboard") for res in ("480p", "768p")} + ours = select_rows(groups, 20261007, {"480p": 5, "768p": 5}) + theirs = _reference_plan(groups, 20261007, 5) + assert [(r.id, r.parquet) for r in ours] == [(c["id"], c["parquet"]) for c in theirs] + excluded = select_rows(groups, 20261007, {"480p": 12}, frozenset({"src0", "src1"})) + assert not any(r.source in {"src0", "src1"} for r in excluded) and all(r.resolution == "480p" for r in excluded) + + +def test_row_stream_feeds_each_data_parallel_group_distinct_rows(): + plan = [RowSpec(f"r{i}", f"s{i}", "first_frame", "480p", f"/p{i}", i) for i in range(10)] + streams = [iter(RowStream(plan, dp_rank=r, dp_size=2, seed=3)) for r in range(2)] + firsts = [[next(s)["row"].id for _ in range(5)] for s in streams] + assert set(firsts[0]).isdisjoint(firsts[1]) and len(set(firsts[0] + firsts[1])) == 10 + replay = iter(RowStream(plan, dp_rank=0, dp_size=2, seed=3)) + assert [next(replay)["row"].id for _ in range(5)] == firsts[0] + + +def test_rung_stats_summaries(): + stats = RungStats(2) + teacher = torch.ones(4, 3) + for rung in range(2): + pair = {m: {"student": teacher * (1.1 if m == "video" else 1.0), "teacher": teacher, + "student_v": teacher, "teacher_v": teacher} for m in ("video", "audio")} + stats.add_rung(rung, pair) + summary = summarize(stats.reduce(contribute=True), 2, {"video": 1.0, "audio": 1.0}) + assert summary["x0_rel_l2/video/rung1"] == pytest.approx(0.1) + assert summary["x0_rel_l2/audio/rung0"] == 0.0 + assert summary["score"] == pytest.approx(0.05) + + +# ----------------------------------------------------------------------------- GPU: deployed numerics +def _runtime_layer(weight: torch.Tensor, input_amax: float | None): + """A ``ReplicatedLinear`` packed the way the H3 NVFP4 converter packs a dense weight.""" + from fastvideo.layers.linear import ReplicatedLinear + from fastvideo.layers.quantization.nvfp4_config import (H3_NVFP4_DIT_INPUT_SF_NAME, NVFP4QuantizeMethod, + _nvfp4_quantize, _require_flashinfer) + sf_layout, _, _ = _require_flashinfer() + layer = ReplicatedLinear(weight.shape[1], weight.shape[0], bias=False).cuda() + global_sf = (448 * 6) / weight.float().abs().nan_to_num().max() + packed, scale = _nvfp4_quantize(weight, global_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + layer._nvfp4_weight, layer._nvfp4_weight_scale = packed, scale + layer._nvfp4_alpha = (1.0 / torch.as_tensor(global_sf, device="cuda", dtype=torch.float32)).float() + layer._weight_global_sf = torch.as_tensor(global_sf, device="cuda").to(torch.bfloat16) + if input_amax is not None: + setattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, torch.tensor(2688.0 / input_amax, device="cuda")) + return layer, NVFP4QuantizeMethod("transformer_blocks.0.ff.fc_out") + + +@needs_fp4 +@pytest.mark.parametrize("input_amax", [None, 40.0]) +def test_qad_linear_forward_is_the_deployed_nvfp4_linear(input_amax): + from fastvideo.layers.quantization.nvfp4_qad import _NVFP4STE, nvfp4_gemm + torch.manual_seed(0) + weight = (torch.randn(512, 1024, device="cuda") * 0.05).to(torch.bfloat16) + x = (torch.randn(300, 1024, device="cuda") * 8).to(torch.bfloat16) + layer, runtime = _runtime_layer(weight, input_amax) + expected = runtime.apply(layer, x) + sf = torch.tensor(1.0 if input_amax is None else 2688.0 / input_amax, device="cuda") + assert torch.equal(nvfp4_gemm(x, weight, None, sf), expected) + weight_master = weight.float().requires_grad_(True) + x_grad = x.clone().requires_grad_(True) + out = _NVFP4STE.apply(x_grad, weight_master, None, sf) + assert torch.equal(out, expected) + grad = torch.randn_like(out) + out.backward(grad) + torch.testing.assert_close(weight_master.grad, (grad.float().t() @ x.float()), rtol=2e-2, atol=1e-2) + torch.testing.assert_close(x_grad.grad.float(), grad.float() @ weight.float(), rtol=2e-2, atol=1e-2) + + +def test_block_reconstruction_is_local_and_zero_at_identity(cpu_loader, single_process_group): # noqa: F811 + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction + teacher = _tiny_model() + student = copy.deepcopy(teacher).train().requires_grad_(True) + teacher.requires_grad_(False) + with BlockReconstruction(teacher, student, scale=1.0) as recon: + _forward(teacher) + assert len(recon.attn_rel) == 2 and float(recon.loss()) == 0.0 + assert all(p.grad is None or float(p.grad.abs().max()) == 0.0 for p in student.parameters()) + # Perturb block 1's FFN only: its FFN unit errs; block 0 and every attention unit stay exact. + with torch.no_grad(): + student.transformer_blocks[1].ff.fc_out.weight.mul_(1.1) + student.zero_grad(set_to_none=True) + with BlockReconstruction(teacher, student, scale=1.0) as recon: + _forward(teacher) + assert float(recon.ff_rel[1]) > 0 and float(recon.ff_rel[0]) == 0.0 + assert all(float(a) == 0.0 for a in recon.attn_rel) + grads = {n: p.grad for n, p in student.named_parameters() if p.grad is not None and p.grad.abs().max() > 0} + assert grads and all(n.startswith("transformer_blocks.1.ff.") for n in grads), sorted(grads) + + +# ----------------------------------------------------------------------------- E0: recipe-review fixes +def test_x0_matches_scheduler_sign(): + """x0 = x_t + (1 - t) v, exactly what MiniMaxH3Scheduler.step denoises to (a final step lands on it).""" + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import denoised + scheduler = MiniMaxH3Scheduler(shift=12.0) + scheduler.set_timesteps(sigmas=torch.tensor([0.7, 0.0])) + x_t, v = torch.randn(5, 3), torch.randn(5, 3) + final = scheduler.step(v, scheduler.timesteps[0], x_t, return_dict=False)[0] + torch.testing.assert_close(denoised(x_t, v, scheduler.timesteps[0]), final) + assert not torch.allclose(final, x_t - 0.7 * v) + + +def test_grouped_rel_error_weights_groups_and_drops_absent(): + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import GroupedRelError + groups = torch.tensor([0, 0, 1, 2, 2, -1]) # video, video, audio, cond, cond, padding + teacher = torch.tensor([[1.0, 0], [1, 0], [0, 2], [0, 100], [100, 0], [7, 7]])[None] + student = teacher.clone() + student[0, 2] += torch.tensor([0.0, 1.0]) # audio rel = 1/4 + student[0, 5] += 9.0 # padding row never counts + loss, per_group = GroupedRelError(groups, {"video": 1.0, "audio": 1.0, "condition": 0.25})(student, teacher) + assert per_group.tolist() == [0.0, 0.25, 0.0] + assert float(loss) == pytest.approx(0.25 / 2.25) # the huge conditioning rows do not dilute audio + no_audio = torch.tensor([0, 0, 2, 2, 2, -1]) + _, per_group = GroupedRelError(no_audio, {"video": 1.0, "audio": 1.0})(student, teacher) + assert math.isnan(float(per_group[1])) + + +def test_grouped_rel_error_sp_gradient_matches_unsharded(monkeypatch): + """FSDP averages gradients over SP ranks; the mean of the shards' gradients must be the unsharded gradient. + + The x0 output loss already has this scale (its all-gather backward sums over SP ranks), so reconstruction + must too, or ``recon_lambda`` is silently divided by ``sp_size``. + """ + import fastvideo.distributed as distributed + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import GroupedRelError + torch.manual_seed(0) + weights = {"video": 1.0, "audio": 1.0, "condition": 0.25} + groups = torch.tensor([0, 0, 0, 1, 2, 2, 0, 1, 2, 0, 0, -1]) # 2 shards of 6 rows, last row is SP padding + x, teacher = torch.randn(12, 5), torch.randn(12, 3) + weight = torch.randn(3, 5, requires_grad=True) + + def grad(loss: torch.Tensor) -> torch.Tensor: + weight.grad = None + loss.backward() + return weight.grad.clone() + + monkeypatch.setattr(distributed, "model_parallel_is_initialized", lambda: False) + full_loss, full_rel = GroupedRelError(groups, weights)(x @ weight.T, teacher) + full_grad = grad(full_loss) + + sp_size, shard = 2, 6 + local_totals: list[torch.Tensor] = [] + + class _Group: # all_reduce = the sum of every shard's totals, recorded in a first pass + + def all_reduce(self, totals: torch.Tensor) -> torch.Tensor: + local_totals.append(totals) + return sum(local_totals[:sp_size]) if len(local_totals) > sp_size else totals + + monkeypatch.setattr(distributed, "model_parallel_is_initialized", lambda: True) + monkeypatch.setattr(distributed, "get_sp_world_size", lambda: sp_size) + monkeypatch.setattr(distributed, "get_sp_group", lambda: _Group()) + rows = [slice(rank * shard, (rank + 1) * shard) for rank in range(sp_size)] + for part in rows: # first pass: record each shard's local totals + GroupedRelError(groups[part], weights)(x[part] @ weight.T, teacher[part]) + shard_grads = [] + for part in rows: + loss, rel = GroupedRelError(groups[part], weights)(x[part] @ weight.T, teacher[part]) + torch.testing.assert_close(rel, full_rel) + shard_grads.append(grad(loss)) + torch.testing.assert_close(sum(shard_grads) / sp_size, full_grad) # FSDP's mean over the SP ranks + + +def test_packed_row_groups_from_layout(): + from types import SimpleNamespace + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import local_row_groups, packed_row_groups + layout = SimpleNamespace(video_indices=torch.tensor([3, 4, 5, 6]), num_condition_video_rows=1, + audio_indices=torch.tensor([7, 8]), num_condition_audio_rows=1) + groups = packed_row_groups(layout, 9) + assert groups.tolist() == [2, 2, 2, 2, 0, 0, 0, 2, 1] + assert torch.equal(local_row_groups(groups, torch.device("cpu")), groups) + + +def test_grouped_block_reconstruction_measure_only_has_no_grads(cpu_loader, single_process_group): # noqa: F811 + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction, GroupedRelError + teacher = _tiny_model() + student = copy.deepcopy(teacher).train().requires_grad_(True) + teacher.requires_grad_(False) + with torch.no_grad(): + student.transformer_blocks[0].attn.to_v.weight.mul_(1.2) + layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2)) + groups = torch.full((int(layout.position_ids.shape[0]), ), 2, dtype=torch.long) + groups[layout.video_indices] = 0 + groups[layout.audio_indices] = 1 + criterion = GroupedRelError(groups, {"video": 1.0, "audio": 1.0, "condition": 0.25}) + with BlockReconstruction(teacher, student, 1.0, criterion, measure_only=True) as recon: + _forward(teacher) + assert float(recon.loss()) > 0 and all(p.grad is None for p in student.parameters()) + with BlockReconstruction(teacher, student, 1.0, criterion) as recon: + _forward(teacher) + grads = {n for n, p in student.named_parameters() if p.grad is not None and p.grad.abs().max() > 0} + assert grads and all(n.startswith("transformer_blocks.0.attn.") for n in grads), sorted(grads) + assert set(recon.metrics("m")) >= {"m/video_rel", "m/audio_rel", "m/condition_rel"} + + +def _bare_method(**attrs): + from fastvideo.train.methods.knowledge_distillation.pdd_qad import OmniRefPDDQADMethod + method = object.__new__(OmniRefPDDQADMethod) + defaults = dict(eligibility_keys=("endpoint/audio", "keyframe/vs_teacher/first"), tolerances={}, + fidelity_tolerance=0.05, min_eligibility_tolerance=0.02, select_metric="score", + best_score=1.0, save_best=True, patience=3, bad_evals=0, bad_score_evals=0, + bad_alignment_evals=0, stop_requested=False, _checkpoint_request=None) + method.__dict__.update({**defaults, **attrs}) + return method + + +def test_qad_data_parallel_groups_follow_the_same_collective_schedule(): + """Different rows must not select different student forwards or reconstruction backwards under FSDP.""" + rows = [f"omniref-case{group}-asset{group}-480p" for group in range(4)] + methods = [_bare_method(dp_size=4, dp_rank=group, sampling_seed=1234, rungs_per_row=4, + rollout_start_step=20, rollout_fraction=0.25) for group in range(4)] + schedules = [] + for iteration in range(1, 201): + choices = [(method._trained_rungs(row, iteration, 8), method._use_rollout(row, iteration)) + for method, row in zip(methods, rows, strict=True)] + assert all(choice == choices[0] for choice in choices) + assert len(choices[0][0]) == 4 + if iteration < 20: + assert not choices[0][1] + schedules.append((tuple(sorted(choices[0][0])), choices[0][1])) + assert len(set(schedules)) > 8 + assert {rollout for _, rollout in schedules[19:]} == {False, True} + + +def test_qad_single_data_parallel_group_retains_row_sampling(): + from fastvideo.train.methods.knowledge_distillation.pdd_qad import _unit_hash + method = _bare_method(dp_size=1, sampling_seed=1234, rungs_per_row=4, + rollout_start_step=20, rollout_fraction=0.25) + for row in ("storyboard-a", "first_frame-b"): + for iteration in (1, 20, 31): + order = sorted(range(8), key=lambda rung: _unit_hash(row, iteration, "rung", rung)) + assert method._trained_rungs(row, iteration, 8) == set(order[:4]) + assert method._use_rollout(row, iteration) == (iteration >= 20 and _unit_hash(row, iteration) < 0.25) + + +def test_noise_floor_sets_two_sigma_tolerances(): + method = _bare_method() + step0 = {"endpoint/audio": 0.40, "keyframe/vs_teacher/first": 0.30, "score": 0.30, "endpoint/video": 0.42} + method._evaluate = lambda rows, rollout, save_tag="": {"endpoint/audio": 0.44, "keyframe/vs_teacher/first": 0.30, + "score": 0.30, "endpoint/video": 0.42} + out = method._noise_floor([], step0) + sigma = 0.04 / (math.sqrt(2) * 0.42) + assert out["init/noise_sigma/endpoint/audio"] == pytest.approx(sigma) + assert method.tolerances["endpoint/audio"] == pytest.approx(2 * sigma) + assert method.tolerances["keyframe/vs_teacher/first"] == 0.02 # floor when two seeds agree + + +def test_best_requires_eligibility_on_full_eval(): + method = _bare_method(step0={"score": 0.30, "endpoint/audio": 0.40, "keyframe/vs_teacher/first": 0.30}, + best_score=0.30, tolerances={"endpoint/audio": 0.10, "keyframe/vs_teacher/first": 0.10}) + teacher_forced_only = {"score": 0.20} + method._guardrails(25, teacher_forced_only) + assert method._checkpoint_request is None # better score, but alignment unknown + audio_regressed = {"score": 0.20, "endpoint/video": 0.4, "endpoint/audio": 0.45, "keyframe/vs_teacher/first": 0.3} + out = method._guardrails(50, audio_regressed) + assert method._checkpoint_request is None and out["eval/eligible"] == 0.0 and method.bad_evals == 1 + fine = {"score": 0.21, "endpoint/video": 0.4, "endpoint/audio": 0.43, "keyframe/vs_teacher/first": 0.31} + out = method._guardrails(75, fine) + assert method._checkpoint_request == "best" and method.best_score == 0.21 and method.bad_evals == 0 + + +def test_teacher_forced_evals_do_not_reset_alignment_failures(): + """eval_every 15 / full_eval_every 30 / patience 2: persistent alignment regressions must still stop the run.""" + method = _bare_method(step0={"score": 0.30, "endpoint/video": 0.40, "endpoint/audio": 0.40}, + best_score=0.30, patience=2) + full_regressed = {"score": 0.25, "endpoint/video": 0.40, "endpoint/audio": 0.60} + method._guardrails(30, full_regressed) + assert (method.bad_alignment_evals, method.bad_evals, method.stop_requested) == (1, 1, False) + method._guardrails(45, {"score": 0.24}) # teacher-forced only: better score, alignment unmeasured + assert (method.bad_score_evals, method.bad_alignment_evals, method.stop_requested) == (0, 1, False) + out = method._guardrails(60, full_regressed) + assert method.stop_requested and out["eval/bad_alignment_evals"] == 2.0 + method = _bare_method(step0={"score": 0.30, "endpoint/audio": 0.40}, patience=2) + method._guardrails(30, {"score": 0.25, "endpoint/video": 0.4, "endpoint/audio": 0.60}) + method._guardrails(60, {"score": 0.25, "endpoint/video": 0.4, "endpoint/audio": 0.40}) # recovered + assert method.bad_alignment_evals == 0 and not method.stop_requested + method._guardrails(75, {"score": 0.31}) + method._guardrails(90, {"score": 0.32, "endpoint/video": 0.4, "endpoint/audio": 0.40}) + assert method.bad_score_evals == 2 and method.stop_requested # the score count still runs on every eval + + +def test_baseline_state_round_trips_through_dcp(tmp_path, single_process_group): # noqa: F811 + """Resume restores the step-0 baseline, guard counters and provisional scales instead of re-measuring.""" + import torch.distributed.checkpoint as dcp + from types import SimpleNamespace + from fastvideo.train.methods.knowledge_distillation.pdd_qad import _BaselineState + + def student(amax): + linears = {"transformer_blocks.0.ff.fc_in": NVFP4QADLinearMethod(object(), "fc_in", "static", amax, + allow_uncalibrated=True)} + return SimpleNamespace(provisional_rows=2, qad_linears=linears) + + saved = _bare_method(normalizers={"video": [0.2, 0.1], "audio": [0.4, 0.3]}, step0={"score": 0.3}, + recon_norm=[0.05, 0.04], tolerances={"endpoint/audio": 0.1}, best_score=0.27, + bad_score_evals=1, bad_alignment_evals=2, student=student(7.5)) + dcp.save({"method.qad_baseline": _BaselineState(saved)}, checkpoint_id=str(tmp_path)) + resumed = _bare_method(normalizers=None, step0={}, recon_norm=None, best_score=math.inf, student=student(None)) + dcp.load({"method.qad_baseline": _BaselineState(resumed)}, checkpoint_id=str(tmp_path)) + for key in ("normalizers", "step0", "recon_norm", "tolerances", "best_score", "bad_score_evals", + "bad_alignment_evals"): + assert getattr(resumed, key) == getattr(saved, key), key + assert resumed.bad_evals == 2 + assert resumed.student.qad_linears["transformer_blocks.0.ff.fc_in"].input_amax == 7.5 + with pytest.raises(ValueError, match="no step-0 baseline"): + _bare_method(student=student(None))._load_baseline_state({**saved._baseline_state(), "normalizers": None}) + + +def test_provisional_calibration_runs_on_a_step0_cache_hit(single_process_group): # noqa: F811 + from types import SimpleNamespace + calls = [] + cached = {"eval/score": 0.3, "eval/x0_mse/video/rung0": 0.2, "eval/x0_mse/audio/rung0": 0.4} + method = _bare_method(normalizers=None, step0={}, eval_every=15, student=SimpleNamespace(num_rungs=1), + _rank_rows=lambda: [], _load_step0=lambda: cached, + _restore_step0=lambda c: {k[len("eval/"):]: v for k, v in c.items()}, + _provisional_calibration=lambda: calls.append(1) or {"init/provisional_calibration_rows": 2}) + metrics = method.on_validation_begin(0) + assert calls == [1] and metrics["init/provisional_calibration_rows"] == 2 + assert method.normalizers == {"video": [0.2], "audio": [0.4]} + method.on_validation_begin(7) # not an eval step (e.g. a resume): calibration still gets its chance + assert calls == [1, 1] + + +def test_manifest_groups_keep_each_resolution_of_a_shared_id(tmp_path): + import json + from fastvideo.train.models.minimax_h3.omniref_data import load_manifest_groups + + def entry(res, frames, case="first_frame", source="a"): + return {"id": f"{case}:{source}", "case": case, "resolution": res, + "parquet": f"/d/bucket=832x480-{frames}f/{source}-{res}.parquet"} + + first, second = tmp_path / "v4.jsonl", tmp_path / "v5.jsonl" + first.write_text("\n".join(json.dumps(e) for e in (entry("480p", 121), entry("768p", 121))) + "\n") + second.write_text("\n".join(json.dumps(e) for e in (entry("480p", 121), entry("480p", 241, source="b"))) + "\n") + groups = load_manifest_groups([str(first), str(second)], ("first_frame", ), 243) + assert {key: [e["id"] for e in entries] for key, entries in groups.items()} == { + ("first_frame", "480p"): ["first_frame:a", "first_frame:b"], # the v5 repeat of (a, 480p) is dropped + ("first_frame", "768p"): ["first_frame:a"], + } + + +def test_eval_top_up_respects_the_eval_frame_cap(): + from fastvideo.train.models.minimax_h3.omniref_data import limit_frames, top_up_eval_rows + groups = {("storyboard", "480p"): [{"id": f"storyboard:s{i}", "case": "storyboard", "resolution": "480p", + "parquet": f"/d/bucket=832x480-{121 if i % 2 else 241}f/s{i}.parquet"} + for i in range(8)]} + assert limit_frames(groups, 0) == groups + capped = limit_frames(groups, 121) + assert len(capped[("storyboard", "480p")]) == 4 and len(groups[("storyboard", "480p")]) == 8 + rows = top_up_eval_rows([], capped, 7, {"480p": 3}) + assert len(rows) == 3 and all(int(row.source[1:]) % 2 for row in rows) + + +def test_eval_top_up_is_held_out_and_interleaved(): + from fastvideo.train.models.minimax_h3.omniref_data import top_up_eval_rows + groups = {(case, res): [{"id": f"{case}:{res}{i}", "case": case, "resolution": res, "parquet": f"/p{i}"} + for i in range(10)] for case in ("first_frame", "storyboard") for res in ("480p", "768p")} + existing = [RowSpec("e1", "480p0", "first_frame", "480p", "/p", 1)] + rows = top_up_eval_rows(existing, groups, 7, {"480p": 3, "768p": 1}) + count = {} + for row in rows: + count[(row.case, row.resolution)] = count.get((row.case, row.resolution), 0) + 1 + assert count == {("first_frame", "480p"): 3, ("first_frame", "768p"): 1, ("storyboard", "480p"): 3, + ("storyboard", "768p"): 1} + assert len({row.source for row in rows if row.resolution == "480p" and row.case == "first_frame"}) == 3 + assert [(r.case, r.resolution) for r in rows[:4]] == sorted({(r.case, r.resolution) for r in rows}) + plan = select_rows(groups, 20261007, {"480p": 10}, frozenset(row.source for row in rows)) + assert not {r.source for r in plan} & {r.source for r in rows} + + +def test_pinned_amax_checks_sha_and_copies(tmp_path, single_process_group): # noqa: F811 + import hashlib + from fastvideo.train.models.minimax_h3.omniref_pdd import _pin_amax + table = tmp_path / "amax.json" + table.write_text('{"b0.ff.fc_in": 3.0}') + digest = hashlib.sha256(table.read_bytes()).hexdigest() + assert _pin_amax(str(table), digest, str(tmp_path / "run")) == digest + assert (tmp_path / "run" / f"amax-{digest[:12]}.json").is_file() + with pytest.raises(ValueError, match="pins"): + _pin_amax(str(table), "0" * 64, "") + + +def test_skip_blocks_keep_whole_blocks_bf16(cpu_loader, single_process_group): # noqa: F811 + methods = install_nvfp4_qad(_tiny_model(), _amax_table(), NVFP4QADPlan(skip_blocks=(1, ))) + assert methods and all(name.startswith("transformer_blocks.0.") for name in methods) + + +def test_fp32_boundary_modules_match_inference_and_stay_frozen(cpu_loader, single_process_group): # noqa: F811 + """The QAD loader keeps H3's FP32 boundary modules FP32 (as inference) and refuses to train them.""" + from types import SimpleNamespace + from fastvideo.train.models.minimax_h3.omniref_pdd import MiniMaxH3OmniRefPDDModel + model = _tiny_model().to(torch.float32) + assert not model.config.uniform_parameter_dtype + fp32 = {n for n, _ in model.named_parameters() if model._get_parameter_dtype(n, torch.bfloat16) == torch.float32} + assert fp32 and {n.split(".", 1)[0] for n in fp32} <= set(model._keep_in_fp32_modules) + model.requires_grad_(False) + fake = SimpleNamespace(transformer=model) + MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake) + model.proj_in.weight.requires_grad_(True) + with pytest.raises(ValueError, match="must stay frozen"): + MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake) + + +def test_block_first_call_keeps_the_fsdp_root(single_process_group): # noqa: F811 + """Stage 2 from a checkpoint calibrates reconstruction (student blocks called directly) before any full student + forward; without pinning the root first, FSDP2 makes block 0 the root and the next full forward fails.""" + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import ensure_fsdp_root_initialized + device = "cuda" if torch.cuda.is_available() else "cpu" + + class Toy(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = torch.nn.ModuleList(torch.nn.Linear(4, 4) for _ in range(2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for block in self.transformer_blocks: + x = block(x) + return x + + def sharded() -> Toy: + mesh = init_device_mesh(device, (1, )) + model = Toy().to(device) + for block in model.transformer_blocks: + fully_shard(block, mesh=mesh) + return fully_shard(model, mesh=mesh) + + x = torch.randn(2, 4, device=device) + broken = sharded() + broken.transformer_blocks[0](x) + with pytest.raises(RuntimeError, match="already been lazily initialized"): + broken(x) + fixed = sharded() + ensure_fsdp_root_initialized(fixed) + ensure_fsdp_root_initialized(fixed) # idempotent + fixed.transformer_blocks[0](x) + assert fixed(x).shape == (2, 4) + ensure_fsdp_root_initialized(torch.nn.Linear(2, 2)) # no FSDP: no-op + + +def test_step0_cache_round_trip_and_fingerprint(tmp_path): + from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0, + step0_fingerprint_inputs) + method = {"init_student_dcp": "/runs/e3/best", "noise_floor_eval": True, "eligibility_keys": ["endpoint/audio"], + "recon_group_weights": {"video": 1.0, "audio": 1.0}, "eval_every": 20} + data = {"eval_manifest": "/m.json", "eval_per_group": {"480p": 4, "768p": 2}, "cases": ["first_frame"]} + qad = {"linears": {"skip_blocks": [42, 43]}, "fp4_attention": {"two_level_p": True}} + models = dict(student_init_from="/ckpt/student", teacher_init_from="/ckpt/teacher") + inputs = step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **models) + metrics = {"eval/score": 0.31, "eval/x0_mse/video/rung0": 0.2, "eval/seconds": 6000.0, + "init/recon_step0/rung0": 0.04, "init/recon_step0/rung1": 0.03, + "init/noise_sigma/endpoint/audio": 0.05, "init/noise_sigma/score": 0.001} + path = tmp_path / "step0.json" + assert load_step0(path, inputs) is None + save_step0(path, inputs, metrics) + # Settings that cannot change step 0 (eval cadence) do not invalidate the cache. + same = step0_fingerprint_inputs({**method, "eval_every": 40}, data, qad, "ab" * 32, 30, **models) + cached = load_step0(path, same) + assert cached is not None and "eval/seconds" not in cached and cached["eval/score"] == 0.31 + summary, recon_norm, tolerances = restore_step0(cached, 2, recon=True, noise_floor=True, min_tolerance=0.02) + assert summary == {"score": 0.31, "x0_mse/video/rung0": 0.2} + assert recon_norm == [0.04, 0.03] + assert tolerances == {"endpoint/audio": pytest.approx(0.1), "score": 0.02} + # Training-plan keys only matter with provisional calibration (no amax_json). + assert load_step0(path, step0_fingerprint_inputs(method, {**data, "seed": 5}, qad, "ab" * 32, 30, **models)) + provisional = {**qad, "linears": {"skip_blocks": [42, 43], "provisional_calibration_rows": 2}} + for changed in (step0_fingerprint_inputs({**method, "init_student_dcp": "/runs/e3b/best"}, data, qad, "ab" * 32, 30, + **models), + step0_fingerprint_inputs(method, {**data, "eval_per_group": 1}, qad, "ab" * 32, 30, **models), + step0_fingerprint_inputs(method, data, {**qad, "linears": {"skip_blocks": []}}, "ab" * 32, 30, + **models), + step0_fingerprint_inputs(method, data, qad, "cd" * 32, 30, **models), + step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "student_init_from": "/x"}), + step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "teacher_init_from": "/x"}), + step0_fingerprint_inputs({**method, "audio_weight": 2.0}, data, qad, "ab" * 32, 30, **models), + step0_fingerprint_inputs(method, data, provisional, "ab" * 32, 30, **models)): + with pytest.raises(ValueError, match="other inputs"): + load_step0(path, changed) diff --git a/fastvideo/tests/train/utils/test_checkpoint.py b/fastvideo/tests/train/utils/test_checkpoint.py index f6dcf0c47..72c743087 100644 --- a/fastvideo/tests/train/utils/test_checkpoint.py +++ b/fastvideo/tests/train/utils/test_checkpoint.py @@ -234,6 +234,15 @@ def test_write_metadata_roundtrip_with_step(tmp_path: Path) -> None: assert loaded == {"step": 7} +def test_write_metadata_marks_tagged_weights_only(tmp_path: Path) -> None: + mgr = _make_manager(tmp_path) + tagged = tmp_path / "best" + tagged.mkdir() + mgr._write_metadata(tagged, step=30, weights_only=True) + loaded = CheckpointManager.load_metadata(tagged) + assert loaded == {"step": 30, "weights_only": True} + + def test_write_metadata_includes_raw_config(tmp_path: Path) -> None: raw = { "models": { diff --git a/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py b/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py index 2a6c1fbc3..3a2f52b94 100644 --- a/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py +++ b/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py @@ -7,6 +7,7 @@ from types import SimpleNamespace from diffusers import ConfigMixin, ModelMixin from diffusers.configuration_utils import register_to_config +import pytest from safetensors.torch import load_file import torch from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper @@ -16,6 +17,7 @@ from fastvideo.models.loader.weight_utils import safetensors_weights_iterator from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Ref2VAModularPipeline from fastvideo.train.entrypoint.dcp_to_diffusers import ( _native_export_state_dict, + _resolve_export_checkpoint, _save_diffusers_safetensors, _save_role_pretrained, _strict_reload_verify, @@ -23,6 +25,20 @@ from fastvideo.train.entrypoint.dcp_to_diffusers import ( from fastvideo.train.utils.lora import enable_lora_training, finalize_lora_training +def test_resolve_export_checkpoint_accepts_tagged_weights_only_dir(tmp_path: Path) -> None: + tagged = tmp_path / "best" + (tagged / "dcp").mkdir(parents=True) + (tagged / "dcp" / ".metadata").write_text("") + assert _resolve_export_checkpoint(str(tagged)) == tagged.resolve() + + +def test_resolve_export_checkpoint_rejects_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "best" + plain.mkdir() + with pytest.raises(ValueError, match="best"): + _resolve_export_checkpoint(str(plain)) + + class _ArchConfig: exclude_lora_layers: list[str] = [] diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index fd1611903..7d87139bf 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -40,6 +40,7 @@ import gc import json import os import sys +from pathlib import Path from typing import Any import fastvideo.envs as envs @@ -449,6 +450,15 @@ def _strict_reload_verify( logger.info("Strict reload verification passed.") +def _resolve_export_checkpoint(checkpoint_dir: str) -> Path: + """A DCP image directory, including a tagged weights-only one, or a regular checkpoint path.""" + raw_path = Path(os.path.expanduser(checkpoint_dir)).resolve() + if (raw_path / "dcp" / ".metadata").is_file(): + return raw_path + from fastvideo.train.utils.checkpoint import _resolve_resume_checkpoint + return _resolve_resume_checkpoint(checkpoint_dir, output_dir=checkpoint_dir) + + def convert( *, checkpoint_dir: str, @@ -468,10 +478,7 @@ def convert( from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, ) from fastvideo.train.utils.builder import build_from_config - from fastvideo.train.utils.checkpoint import ( - CheckpointManager, - _resolve_resume_checkpoint, - ) + from fastvideo.train.utils.checkpoint import CheckpointManager from fastvideo.train.utils.config import ( RunConfig, load_run_config, @@ -480,20 +487,18 @@ def convert( import torch.distributed.checkpoint as dcp # -- Resolve checkpoint directory -- - resolved = _resolve_resume_checkpoint( - checkpoint_dir, - output_dir=checkpoint_dir, - ) + resolved = _resolve_export_checkpoint(checkpoint_dir) dcp_dir = resolved / "dcp" if not dcp_dir.is_dir(): raise FileNotFoundError(f"Missing dcp/ under {resolved}") + metadata = CheckpointManager.load_metadata(resolved) if (resolved / "metadata.json").is_file() else {} + weights_only = bool(metadata.get("weights_only", False)) # -- Obtain config -- cfg: RunConfig if config_path is not None: cfg = load_run_config(config_path) else: - metadata = CheckpointManager.load_metadata(resolved) raw_config = metadata.get("config") if raw_config is None: raise ValueError("Checkpoint metadata.json does not " @@ -526,6 +531,8 @@ def convert( # -- Load DCP weights into the model -- states = method.checkpoint_state() + if weights_only: + states = {key: value for key, value in states.items() if key.startswith("roles.")} logger.info( "Loading DCP checkpoint from %s", resolved, diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py new file mode 100644 index 000000000..d605792a2 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-forward quantization-aware distillation (QAD) of a FastH3 OmniRef PDD student. + +Student = the teacher's own weights + NVFP4 fake-quantized block linears +(Stage A calibrated static FFN scales, unit attention/gate scales) + the +deployed sparse-FP4 attention numerics. Loss = per-forward regression of the +denoiser output (x0) to the frozen bf16 teacher, at identical inputs on the +eight-forward PDD ladder; no DMD, no adversarial term. + +Each optimizer micro-step takes one OmniRef row: the teacher walks its PDD +trajectory (teacher forcing), or with probability ``rollout_fraction`` the +student walks its own (DAgger-style exposure-bias correction); at every +forward both models see the same state and the student's x0 error is +backpropagated at once, inside that forward's scope. + +Guardrails: step-0 checks (quantization disabled -> bit-identical to the +teacher; quantized error vs the Stage A PTQ number), per-(rung, modality) +loss normalization by the step-0 PTQ error, held-out evaluation every +``eval_every`` steps (teacher-forced x0 rel-L2 per forward, student-rollout +endpoint, latent keyframe fidelity), best-checkpoint requests, and a stop +request after ``patience`` consecutive evaluations worse than step 0 (teacher +alignment counts full evaluations only). The step-0 baseline and the guard +state are checkpointed (``method.qad_baseline``), so a resumed run keeps them. +""" +from __future__ import annotations + +import hashlib +import json +import math +import time +from pathlib import Path +from typing import Any + +import torch + +from fastvideo.distributed import get_world_group +from fastvideo.logger import init_logger +from fastvideo.train.methods.base import LogScalar, TrainingMethod +from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction +from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0, + step0_fingerprint_inputs) +from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import (MODALITIES, RungStats, keyframe_latents, + rel_l2, step_packed, summarize, + target_slices, x0_pair) +from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler + +logger = init_logger(__name__) + +# Teacher alignment: the student's own sample (same seed) and its keyframes vs the teacher's; full evaluations only. +ALIGNMENT_KEYS = ("endpoint/video", "endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last") +BASELINE_FORMAT = 1 + + +class _BaselineState: + """DCP-stateful step-0 baseline and guard state of an ``OmniRefPDDQADMethod`` (one JSON string).""" + + def __init__(self, method: OmniRefPDDQADMethod) -> None: + self.method = method + + def state_dict(self) -> dict[str, Any]: + return {"json": json.dumps(self.method._baseline_state(), sort_keys=True)} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.method._load_baseline_state(json.loads(state_dict["json"])) + + +def _unit_hash(*parts: Any) -> float: + digest = hashlib.sha256(":".join(str(p) for p in parts).encode()).digest() + return int.from_bytes(digest[:8], "little") / 2.0**64 + + +class OmniRefPDDQADMethod(TrainingMethod): + """Teacher-forced (+ optional student-rollout) per-forward x0 regression; see module docstring.""" + + def __init__(self, *, cfg: Any, role_models: dict[str, Any]) -> None: + super().__init__(cfg=cfg, role_models=role_models) + if "teacher" not in role_models: + raise ValueError("OmniRefPDDQADMethod needs a frozen 'teacher' role") + self.teacher = role_models["teacher"] + if getattr(self.teacher, "_trainable", False) or self.teacher.qad_linears: + raise ValueError("the teacher must be frozen and unquantized") + mc = self.method_config + self.weights = {"video": float(mc.get("video_weight", 1.0)), "audio": float(mc.get("audio_weight", 1.0))} + self.rollout_fraction = float(mc.get("rollout_fraction", 0.0)) + self.rollout_start_step = int(mc.get("rollout_start_step", 0)) + # Student forward+backward on this many of the 8 forwards per row (the teacher walks all of them). + self.rungs_per_row = int(mc.get("rungs_per_row", 0)) + # "output": per-forward x0 regression. "recon": layer-wise reconstruction only. + # "output+recon": x0 regression + lambda * reconstruction. + self.objective = str(mc.get("objective", "output")) + if self.objective not in ("output", "recon", "output+recon"): + raise ValueError(f"objective must be output, recon or output+recon, got {self.objective!r}") + recon_lambda = float(mc.get("recon_lambda", 1.0)) + self.output_weight = {"output": 1.0, "recon": 0.0, "output+recon": 1.0 / (1.0 + recon_lambda)}[self.objective] + self.recon_weight = { + "output": 0.0, + "recon": 1.0, + "output+recon": recon_lambda / (1.0 + recon_lambda) + }[self.objective] + if self.objective == "recon" and float(mc.get("rollout_fraction", 0.0)) > 0: + raise ValueError("reconstruction trains on teacher states only; set rollout_fraction: 0") + # Best checkpoint by this eval key (lower is better): "score" (teacher-forced x0 rel-L2) or + # "endpoint/video" (the student's own sample vs the teacher's: teacher alignment). + self.select_metric = str(mc.get("select_metric", "score")) + # Save each eval's final latents (student's own sample, teacher's) for side-by-side decoding. + self.eval_save_dir = mc.get("eval_save_dir") + # Full evaluation (with the student's own 8-step sample) every this many steps; teacher-forced only otherwise. + self.full_eval_every = int(mc.get("full_eval_every", 0)) or int(mc.get("eval_every", 50)) + # Start the student from another run's tagged weights (e.g. Stage 1 best) instead of the teacher's. + self.init_student_dcp = mc.get("init_student_dcp") + self._eval_iteration = 0 + self.eval_every = int(mc.get("eval_every", 50)) + self.patience = int(mc.get("patience", 3)) + self.fidelity_tolerance = float(mc.get("fidelity_tolerance", 0.05)) + self.ptq_reference_json = mc.get("ptq_reference_json") + self.ptq_tolerance = float(mc.get("ptq_tolerance", 0.05)) + self.require_bitwise_init = bool(mc.get("require_bitwise_init", True)) + self.max_eval_rows = int(mc.get("max_eval_rows", 0)) + self.save_best = bool(mc.get("save_best", True)) + # Reconstruction per row group (target video, target audio, conditioning): over the whole packed sequence + # audio is ~1.4% of rows and text tokens dominate the norm, so an ungrouped ratio barely trains audio. + self.recon_group_weights = dict( + mc.get("recon_group_weights") or { + "video": 1.0, + "audio": 1.0, + "condition": 0.25 + }) + # "step0": divide each forward's reconstruction loss by its step-0 value (as the x0 terms are), so + # recon_lambda means what it says; raw reconstruction gradients are ~1e3-1e4x smaller than x0's. + self.recon_normalize = str(mc.get("recon_normalize", "none")) + if self.recon_normalize not in ("none", "step0"): + raise ValueError(f"recon_normalize must be none or step0, got {self.recon_normalize!r}") + self.recon_norm_rows = int(mc.get("recon_norm_rows", 4)) + self.recon_norm: list[float] | None = None + # Full evaluations roll out (8 more student forwards) on at most this many rows; others are teacher-forced. + self.max_rollout_rows = int(mc.get("max_rollout_rows", 0)) + # A checkpoint is "best" only if these full-eval metrics stay <= step0 * (1 + tolerance). With + # noise_floor_eval the tolerance is 2 sigma of a seed+1 re-evaluation at step 0 (>= min_eligibility_tolerance); + # otherwise fidelity_tolerance. + self.eligibility_keys = tuple( + mc.get("eligibility_keys") + or ("endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last", "x0_rel_l2/audio/mean")) + self.noise_floor_eval = bool(mc.get("noise_floor_eval", False)) + self.min_eligibility_tolerance = float(mc.get("min_eligibility_tolerance", 0.02)) + self.tolerances: dict[str, float] = {} + # Stage A's per-rung reference metric for the step-0 PTQ match (v_rel is sign-independent). + self.ptq_reference_metric = str(mc.get("ptq_reference_metric", "v_rel_l2")) + self._recon_criteria: dict[str, Any] = {} + # Reuse (or write) this run's step-0 evaluation, keyed by a fingerprint of everything it depends on. + self.step0_cache = mc.get("step0_cache") + self._method_config_snapshot = mc + self.student.init_preprocessors(cfg.training) + tc = self.training_config + params = [p for p in self.student.transformer.parameters() if p.requires_grad] + self.optimizer, self.lr_scheduler = build_optimizer_and_scheduler(params=params, + optimizer_config=tc.optimizer, + loop_config=tc.loop, + learning_rate=float( + tc.optimizer.learning_rate), + betas=tuple(tc.optimizer.betas), + scheduler_name=str(tc.optimizer.lr_scheduler)) + self.normalizers: dict[str, list[float]] | None = None + self.step0: dict[str, float] = {} + self.best_score = math.inf + # Consecutive evaluations with a worse teacher-forced score (every eval) and with worse teacher alignment + # (full evals only); ``bad_evals`` is the larger. + self.bad_score_evals = 0 + self.bad_alignment_evals = 0 + self.bad_evals = 0 + self.stop_requested = False + self._checkpoint_request: str | None = None + if self.init_student_dcp: + import torch.distributed.checkpoint as dcp + + # The same DCP wrapper the checkpoint manager saves tagged/regular student weights with. + states = { + key: value + for key, value in self.checkpoint_state().items() if key == "roles.student.transformer" + } + if not states: + raise ValueError("init_student_dcp needs a trainable 'student' role with a transformer") + dcp.load(states, checkpoint_id=str(Path(self.init_student_dcp) / "dcp")) + logger.info("QAD student initialized from %s", self.init_student_dcp) + world = get_world_group() + self.sp_size = int(tc.distributed.sp_size or 1) + self.dp_rank, self.dp_size = world.rank // self.sp_size, world.world_size // self.sp_size + self.is_sp_leader = world.rank % self.sp_size == 0 + self.sampling_seed = int(tc.data.seed or 0) + + # ------------------------------------------------------------------ optimizer plumbing + def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]: + return [self.optimizer] + + def get_lr_schedulers(self, iteration: int) -> list[Any]: + return [self.lr_scheduler] + + @property + def _optimizer_dict(self) -> dict[str, Any]: + return {"student": self.optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self.lr_scheduler} + + def backward(self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int = 1) -> None: + """No-op: each forward's loss was backpropagated inside its own rung scope.""" + + def consume_checkpoint_request(self) -> str | None: + request, self._checkpoint_request = self._checkpoint_request, None + return request + + # ------------------------------------------------------------------ resume + def checkpoint_state(self) -> dict[str, Any]: + """Roles, optimizer and scheduler, plus the step-0 baseline and guard state. + + Without the baseline a resumed run would re-measure step 0 on its trained weights (wrong loss scale, + stop history and best-checkpoint reference); a checkpoint without ``method.qad_baseline`` cannot be + resumed (DCP reports the missing key). + """ + return {**super().checkpoint_state(), "method.qad_baseline": _BaselineState(self)} + + def _baseline_state(self) -> dict[str, Any]: + provisional = {} + if getattr(self.student, "provisional_rows", 0) > 0: + provisional = { + name: m.input_amax + for name, m in self.student.qad_linears.items() if m.act_scale == "static" and m.input_amax is not None + } + return { + "format": BASELINE_FORMAT, + "normalizers": self.normalizers, + "step0": self.step0, + "recon_norm": self.recon_norm, + "tolerances": self.tolerances, + "best_score": self.best_score, + "bad_score_evals": self.bad_score_evals, + "bad_alignment_evals": self.bad_alignment_evals, + "provisional_amax": provisional, + } + + def _load_baseline_state(self, state: dict[str, Any]) -> None: + if state.get("format") != BASELINE_FORMAT: + raise ValueError(f"unsupported QAD baseline state format {state.get('format')!r}") + if state["normalizers"] is None: + raise ValueError("the checkpoint has no step-0 baseline (saved before the step-0 evaluation)") + self.normalizers = {m: [float(v) for v in values] for m, values in state["normalizers"].items()} + self.step0 = {k: float(v) for k, v in state["step0"].items()} + self.recon_norm = None if state["recon_norm"] is None else [float(v) for v in state["recon_norm"]] + self.tolerances = {k: float(v) for k, v in state["tolerances"].items()} + self.best_score = float(state["best_score"]) + self.bad_score_evals = int(state["bad_score_evals"]) + self.bad_alignment_evals = int(state["bad_alignment_evals"]) + self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals) + for name, amax in state.get("provisional_amax", {}).items(): + method = self.student.qad_linears.get(name) + if method is None: + raise KeyError(f"checkpoint has a provisional scale for {name}, which is not a QAD linear here") + if method.input_amax is None: + method.set_input_amax(float(amax)) + if get_world_group().rank == 0: + logger.info("QAD step-0 baseline restored from the checkpoint (score %.5f, bad evals %d/%d)", + self.step0.get("score", float("nan")), self.bad_score_evals, self.bad_alignment_evals) + + # ------------------------------------------------------------------ one training row + def _row(self, spec: Any) -> Any: + return self.student.load_row(spec) + + def single_train_step(self, batch: dict[str, Any], iteration: int) -> tuple[dict, dict, dict[str, LogScalar]]: + if self.normalizers is None: + raise RuntimeError("step-0 evaluation must run before training (it sets the loss normalizers)") + spec = batch["row"] + prepared = self._row(spec) + rollout = self._use_rollout(spec.id, iteration) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1)) + num_rungs = self.student.num_rungs + trained = self._trained_rungs(spec.id, iteration, num_rungs) + video, audio = prepared.video, prepared.audio + total = torch.zeros((), device=self.student.device) + metrics: dict[str, LogScalar] = { + "train/rollout": float(rollout), + "train/tokens": float(prepared.layout.sequence_length) + } + started = time.perf_counter() + for rung in range(num_rungs): + recon = None + with torch.no_grad(), self.teacher.rung_scope(prepared, rung, timesteps) as rt: + if rung in trained and self.recon_weight > 0: + # Student blocks train inside the teacher forward, on the teacher's sub-layer inputs. + norm = self.recon_norm[rung] if self.recon_norm else 1.0 + scale = self.recon_weight / (norm * 2 * len(self.student.transformer.transformer_blocks) * + len(trained) * accum) + with BlockReconstruction(self.teacher.transformer, self.student.transformer, scale, + self._recon_criterion(prepared)) as recon: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + else: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + if recon is not None: + recon_loss = recon.loss() + if not torch.isfinite(recon_loss): + raise FloatingPointError(f"non-finite reconstruction loss at step {iteration}, rung {rung}") + norm = self.recon_norm[rung] if self.recon_norm else 1.0 + total += self.recon_weight * recon_loss / (norm * len(trained)) + metrics[f"train/recon_vs_step0/rung{rung}"] = recon_loss / norm + for key, value in recon.metrics("train/recon").items(): + metrics[f"{key}/rung{rung}"] = value + if rung not in trained or self.output_weight == 0: + if rollout: # the student's own trajectory still needs its (no-grad) forward here + with torch.no_grad(), self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, + *(student_out if rollout else teacher_out)) + continue + with self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + pair = x0_pair(prepared, schedulers, rung, video, audio, { + "student": student_out, + "teacher": teacher_out + }) + loss = self._rung_loss(pair, rung, metrics) + if not torch.isfinite(loss): + raise FloatingPointError(f"non-finite QAD loss at step {iteration}, row {spec.id}, rung {rung}") + (self.output_weight * loss / (len(trained) * accum)).backward() + total += self.output_weight * loss.detach() / len(trained) + step_out = tuple(v.detach() for v in student_out) if rollout else teacher_out + video, audio = step_packed(schedulers, prepared, rung, video, audio, *step_out) + metrics["train/row_seconds"] = time.perf_counter() - started + if get_world_group().rank == 0: + logger.info("QAD step %d row %s (%s, %d tokens, rollout=%s): loss %.5f in %.1fs", iteration, spec.id, + spec.case, prepared.layout.sequence_length, rollout, float(total), metrics["train/row_seconds"]) + return {"total_loss": total}, {}, metrics + + def _recon_criterion(self, prepared: Any) -> Any: + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import (GroupedRelError, local_row_groups, + packed_row_groups) + key = prepared.spec.id + if key not in self._recon_criteria: + self._recon_criteria = {} # one row at a time + groups = packed_row_groups(prepared.layout, int(prepared.layout.sequence_length)) + self._recon_criteria[key] = GroupedRelError(local_row_groups(groups, self.student.device), + self.recon_group_weights) + return self._recon_criteria[key] + + @torch.no_grad() + def _calibrate_recon(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]: + """Step-0 reconstruction error per forward (measure-only), the Stage-2 normalizer.""" + num_rungs = self.student.num_rungs + sums = torch.zeros(num_rungs, dtype=torch.float64, device=self.student.device) + count = torch.zeros((), dtype=torch.float64, device=self.student.device) + per_rank = -(-self.recon_norm_rows // self.dp_size) + for spec, counted in rows[:per_rank]: + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(num_rungs): + with self.teacher.rung_scope(prepared, rung, + timesteps) as rt, BlockReconstruction(self.teacher.transformer, + self.student.transformer, + 0.0, + self._recon_criterion(prepared), + measure_only=True) as recon: + out = self.teacher.forward_rung(prepared, video, audio, rt) + if counted: + sums[rung] += recon.loss().double() + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + count += float(counted) + if not self.is_sp_leader: + sums.zero_() + count.zero_() + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(sums) + torch.distributed.all_reduce(count) + self.recon_norm = [max(float(v), 1e-12) for v in (sums / count.clamp_min(1.0))] + return {f"init/recon_step0/rung{rung}": value for rung, value in enumerate(self.recon_norm)} + + def _noise_floor(self, rows: list[tuple[Any, bool]], step0: dict[str, float]) -> dict[str, LogScalar]: + """Re-evaluate step 0 at seed+1; per-metric tolerance = max(2 sigma, minimum) with sigma relative.""" + import dataclasses + shifted = [(dataclasses.replace(spec, seed=spec.seed + 1), counted) for spec, counted in rows] + other = self._evaluate(shifted, rollout=True, save_tag="seed1") + out: dict[str, LogScalar] = {} + for key in set(self.eligibility_keys) | {"endpoint/video", "score"}: + if key in step0 and key in other: + mean = 0.5 * (step0[key] + other[key]) + sigma = abs(step0[key] - other[key]) / (math.sqrt(2.0) * max(mean, 1e-12)) + self.tolerances[key] = max(2.0 * sigma, self.min_eligibility_tolerance) + out[f"init/noise_sigma/{key}"] = sigma + out[f"init/seed1/{key}"] = other[key] + return out + + def _tolerance(self, key: str) -> float: + return self.tolerances.get(key, self.fidelity_tolerance) + + def _sampling_key(self, row_id: str) -> str: + """Keep forward/backward choices in lockstep across the FSDP shard group. + + FSDP spans data-parallel groups, whose rows differ. A row-dependent choice can make one group gather + teacher block weights while another gathers student block weights, corrupting both forwards. Single-group + runs retain their existing row-dependent sampling. + """ + return row_id if self.dp_size == 1 else f"fsdp-shared:{self.sampling_seed}" + + def _use_rollout(self, row_id: str, iteration: int) -> bool: + return (iteration >= self.rollout_start_step + and _unit_hash(self._sampling_key(row_id), iteration) < self.rollout_fraction) + + def _trained_rungs(self, row_id: str, iteration: int, num_rungs: int) -> set[int]: + """All forwards, or a deterministic subset shared by every FSDP rank.""" + if self.rungs_per_row <= 0 or self.rungs_per_row >= num_rungs: + return set(range(num_rungs)) + order = sorted(range(num_rungs), + key=lambda rung: _unit_hash(self._sampling_key(row_id), iteration, "rung", rung)) + return set(order[:self.rungs_per_row]) + + def _rung_loss(self, pair: dict[str, dict[str, torch.Tensor]], rung: int, metrics: dict[str, + LogScalar]) -> torch.Tensor: + assert self.normalizers is not None + loss = torch.zeros((), device=self.student.device) + for modality in MODALITIES: + student, teacher = pair[modality]["student"], pair[modality]["teacher"] + mse = (student - teacher).pow(2).mean() + normalized = mse / self.normalizers[modality][rung] + loss = loss + self.weights[modality] * normalized + metrics[f"train/rel_mse_vs_ptq/{modality}/rung{rung}"] = normalized.detach() + return loss / sum(self.weights.values()) + + # ------------------------------------------------------------------ evaluation and guardrails + def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]: + # Smoke runs without Stage A's table: calibrate the FFN scales before anything uses them, whether step 0 + # is measured, restored from the cache, or skipped by a resume. A no-op once every scale is set. + calibration = self._provisional_calibration() + first = self.normalizers is None + if not first and (self.eval_every <= 0 or iteration % self.eval_every): + return dict(calibration) + rows = self._rank_rows() + metrics: dict[str, LogScalar] = dict(calibration) + started = time.perf_counter() + self._eval_iteration = iteration + cached = self._load_step0() if first else None + if cached is not None: + summary = self._restore_step0(cached) + metrics.update(cached) + else: + if first and iteration == 0 and not self.init_student_dcp: + # Fresh start only: a resumed run (e.g. Stage 2 from a Stage-1 checkpoint) is no longer the teacher. + metrics.update(self._initial_checks(rows)) + if first and self._needs_recon_norm(): + metrics.update(self._calibrate_recon(rows)) + full = first or (iteration % self.full_eval_every == 0) + summary = self._evaluate(rows, rollout=full) + metrics.update({f"eval/{k}": v for k, v in summary.items()}) + if first and self.noise_floor_eval: + metrics.update(self._noise_floor(rows, summary)) + if first and self.step0_cache and get_world_group().rank == 0: + save_step0(self.step0_cache, self._step0_inputs(), {k: float(v) for k, v in metrics.items()}) + if first: + self.normalizers = { + m: [max(summary[f"x0_mse/{m}/rung{r}"], 1e-20) for r in range(self.student.num_rungs)] + for m in MODALITIES + } + self.step0 = dict(summary) + self.best_score = summary.get(self.select_metric, summary["score"]) + else: + metrics.update(self._guardrails(iteration, summary)) + metrics["eval/seconds"] = time.perf_counter() - started + metrics["eval/score_vs_ptq"] = summary["score"] / max(self.step0.get("score", summary["score"]), 1e-30) + if get_world_group().rank == 0: + logger.info("QAD eval @%d: score %.5f (PTQ %.5f) %s", iteration, summary["score"], + self.step0.get("score", float("nan")), + json.dumps({ + k: round(float(v), 5) + for k, v in metrics.items() + })) + return metrics + + def _needs_recon_norm(self) -> bool: + return self.recon_normalize == "step0" and self.recon_weight > 0 + + def _step0_inputs(self) -> dict[str, Any]: + return step0_fingerprint_inputs(self._method_config_snapshot, + self.student.data_config, + self.student.qad_config, + self.student.amax_sha256, + len(self.student.eval_rows), + student_init_from=self.student._init_from, + teacher_init_from=self.teacher._init_from) + + def _load_step0(self) -> dict[str, float] | None: + if not self.step0_cache: + return None + cached = load_step0(self.step0_cache, self._step0_inputs()) + if cached is not None and get_world_group().rank == 0: + logger.info("QAD step 0 restored from %s (no step-0 evaluation)", self.step0_cache) + return cached + + def _restore_step0(self, cached: dict[str, float]) -> dict[str, float]: + summary, recon_norm, tolerances = restore_step0(cached, self.student.num_rungs, self._needs_recon_norm(), + self.noise_floor_eval, self.min_eligibility_tolerance) + if recon_norm is not None: + self.recon_norm = recon_norm + self.tolerances.update(tolerances) + return summary + + def _rank_rows(self) -> list[tuple[Any, bool]]: + """This data-parallel group's eval rows, padded so every group runs the same number of forwards. + + FSDP all-gathers span every rank, so groups must stay in lockstep; padded + repeats run but are not counted (``(row, counted)`` pairs). + """ + rows = self.student.eval_rows[:self.max_eval_rows] if self.max_eval_rows else self.student.eval_rows + if not rows: + return [] + per_rank = -(-len(rows) // self.dp_size) + return [(rows[(i * self.dp_size + self.dp_rank) % len(rows)], i * self.dp_size + self.dp_rank < len(rows)) + for i in range(per_rank)] + + def _eligible(self, summary: dict[str, float]) -> bool: + """Full-eval audio/keyframe metrics within step0 * (1 + tolerance); a teacher-forced-only eval cannot tell.""" + if "endpoint/video" not in summary: + return False + return all(summary[key] <= self.step0[key] * (1.0 + self._tolerance(key)) for key in self.eligibility_keys + if key in self.step0 and key in summary) + + def _guardrails(self, iteration: int, summary: dict[str, float]) -> dict[str, LogScalar]: + worse_score = summary["score"] > self.step0["score"] + self.bad_score_evals = self.bad_score_evals + 1 if worse_score else 0 + # Teacher alignment is measured by full evaluations only; a teacher-forced-only one leaves its count as is. + measured = [key for key in ALIGNMENT_KEYS if key in summary and key in self.step0] + if measured: + misaligned = any(summary[key] > self.step0[key] * (1.0 + self._tolerance(key)) for key in measured) + self.bad_alignment_evals = self.bad_alignment_evals + 1 if misaligned else 0 + self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals) + selected = summary.get(self.select_metric) + eligible = self._eligible(summary) + if selected is not None and selected < self.best_score and eligible: + self.best_score = selected + self._checkpoint_request = "best" if self.save_best else None + if self.bad_evals >= self.patience: + self.stop_requested = True + logger.warning("QAD stop: %d consecutive evaluations worse than the PTQ start (step %d)", self.bad_evals, + iteration) + return { + "eval/eligible": float(eligible), + "eval/bad_evals": float(self.bad_evals), + "eval/bad_score_evals": float(self.bad_score_evals), + "eval/bad_alignment_evals": float(self.bad_alignment_evals), + "eval/best_score": self.best_score, + "eval/stop_requested": float(self.stop_requested) + } + + @torch.no_grad() + def _evaluate(self, + rows: list[tuple[Any, bool]], + *, + rollout: bool, + attention: bool = True, + linears: bool = True, + save_tag: str = "") -> dict[str, float]: + num_rungs = self.student.num_rungs + stats = RungStats(num_rungs) + counted_stats = stats + # Same rollout count on every data-parallel group (lockstep), covering the first max_rollout_rows rows. + rollout_rows = -(-self.max_rollout_rows // self.dp_size) if self.max_rollout_rows else len(rows) + with self.student.quantization(linears=linears, attention=attention): + for position, (spec, counted) in enumerate(rows): + stats = counted_stats if counted else RungStats(num_rungs) + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(num_rungs): + with self.teacher.rung_scope(prepared, rung, timesteps) as rt: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + with self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + stats.add_rung( + rung, + x0_pair(prepared, schedulers, rung, video, audio, { + "student": student_out, + "teacher": teacher_out + })) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *teacher_out) + if rollout and position < rollout_rows: + self._rollout_metrics(stats, + prepared, + schedulers, + timesteps, + video, + audio, + save=counted, + save_tag=save_tag) + means = counted_stats.reduce(contribute=self.is_sp_leader) + summary = summarize(means, num_rungs, self.weights) + for name in ("first", "last"): + student, teacher = summary.get(f"keyframe/student/{name}"), summary.get(f"keyframe/teacher/{name}") + if student is not None and teacher is not None: + summary[f"keyframe/drop/{name}"] = student / max(teacher, 1e-12) - 1.0 + return summary + + def _rollout_metrics(self, + stats: RungStats, + prepared: Any, + schedulers: tuple[Any, Any], + timesteps: Any, + teacher_video: torch.Tensor, + teacher_audio: torch.Tensor, + save: bool = False, + save_tag: str = "") -> None: + """Student's own 8-forward sample vs the teacher's: endpoint rel-L2 and latent keyframe fidelity.""" + video, audio = prepared.video, prepared.audio + for rung in range(self.student.num_rungs): + with self.student.rung_scope(prepared, rung, timesteps) as rt: + out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + cut = target_slices(prepared) + if save and self.eval_save_dir and self.is_sp_leader: + step_dir = f"step{self._eval_iteration:05d}" + (f"_{save_tag}" if save_tag else "") + path = Path(self.eval_save_dir) / step_dir / f"{prepared.spec.id}.pt" + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "student_video": video[cut["video"]].float().cpu(), + "teacher_video": teacher_video[cut["video"]].float().cpu(), + # Audio for decoded stats (spectral centroid, >8 kHz share, dynamic range, log-mel distance). + "student_audio": audio[cut["audio"]].float().cpu(), + "teacher_audio": teacher_audio[cut["audio"]].float().cpu(), + "latent_shape": prepared.latent_shape, + "case": prepared.spec.case, + "seed": prepared.spec.seed + }, + path) + for modality, student, teacher in (("video", video, teacher_video), ("audio", audio, teacher_audio)): + s, t = student[cut[modality]].float(), teacher[cut[modality]].float() + stats.add_scalar(f"endpoint.{modality}", (s - t).norm() / t.norm().clamp_min(1e-12)) + teacher_frames = keyframe_latents(prepared, teacher_video) + for name, (frame, reference) in keyframe_latents(prepared, video).items(): + stats.add_scalar(f"keyframe.vs_teacher.{name}", rel_l2(frame, teacher_frames[name][0])) + if reference is not None: + stats.add_scalar(f"keyframe.student.{name}", rel_l2(frame, reference)) + stats.add_scalar(f"keyframe.teacher.{name}", rel_l2(teacher_frames[name][0], reference)) + + def _initial_checks(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]: + """Bit-identity with quantization off, and the PTQ error with bf16 attention (Stage A's protocol).""" + out: dict[str, LogScalar] = {} + diff = torch.zeros((), device=self.student.device, dtype=torch.float64) + if rows: + prepared = self._row(rows[0][0]) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + with torch.no_grad(), self.student.quantization(linears=False, attention=False): + for rung in (0, self.student.num_rungs - 1): + outs = [] + for model in (self.teacher, self.student): + with model.rung_scope(prepared, rung, timesteps) as rt: + outs.append(model.forward_rung(prepared, prepared.video, prepared.audio, rt)) + for teacher_t, student_t in zip(*outs, strict=True): + diff = torch.maximum(diff, (teacher_t.double() - student_t.double()).abs().max()) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(diff, op=torch.distributed.ReduceOp.MAX) + out["init/noquant_vs_teacher_max_abs"] = float(diff) + sha = getattr(self.student, "amax_sha256", "") + if sha: # first 32 bits of the pinned calibration table's sha256, for W&B + out["init/amax_sha256_u32"] = float(int(sha[:8], 16)) + if self.require_bitwise_init and float(diff) != 0.0: + raise RuntimeError(f"student with quantization disabled differs from the teacher (max |diff| " + f"{float(diff):.3e}); refusing to train from a student that is not the teacher") + weights_only = self._evaluate(rows, rollout=False, attention=False) + out.update({ + f"init/ptq_bf16_attn/{k}": v + for k, v in weights_only.items() if k.startswith(("x0_rel", "v_rel", "score")) + }) + out.update(self._compare_ptq_reference(weights_only)) + if getattr(self.student, "compression", None) is not None: + # Training-free token compression alone (bf16 linears): how much of the step-0 error it contributes. + compress_only = self._evaluate(rows, rollout=False, linears=False) + out.update({ + f"init/compress_only/{k}": v + for k, v in compress_only.items() if k.startswith(("x0_rel", "v_rel", "score")) + }) + return out + + @torch.no_grad() + def _provisional_calibration(self) -> dict[str, LogScalar]: + """Smoke runs without Stage A's table: max-calibrate the FFN inputs on dense teacher trajectories.""" + from fastvideo.layers.quantization.nvfp4_qad import calibrating, nvfp4_qad_export_scales + pending = [m for m in self.student.qad_linears.values() if m.act_scale == "static" and m.input_amax is None] + if not pending: + return {} + count = int(getattr(self.student, "provisional_rows", 0)) + rows = self.student.train_rows[:count] + per_rank = -(-len(rows) // self.dp_size) + mine = [rows[(i * self.dp_size + self.dp_rank) % len(rows)] for i in range(per_rank)] + with self.student.quantization(linears=False, attention=False), calibrating(self.student.qad_linears): + for spec in mine: + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(self.student.num_rungs): + with self.student.rung_scope(prepared, rung, timesteps) as rt: + out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + table = nvfp4_qad_export_scales(self.student.qad_linears) + if get_world_group().rank == 0: + path = Path(self.training_config.checkpoint.output_dir) / "provisional_amax.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(table, indent=1, sort_keys=True)) + logger.warning("PROVISIONAL FFN scales from %d rows written to %s; use Stage A's table for real runs", + len(rows), path) + return {"init/provisional_calibration_rows": float(len(rows))} + + def _compare_ptq_reference(self, weights_only: dict[str, float]) -> dict[str, LogScalar]: + """Compare step-0 weights-only x0 rel-L2 per rung to Stage A's (``{"video": [...], "audio": [...]}``).""" + if not self.ptq_reference_json: + return {} + reference = json.loads(Path(self.ptq_reference_json).read_text()) + deviations = [] + for modality in MODALITIES: + for rung, value in enumerate(reference.get(modality, [])): + ours = weights_only.get(f"{self.ptq_reference_metric}/{modality}/rung{rung}") + if ours is not None and value: + deviations.append(abs(ours / float(value) - 1.0)) + worst = max(deviations, default=float("nan")) + if deviations and worst > self.ptq_tolerance: + raise RuntimeError(f"step-0 PTQ error deviates {worst:.1%} from Stage A's ({self.ptq_reference_json}); " + "the QAD student does not start at the calibrated PTQ point") + return {"init/ptq_vs_stage_a_max_rel_dev": worst} + + +__all__ = ["OmniRefPDDQADMethod"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py new file mode 100644 index 000000000..7ed74bee6 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-forward (T1) error and reference-fidelity bookkeeping for PDD QAD. + +All quantities are on the target rows only (conditioning rows are inputs). +``x0`` is the scheduler's own denoised estimate, ``x_t + sigma_t * v`` with +``sigma_t = 1 - timestep``, so x0 error is ``sigma_t`` times velocity error. +""" +from __future__ import annotations + +from collections import defaultdict +from typing import Any + +import torch +import torch.distributed as dist + +MODALITIES = ("video", "audio") + + +def target_slices(prepared: Any) -> dict[str, slice]: + return { + "video": slice(prepared.layout.num_condition_video_rows, None), + "audio": slice(prepared.layout.num_condition_audio_rows, None) + } + + +def denoised(state: torch.Tensor, velocity: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + """x0 exactly as ``MiniMaxH3Scheduler.step`` forms it (fp32).""" + return state.float() + (1 - timestep.to(torch.float32)) * velocity.float() + + +def scheduler_step(scheduler: Any, rung: int, velocity: torch.Tensor, state: torch.Tensor) -> torch.Tensor: + """``scheduler.step`` at forward ``rung`` without relying on its internal counter.""" + scheduler._step_index = rung + return scheduler.step(velocity.float(), scheduler.timesteps[rung], state, return_dict=False)[0] + + +def step_packed(schedulers: tuple[Any, Any], prepared: Any, rung: int, video: torch.Tensor, audio: torch.Tensor, + video_v: torch.Tensor, audio_v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Advance the packed target rows by one PDD forward (conditioning rows untouched), as the sampler does.""" + cut = target_slices(prepared) + video, audio = video.clone(), audio.clone() + video[cut["video"]] = scheduler_step(schedulers[0], rung, video_v[cut["video"]], video[cut["video"]]) + audio[cut["audio"]] = scheduler_step(schedulers[1], rung, audio_v[cut["audio"]], audio[cut["audio"]]) + return video, audio + + +def x0_pair(prepared: Any, schedulers: tuple[Any, Any], rung: int, video: torch.Tensor, audio: torch.Tensor, + outputs: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> dict[str, dict[str, torch.Tensor]]: + """``{modality: {"student": x0, "teacher": x0, "student_v": v, "teacher_v": v}}`` on target rows.""" + cut = target_slices(prepared) + states = {"video": video, "audio": audio} + result: dict[str, dict[str, torch.Tensor]] = {} + for index, modality in enumerate(MODALITIES): + timestep = schedulers[index].timesteps[rung] + entry = {} + for role, (video_v, audio_v) in outputs.items(): + velocity = (video_v, audio_v)[index][cut[modality]] + entry[role] = denoised(states[modality][cut[modality]], velocity, timestep) + entry[f"{role}_v"] = velocity.float() + result[modality] = entry + return result + + +class RungStats: + """Sums of per-rung error terms, reduced over data-parallel groups (one SP rank contributes).""" + + KEYS = ("mse", "err_sq", "ref_sq", "verr_sq", "vref_sq", "dot", "s_sq") + EXTRA_KEYS = ("endpoint.video", "endpoint.audio", "keyframe.vs_teacher.first", "keyframe.vs_teacher.last", + "keyframe.student.first", "keyframe.teacher.first", "keyframe.student.last", "keyframe.teacher.last") + + def __init__(self, num_rungs: int) -> None: + self.num_rungs = num_rungs + self.sums: dict[str, torch.Tensor] = {} + self.counts: dict[str, float] = defaultdict(float) + # Ranks hold different rows (cases), so every rank reduces this same fixed key list. + self.all_keys = [f"{m}.{r}.{k}" for m in MODALITIES for r in range(num_rungs) for k in self.KEYS] + self.all_keys += list(self.EXTRA_KEYS) + + def _add(self, key: str, value: torch.Tensor) -> None: + self.sums[key] = self.sums.get(key, 0.0) + value.detach().double() + + def add_rung(self, rung: int, pair: dict[str, dict[str, torch.Tensor]]) -> None: + for modality, entry in pair.items(): + s, t, sv, tv = entry["student"], entry["teacher"], entry["student_v"], entry["teacher_v"] + prefix = f"{modality}.{rung}" + self._add(f"{prefix}.mse", (s - t).pow(2).mean()) + self._add(f"{prefix}.err_sq", (s - t).pow(2).sum()) + self._add(f"{prefix}.ref_sq", t.pow(2).sum()) + self._add(f"{prefix}.verr_sq", (sv - tv).pow(2).sum()) + self._add(f"{prefix}.vref_sq", tv.pow(2).sum()) + self._add(f"{prefix}.dot", (s * t).sum()) + self._add(f"{prefix}.s_sq", s.pow(2).sum()) + for key in self.KEYS: + self.counts[f"{prefix}.{key}"] += 1 + + def add_scalar(self, key: str, value: torch.Tensor | float) -> None: + if key not in self.EXTRA_KEYS: + raise KeyError(f"{key} is not a reduced RungStats key") + self._add(key, torch.as_tensor(value, dtype=torch.float64, device=self._device())) + self.counts[key] += 1 + + def _device(self) -> torch.device: + return torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + + def reduce(self, contribute: bool) -> dict[str, float]: + """All-reduce sums and counts (call on every rank in the same key order); returns means.""" + keys = self.all_keys + device = self._device() + values = torch.tensor([float(self.sums.get(k, 0.0)) for k in keys] + [self.counts.get(k, 0.0) for k in keys], + dtype=torch.float64, + device=device) + if not contribute: + values.zero_() + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(values) + sums, counts = values[:len(keys)], values[len(keys):] + return {key: float(sums[i]) / float(counts[i]) for i, key in enumerate(keys) if float(counts[i]) > 0} + + +def summarize(means: dict[str, float], num_rungs: int, weights: dict[str, float]) -> dict[str, float]: + """Per-rung rel-L2 / cosine / MSE plus the weighted mean x0 rel-L2 ``score`` (lower is better).""" + out: dict[str, float] = {} + score_num = score_den = 0.0 + for modality in MODALITIES: + rel_all = [] + for rung in range(num_rungs): + p = f"{modality}.{rung}" + if f"{p}.ref_sq" not in means: + continue + rel = (means[f"{p}.err_sq"] / max(means[f"{p}.ref_sq"], 1e-30))**0.5 + vrel = (means[f"{p}.verr_sq"] / max(means[f"{p}.vref_sq"], 1e-30))**0.5 + cos = means[f"{p}.dot"] / max((means[f"{p}.s_sq"] * means[f"{p}.ref_sq"])**0.5, 1e-30) + out[f"x0_rel_l2/{modality}/rung{rung}"] = rel + out[f"v_rel_l2/{modality}/rung{rung}"] = vrel + out[f"x0_cos/{modality}/rung{rung}"] = cos + out[f"x0_mse/{modality}/rung{rung}"] = means[f"{p}.mse"] + rel_all.append(rel) + if rel_all: + out[f"x0_rel_l2/{modality}/mean"] = sum(rel_all) / len(rel_all) + score_num += weights[modality] * out[f"x0_rel_l2/{modality}/mean"] + score_den += weights[modality] + out["score"] = score_num / max(score_den, 1e-30) + for key, value in means.items(): + if key.startswith(("endpoint.", "keyframe.")): + out[key.replace(".", "/")] = value + return out + + +def keyframe_latents(prepared: Any, final_video: torch.Tensor) -> dict[str, tuple[torch.Tensor, torch.Tensor | None]]: + """Generated keyframe latents ``[C, H, W]`` (first; last for first_last_frame) and their reference image. + + OmniRef encodes keyframe references at high resolution with the target's + aspect ratio, so the reference latent is area-resized to the target grid: a + rough proxy, logged only. The guarded keyframe metric compares the student's + keyframe with the teacher's (``keyframe/vs_teacher``). + """ + import torch.nn.functional as F + + from fastvideo.pipelines.basic.minimax_h3.packing import unpatchify_video_tokens + + channels, frames, height, width = prepared.latent_shape + patch = (1, 2, 2) + cut = target_slices(prepared)["video"] + video = unpatchify_video_tokens(final_video[cut][None].float(), frames, height, width, channels, patch)[0] + images: list[torch.Tensor | None] = [] + offset = 0 + for ref in prepared.references: + if ref.media_type == "audio": + continue + count = ref.num_latent_frames * (ref.latent_height // 2) * (ref.latent_width // 2) + image = None + same_aspect = abs(ref.latent_height * width - ref.latent_width * height) <= 0.02 * ref.latent_width * height + if ref.media_type == "image" and same_aspect: + rows = prepared.reference_video_rows[offset:offset + count].to(video.device) + latent = unpatchify_video_tokens(rows[None].float(), 1, ref.latent_height, ref.latent_width, channels, + patch)[:, :, 0] + if (ref.latent_height, ref.latent_width) != (height, width): + latent = F.interpolate(latent, size=(height, width), mode="area") + image = latent[0] + images.append(image) + offset += count + pairs = {"first_frame": [("first", 0, 0)], "first_last_frame": [("first", 0, 0), ("last", -1, 1)]} + return { + name: (video[:, frame], images[ref] if ref < len(images) else None) + for name, frame, ref in pairs.get(prepared.spec.case, []) + } + + +def rel_l2(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12) + + +__all__ = ["RungStats", "keyframe_latents", "rel_l2", "step_packed", "summarize", "target_slices", "x0_pair"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py new file mode 100644 index 000000000..5cda222cc --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Layer-wise reconstruction for the FastH3 OmniRef QAD student, streamed per block. + +Supervised units R, per transformer block: the attention sub-layer (q/k/v/gate/out NVFP4 projections, QK +RMSNorm, RoPE and the sparse FP4 VSA-128 attention, as deployed) and the FFN sub-layer (NVFP4 fc_in/fc_out). +Each student unit receives the frozen teacher's input to that unit (stop-gradient) and matches the teacher +unit's output, so quantization error is corrected locally and cannot accumulate into a drift away from the +teacher. + +Locality makes this cheap: inside one teacher forward, right after teacher block ``b`` runs, student block +``b`` runs on the captured inputs and is backpropagated at once. There is no end-to-end student forward or +backward, no activation recompute, and memory holds one block's activations. + +Attention is sparse (VSA-128, the teacher's trained selection), not dense as in the paper's warmup: our +teacher was trained sparse and the student is deployed sparse, so dense attention would supervise a +configuration that never runs. + +The loss per unit is the relative squared error ``||s - t||^2 / ||t||^2`` (late FFN outputs reach ~1e6, so an +unnormalized sum would only train those), averaged over R. With a ``GroupedRelError`` it is computed per row group +(target video, target audio, conditioning) and weighted: over the whole packed sequence audio is ~1.4% of the rows +and massive text tokens dominate ``||t||^2``, so an ungrouped ratio barely trains audio. +""" +from __future__ import annotations + +from typing import Any + +import torch + + +def ensure_fsdp_root_initialized(module: torch.nn.Module) -> None: + """Make ``module`` its FSDP2 tree's root before any of its sharded blocks runs on its own. + + FSDP2 picks the root lazily, at the first forward of any sharded module. If a block runs first (as in + reconstruction, which calls student blocks directly), that block claims the root and the model's next full + forward fails with "FSDP state has already been lazily initialized". Idempotent; a no-op without FSDP2. + """ + get_state = getattr(module, "_get_fsdp_state", None) + if get_state is not None: + get_state()._lazy_init() + + +def _unwrap(module: torch.nn.Module) -> torch.nn.Module: + return getattr(module, "_checkpoint_wrapped_module", module) + + +GROUPS = ("video", "audio", "condition") # target video rows, target audio rows, text + reference rows + + +def packed_row_groups(layout: Any, sequence_length: int) -> torch.Tensor: + """Group id per packed row: 0 target video, 1 target audio, 2 conditioning (text, references).""" + groups = torch.full((sequence_length, ), 2, dtype=torch.long) + groups[layout.video_indices[layout.num_condition_video_rows:]] = 0 + groups[layout.audio_indices[layout.num_condition_audio_rows:]] = 1 + return groups + + +def local_row_groups(groups: torch.Tensor, device: torch.device) -> torch.Tensor: + """This sequence-parallel rank's rows of ``groups`` as the transformer shards them; padding is -1.""" + from fastvideo.distributed import get_sp_world_size, model_parallel_is_initialized + from fastvideo.distributed.communication_op import sequence_model_parallel_shard + + shifted = (groups + 1).to(device) + if model_parallel_is_initialized() and get_sp_world_size() > 1: + shifted, _ = sequence_model_parallel_shard(shifted, dim=0) # pads with 0 -> -1 after the shift + return shifted - 1 + + +class GroupedRelError: + """Weighted mean over row groups of ``||s - t||^2 / ||t||^2`` with sequence-parallel-global sums. + + Each rank backpropagates ``sp_size * w_g * num_local / den_global``. Summed over ranks, ``w_g * num_local / + den_global`` is the global ratio, so a group living on one rank (e.g. target audio) is weighted exactly as + if unsharded. The ``sp_size`` factor matches the x0 output loss: FSDP averages gradients over every rank, + SP ranks included, and the output loss is identical on every SP rank, where the transformer's output + all-gather backward (a SUM reduce-scatter) gives each rank ``sp_size`` times its share of the gradient. + Reconstruction has no gather, so without the factor the averaged gradient would be ``1 / sp_size`` of the + global ratio's. Groups absent from the sequence are dropped from the weight normalization. + """ + + def __init__(self, groups: torch.Tensor, weights: dict[str, float]) -> None: + self.groups = groups + self.weights = torch.tensor([float(weights.get(name, 0.0)) for name in GROUPS], device=groups.device) + + def __call__(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(local loss to backprop, global per-group relative error [3], detached)``. + + The local loss is this rank's ``sp_size``-scaled share (see the class docstring); the detached global + per-group errors are what to log. + """ + from fastvideo.distributed import get_sp_group, get_sp_world_size, model_parallel_is_initialized + + teacher = teacher.float().reshape(-1, teacher.shape[-1]) + err = (student.float().reshape(-1, teacher.shape[-1]) - teacher).pow(2).sum(-1) + ref = teacher.pow(2).sum(-1) + one_hot = (self.groups[:, None] == torch.arange(len(GROUPS), device=err.device)[None, :]).float() + num = err @ one_hot # [3], differentiable + totals = torch.stack((num.detach(), ref @ one_hot)) # [2, 3] + sp_size = get_sp_world_size() if model_parallel_is_initialized() else 1 + if sp_size > 1: + totals = get_sp_group().all_reduce(totals) + present = totals[1] > 0 + den = totals[1].clamp_min(1e-30) + weights = torch.where(present, self.weights, torch.zeros_like(self.weights)) + norm = weights.sum().clamp_min(1e-30) + loss = sp_size * (weights * num / den).sum() / norm + return loss, torch.where(present, totals[0] / den, torch.full_like(den, float("nan"))) + + +class _Replace: + """Forward pre-hook replacing a sub-layer's first input with the teacher's, plus an output capture.""" + + def __init__(self) -> None: + self.value: torch.Tensor | None = None + self.output: torch.Tensor | None = None + + def pre(self, module: torch.nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...]: + return (self.value, *args[1:]) + + def post(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self.output = output + + +class _Capture: + + def __init__(self) -> None: + self.input: torch.Tensor | None = None + self.output: torch.Tensor | None = None + + def __call__(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self.input, self.output = args[0], output + + +class BlockReconstruction: + """Stream reconstruction losses for every block during one teacher forward. + + ``scale`` multiplies each block's loss before its backward (e.g. ``1 / (units x rungs x accum)``). + Use as a context manager around the teacher's ``forward_rung``, with grad disabled for the teacher. + """ + + def __init__(self, + teacher: torch.nn.Module, + student: torch.nn.Module, + scale: float, + criterion: GroupedRelError | None = None, + measure_only: bool = False) -> None: + ensure_fsdp_root_initialized(teacher) + ensure_fsdp_root_initialized(student) + self.teacher_blocks = list(teacher.transformer_blocks) + self.student_blocks = list(student.transformer_blocks) + if len(self.teacher_blocks) != len(self.student_blocks): + raise ValueError("teacher and student block counts differ") + self.scale = scale + self.criterion = criterion + self.measure_only = measure_only + self.attn_rel: list[torch.Tensor] = [] + self.ff_rel: list[torch.Tensor] = [] + self.group_rel: list[torch.Tensor] = [] # [3] per unit, global, detached + self._handles: list[Any] = [] + + def _weighted(self, rel: torch.Tensor) -> torch.Tensor: + """Group-weighted mean of per-group relative errors ``[..., 3]`` (NaN = absent group, dropped).""" + assert self.criterion is not None + weights = self.criterion.weights.expand_as(rel) * (~rel.isnan()) + return (weights * rel.nan_to_num()).sum(-1) / weights.sum(-1).clamp_min(1e-30) + + def _unit_loss(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(loss to backprop, sequence-parallel-global relative error to log, detached)``.""" + if self.criterion is None: + teacher_f = teacher.float() + loss = (student.float() - teacher_f).pow(2).mean() / teacher_f.pow(2).mean().clamp_min(1e-20) + return loss, loss.detach() + loss, per_group = self.criterion(student, teacher) + self.group_rel.append(per_group) + return loss, self._weighted(per_group) + + def __enter__(self) -> BlockReconstruction: + for index, block in enumerate(self.teacher_blocks): + inner = _unwrap(block) + attn, ff = _Capture(), _Capture() + self._handles.append(inner.attn.register_forward_hook(attn)) + self._handles.append(inner.ff.register_forward_hook(ff)) + self._handles.append(block.register_forward_hook(self._after_block(index, attn, ff))) + return self + + def __exit__(self, *exc: Any) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + def _after_block(self, index: int, attn: _Capture, ff: _Capture): + + def hook(module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self._student_block(index, args, attn, ff) + attn.input = attn.output = ff.input = ff.output = None + + return hook + + def _student_block(self, index: int, block_args: tuple[Any, ...], attn: _Capture, ff: _Capture) -> None: + block = self.student_blocks[index] + inner = _unwrap(block) + replace_attn, replace_ff = _Replace(), _Replace() + replace_attn.value, replace_ff.value = attn.input, ff.input + handles = [ + inner.attn.register_forward_pre_hook(replace_attn.pre), + inner.attn.register_forward_hook(replace_attn.post), + inner.ff.register_forward_pre_hook(replace_ff.pre), + inner.ff.register_forward_hook(replace_ff.post) + ] + try: + with torch.set_grad_enabled(not self.measure_only): + # The block's own hidden input only feeds norm1 and the residual, whose results are unused; the + # teacher attention input has the right shape. The call goes through the block, and the block + # input and output join the graph (zero weight), so FSDP's pre-/post-backward hooks re-gather the + # block's weights for backward and reduce-scatter its gradients. + hidden = attn.input.detach().requires_grad_(not self.measure_only) + block_out = block(hidden, *block_args[1:]) + attn_loss, attn_rel = self._unit_loss(replace_attn.output, attn.output) + ff_loss, ff_rel = self._unit_loss(replace_ff.output, ff.output) + if not self.measure_only: + ((attn_loss + ff_loss) * self.scale + 0.0 * block_out.float().sum()).backward() + finally: + for handle in handles: + handle.remove() + self.attn_rel.append(attn_rel) + self.ff_rel.append(ff_rel) + + def loss(self) -> torch.Tensor: + """Mean (group-weighted) relative error over all units of this forward, sequence-parallel-global, detached.""" + if self.criterion is not None and self.group_rel: + return self._weighted(torch.stack(self.group_rel)).mean() # [units, 3], NaN for absent groups + return (torch.stack(self.attn_rel).sum() + torch.stack(self.ff_rel).sum()) / (2 * len(self.attn_rel)) + + def metrics(self, prefix: str) -> dict[str, torch.Tensor]: + attn, ff = torch.stack(self.attn_rel), torch.stack(self.ff_rel) + third = max(1, len(attn) // 3) + groups = {} + if self.group_rel: + per_group = torch.stack(self.group_rel).nanmean(dim=0) + groups = {f"{prefix}/{name}_rel": per_group[i] for i, name in enumerate(GROUPS)} + return { + **groups, f"{prefix}/attn_rel": attn.mean(), + f"{prefix}/ff_rel": ff.mean(), + f"{prefix}/attn_rel_first_third": attn[:third].mean(), + f"{prefix}/attn_rel_last_third": attn[-third:].mean(), + f"{prefix}/ff_rel_last_third": ff[-third:].mean() + } + + +__all__ = ["BlockReconstruction", "GROUPS", "GroupedRelError", "local_row_groups", "packed_row_groups"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py new file mode 100644 index 000000000..f98f3c2fe --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reuse a QAD run's step-0 evaluation (PTQ baseline, seed+1 noise floor, reconstruction normalizers). + +A Stage-2 run from a fixed checkpoint spends most of its first ~1.75 h on step 0: the full held-out evaluation, +the same evaluation at seed+1 for the noise floor, and the reconstruction calibration. All of it is a deterministic +function of the teacher and student checkpoints (``init_from``, plus ``init_student_dcp``), the calibration table +(or the provisional-calibration rows), the quantization plan, the evaluation set and the score weights, so a rerun +with the same inputs can load it. The cache is keyed by a fingerprint of exactly those inputs; any change (another +init checkpoint, eval rows, tolerances, modality weights, ...) is a mismatch and the run refuses to start. + +No torch import, so the fingerprint can be computed offline from a run's ``cfg.yaml``. +""" +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any + +STEP0_FORMAT = 2 +METHOD_KEYS = ("init_student_dcp", "max_eval_rows", "max_rollout_rows", "noise_floor_eval", "min_eligibility_tolerance", + "eligibility_keys", "fidelity_tolerance", "objective", "recon_normalize", "recon_norm_rows", + "recon_group_weights", "rungs_per_row") +DATA_KEYS = ("manifests", "cases", "max_frames", "eval_manifest", "eval_resolutions", "eval_per_group", + "eval_max_frames", "eval_seed", "width", "height", "num_frames") +# The training-row plan; it only reaches step 0 through provisional FFN calibration (no amax_json). +PROVISIONAL_DATA_KEYS = ("seed", "per_group", "max_train_rows", "train_manifest") +QAD_KEYS = ("linears", "fp4_attention") +# Logged per evaluation but not part of the step-0 result. +VOLATILE = ("eval/seconds", "eval/score_vs_ptq") + + +def _plain(value: Any) -> Any: + """OmegaConf / mapping / sequence -> JSON-native containers.""" + if hasattr(value, "items"): + return {str(k): _plain(v) for k, v in value.items()} + if isinstance(value, list | tuple) or (hasattr(value, "__iter__") and not isinstance(value, str | bytes)): + return [_plain(v) for v in value] + return value + + +def step0_fingerprint_inputs(method_config: Any, data_config: Any, qad_config: Any, amax_sha256: str, + num_eval_rows: int, *, student_init_from: str, teacher_init_from: str) -> dict[str, Any]: + """Everything step 0 depends on. The init paths stand for their weights and ``fastvideo_inference.json``.""" + method, data, qad = _plain(method_config), _plain(data_config), _plain(qad_config) + linears = qad.get("linears") or {} + provisional = int(linears.get("provisional_calibration_rows") or 0) > 0 and not linears.get("amax_json") + return { + "format": STEP0_FORMAT, + "models": { + "student_init_from": str(student_init_from), + "teacher_init_from": str(teacher_init_from) + }, + "method": { + key: method.get(key) + for key in METHOD_KEYS + }, + # ``summarize`` weights the score with these; best_score and the guards compare against it. + "score_weights": { + modality: float(method.get(f"{modality}_weight", 1.0)) + for modality in ("video", "audio") + }, + "data": { + key: data.get(key) + for key in DATA_KEYS + }, + "provisional_plan": { + key: data.get(key) + for key in PROVISIONAL_DATA_KEYS + } if provisional else None, + "qad": { + key: qad.get(key) + for key in QAD_KEYS + }, + "amax_sha256": amax_sha256, + "num_eval_rows": int(num_eval_rows), + } + + +def fingerprint(inputs: dict[str, Any]) -> str: + return hashlib.sha256(json.dumps(inputs, sort_keys=True).encode()).hexdigest() + + +def save_step0(path: str | Path, inputs: dict[str, Any], metrics: dict[str, float]) -> None: + payload = { + "format": STEP0_FORMAT, + "fingerprint": fingerprint(inputs), + "inputs": inputs, + "metrics": { + key: float(value) + for key, value in metrics.items() if key not in VOLATILE + }, + } + Path(path).parent.mkdir(parents=True, exist_ok=True) + tmp = Path(f"{path}.tmp") + tmp.write_text(json.dumps(payload, indent=1, sort_keys=True)) + tmp.replace(path) + + +def load_step0(path: str | Path, inputs: dict[str, Any]) -> dict[str, float] | None: + """The cached step-0 metrics, None if there is no cache; raises on a fingerprint mismatch.""" + path = Path(path) + if not path.exists(): + return None + payload = json.loads(path.read_text()) + expected = fingerprint(inputs) + if payload.get("format") != STEP0_FORMAT or payload.get("fingerprint") != expected: + cached = payload.get("inputs", {}) + diff = sorted(k for k in set(cached) | set(inputs) if cached.get(k) != inputs.get(k)) + raise ValueError(f"step-0 cache {path} was computed for other inputs (differs in {diff}); " + f"delete it or point step0_cache elsewhere") + return {key: float(value) for key, value in payload["metrics"].items()} + + +def restore_step0(metrics: dict[str, float], num_rungs: int, recon: bool, noise_floor: bool, + min_tolerance: float) -> tuple[dict[str, float], list[float] | None, dict[str, float]]: + """(eval summary, reconstruction normalizers, eligibility tolerances) from cached step-0 metrics.""" + summary = {key[len("eval/"):]: value for key, value in metrics.items() if key.startswith("eval/")} + if "score" not in summary: + raise ValueError("step-0 cache has no eval/score") + recon_norm = None + if recon: + recon_norm = [max(metrics[f"init/recon_step0/rung{rung}"], 1e-12) for rung in range(num_rungs)] + tolerances: dict[str, float] = {} + if noise_floor: + prefix = "init/noise_sigma/" + tolerances = { + key[len(prefix):]: max(2.0 * value, min_tolerance) + for key, value in metrics.items() if key.startswith(prefix) and math.isfinite(value) + } + if not tolerances: + raise ValueError("noise_floor_eval is on but the step-0 cache has no noise-floor sigmas") + return summary, recon_norm, tolerances diff --git a/fastvideo/train/models/minimax_h3/omniref_data.py b/fastvideo/train/models/minimax_h3/omniref_data.py new file mode 100644 index 000000000..7cae77c19 --- /dev/null +++ b/fastvideo/train/models/minimax_h3/omniref_data.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniRef (FastH3 Ref2VA) rows for per-forward distillation of the PDD student. + +A row is one OmniRef training parquet row (precomputed Qwen3-VL presentation, +clean reference latent rows and their geometry); nothing is re-encoded. The +initial packed state of a (row, seed) pair is built exactly as +``generate_omniref_latents.py`` and ``MiniMaxH3LatentPreparationStage`` build +it (reference rows noise-augmented to 0.999, target video/audio noise from the +same CPU generator, PDD start scaling), so trajectories, teacher latents and +Stage A's teacher-forced error are all reproducible from the row id and seed. + +Row selection reuses the generator's deterministic plan (per (case, +resolution): seeded shuffle, then interleave), so the first rows of each group +are the rows the OmniRef latent generator and Stage A calibration used. +Held-out sources are excluded at both resolutions. +""" +from __future__ import annotations + +import hashlib +import json +import random +import re +from collections import defaultdict +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +OMNIREF_CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot", + "full_scene_audio_image_reference") +RESOLUTIONS = ("480p", "768p") +_BUCKET_FRAMES = re.compile(r"bucket=\d+x\d+-(\d+)f") +ROW_COLUMNS = [ + "id", "vae_latent_shape", "audio_latent_shape", "text_embedding_bytes", "text_embedding_shape", + "text_token_tags_bytes", "text_token_tags_shape", "reference_video_latent_rows_bytes", + "reference_video_latent_rows_shape", "reference_audio_latent_rows_bytes", "reference_audio_latent_rows_shape", + "reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights", + "reference_latent_widths", "reference_num_audio_latents", "width", "height", "num_frames" +] + + +@dataclass(frozen=True) +class RowSpec: + id: str + source: str + case: str + resolution: str + parquet: str + seed: int + + +def clip_seed(base_seed: int, clip_id: str) -> int: + digest = hashlib.sha256(f"{base_seed}:{clip_id}".encode()).digest() + return int.from_bytes(digest[:4], "little") & 0x7FFFFFFF + + +def load_manifest_groups(paths: list[str], cases: tuple[str, ...], + max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]: + """Manifest entries grouped by (case, resolution); several manifests (v4 + v5 additions) merge. + + An entry repeated across manifests is kept once per (id, resolution): the same source at 480p and 768p is + two rows. + """ + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + seen: set[tuple[str, str]] = set() + for path in paths: + with open(path) as handle: + for line in handle: + entry = json.loads(line) + match = _BUCKET_FRAMES.search(entry["parquet"]) + key = (entry["id"], entry["resolution"]) + if (match is None or entry["case"] not in cases or entry["resolution"] not in RESOLUTIONS or key in seen + or int(match.group(1)) > max_frames): + continue + seen.add(key) + groups[(entry["case"], entry["resolution"])].append(entry) + return groups + + +def limit_frames(groups: dict[tuple[str, str], list[dict[str, Any]]], + max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]: + """``groups`` without entries whose bucket is longer than ``max_frames`` (0 = no limit); a new mapping.""" + if not max_frames: + return dict(groups) + + def frames(entry: dict[str, Any]) -> int: + match = _BUCKET_FRAMES.search(entry["parquet"]) + return int(match.group(1)) if match else max_frames + 1 + + return {key: [entry for entry in entries if frames(entry) <= max_frames] for key, entries in groups.items()} + + +def select_rows(groups: dict[tuple[str, str], list[dict[str, Any]]], + seed: int, + per_group: dict[str, int], + exclude_sources: frozenset[str] = frozenset()) -> list[RowSpec]: + """The generator's plan: per group a seeded shuffle, minus excluded sources, then interleaved.""" + picked = [] + for key in sorted(groups): + entries = sorted(groups[key], key=lambda entry: entry["id"]) + random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries) + entries = [entry for entry in entries if entry["id"].split(":", 1)[-1] not in exclude_sources] + picked.append(entries[:int(per_group.get(key[1], 0))]) + plan = [] + for index in range(max((len(entries) for entries in picked), default=0)): + for entries in picked: + if index < len(entries): + entry = entries[index] + source = entry["id"].split(":", 1)[-1] + clip_id = f"omniref-{entry['case']}-{source}-{entry['resolution']}" + plan.append( + RowSpec(clip_id, source, entry["case"], entry["resolution"], entry["parquet"], + clip_seed(seed, clip_id))) + return plan + + +def load_eval_rows(manifest_json: str, + cases: tuple[str, ...], + resolutions: tuple[str, ...], + per_group: int | dict[str, int], + max_frames: int = 0) -> list[RowSpec]: + """Held-out rows from a ``generate_omniref_latents.py --eval-dir`` manifest (ids, parquets and seeds).""" + items = json.loads(Path(manifest_json).read_text()) + counts: dict[tuple[str, str], int] = defaultdict(int) + rows = [] + for item in sorted(items, key=lambda item: item["id"]): + key = (item["case"], item["resolution"]) + limit = int(per_group.get(item["resolution"], 0)) if isinstance(per_group, dict) else int(per_group) + if item["case"] not in cases or item["resolution"] not in resolutions or counts[key] >= limit: + continue + if max_frames and int(item.get("num_frames", 0)) > max_frames: + continue + counts[key] += 1 + rows.append( + RowSpec(item["id"], item["source"], item["case"], item["resolution"], item["parquet"], int(item["seed"]))) + return rows + + +def top_up_eval_rows(eval_rows: list[RowSpec], groups: dict[tuple[str, str], list[dict[str, Any]]], seed: int, + per_group: dict[str, int]) -> list[RowSpec]: + """Fill each (case, resolution) up to ``per_group`` with fresh held-out rows (call before the training plan). + + Extra rows come from the generator's plan rule at ``seed``, skipping sources already in ``eval_rows``; + the training plan must then exclude every eval source. Result is interleaved across groups. + """ + have: dict[tuple[str, str], int] = defaultdict(int) + for row in eval_rows: + have[(row.case, row.resolution)] += 1 + taken = frozenset(row.source for row in eval_rows) + extra = [] + for key in sorted(groups): + need = int(per_group.get(key[1], 0)) - have[key] + if need > 0: + extra += select_rows({key: groups[key]}, seed, {key[1]: need}, taken) + return interleave_by_group(eval_rows + extra) + + +def interleave_by_group(rows: list[RowSpec]) -> list[RowSpec]: + """Round-robin over (case, resolution) groups, so any prefix covers every case.""" + buckets: dict[tuple[str, str], list[RowSpec]] = defaultdict(list) + for row in rows: + buckets[(row.case, row.resolution)].append(row) + ordered = [buckets[key] for key in sorted(buckets)] + return [ + bucket[i] for i in range(max((len(b) for b in ordered), default=0)) for bucket in ordered if i < len(bucket) + ] + + +class RowStream: + """Infinite, deterministic per-data-parallel-rank row stream (every SP rank of a group sees the same row).""" + + def __init__(self, plan: list[RowSpec], *, dp_rank: int, dp_size: int, seed: int) -> None: + if not plan: + raise ValueError("the training row plan is empty") + self.plan = plan + self.dp_rank = dp_rank + self.dp_size = dp_size + self.seed = seed + self.position = 0 + + def __iter__(self) -> Iterator[dict[str, Any]]: + while True: + epoch, offset = divmod(self.position * self.dp_size + self.dp_rank, len(self.plan)) + order = list(range(len(self.plan))) + random.Random(f"{self.seed}:epoch{epoch}").shuffle(order) + self.position += 1 + yield {"row": self.plan[order[offset]], "epoch": epoch} + + def state_dict(self) -> dict[str, int]: + return {"position": self.position} + + def load_state_dict(self, state: dict[str, int]) -> None: + self.position = int(state["position"]) + + +# --------------------------------------------------------------------------- one row -> initial packed state +def _tensor(row: dict[str, Any], name: str) -> torch.Tensor: + shape = [int(value) for value in row[f"{name}_shape"]] + array = np.frombuffer(row[f"{name}_bytes"], dtype=np.float32) + if array.size != int(np.prod(shape, dtype=np.int64)): + raise ValueError(f"{name}: {array.size} values do not fill shape {shape}") + return torch.from_numpy(array.reshape(shape).copy()) + + +def read_row(parquet: str) -> dict[str, Any]: + import pyarrow.parquet as pq + + rows = pq.read_table(parquet, columns=ROW_COLUMNS).to_pylist() + if len(rows) != 1: + raise ValueError(f"{parquet}: expected one row, got {len(rows)}") + return rows[0] + + +def prepared_references(row: dict[str, Any]) -> list[Any]: + from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference + + fields = ("reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights", + "reference_latent_widths", "reference_num_audio_latents") + return [ + MiniMaxH3PreparedReference(media_type=media, + has_audio=bool(audio), + num_latent_frames=int(frames), + latent_height=int(height), + latent_width=int(width), + num_audio_latents=int(count)) + for media, audio, frames, height, width, count in zip(*(row[name] for name in fields), strict=True) + ] + + +@dataclass +class PreparedRow: + """Everything one PDD trajectory of a row needs; ``video``/``audio`` are the packed initial rows (fp32).""" + spec: RowSpec + layout: Any + prompt_embeds: torch.Tensor + video: torch.Tensor + audio: torch.Tensor + latent_shape: tuple[int, int, int, int] # (channels, frames, height, width) of the target video + references: list[Any] + reference_video_rows: torch.Tensor # clean, CPU + + +def prepare_row(spec: RowSpec, row: dict[str, Any], *, patch_size: tuple[int, int, int], latent_channels: int, + audio_channels: int, pdd_start_sigmas: tuple[float, float], device: torch.device) -> PreparedRow: + """Initial packed state, bit-identical to ``OmniRefLatentGenerator._batch`` + ``scale_pdd_initial_noise``.""" + from diffusers.utils.torch_utils import randn_tensor + + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_NOISE_AUG, + audio_latent_num_frames, build_ref2va_packed_sequence, + keyframe_condition_noise, patchify_video_latents, + video_latent_num_frames) + + num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"]) + frames, lat_h, lat_w = video_latent_num_frames(num_frames), height // 16, width // 16 + expected = tuple(int(value) for value in row["vae_latent_shape"][1:]) + if (frames, lat_h, lat_w) != expected or int(row["vae_latent_shape"][0]) != latent_channels: + raise ValueError(f"{spec.id}: target geometry {(frames, lat_h, lat_w)} disagrees with {expected}") + num_audio = audio_latent_num_frames(num_frames) + references = prepared_references(row) + tags = _tensor(row, "text_token_tags").round().to(torch.long) + embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None] + layout = build_ref2va_packed_sequence(tags, references, frames, lat_h, lat_w, num_audio, patch_size) + + generator = torch.Generator("cpu").manual_seed(spec.seed) + shapes = tuple( + (ref.num_latent_frames, ref.latent_height, ref.latent_width) for ref in references if ref.media_type != "audio") + noise = keyframe_condition_noise(shapes, patch_size, latent_channels, generator=generator, device=device) + clean = _tensor(row, "reference_video_latent_rows") + condition_video = MiniMaxH3Scheduler().scale_noise(clean.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) + video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w), + generator=generator, + device=device, + dtype=torch.float32) + audio_rows = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels), + generator=generator, + device=device, + dtype=torch.float32) + condition_audio = _tensor(row, "reference_audio_latent_rows").to(device) + video = torch.cat((condition_video, patchify_video_latents(video_noise, patch_size))) + audio = torch.cat((condition_audio, audio_rows)) + if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"{spec.id}: packed rows do not match the Ref2VA layout") + video[layout.num_condition_video_rows:] *= pdd_start_sigmas[0] + audio[layout.num_condition_audio_rows:] *= pdd_start_sigmas[1] + return PreparedRow(spec, layout, embeds, video, audio, (latent_channels, frames, lat_h, lat_w), references, clean) + + +__all__ = [ + "OMNIREF_CASES", "PreparedRow", "RowSpec", "RowStream", "interleave_by_group", "load_eval_rows", + "load_manifest_groups", "top_up_eval_rows", "prepare_row", "read_row", "select_rows" +] diff --git a/fastvideo/train/models/minimax_h3/omniref_pdd.py b/fastvideo/train/models/minimax_h3/omniref_pdd.py new file mode 100644 index 000000000..ca5bb436b --- /dev/null +++ b/fastvideo/train/models/minimax_h3/omniref_pdd.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastH3 OmniRef PDD transformer as a per-forward (rung) model for distillation. + +One instance wraps ``transformer_ref`` of a FastH3 OmniRef PDD export, loaded +through the training FSDP loader with VSA-H3 attention at the trained tile +size. It exposes the eight PDD forwards of the deployed sampler one at a +time: ``rung_scope`` installs exactly what ``MiniMaxH3DenoisingStage`` sets +for forward ``k`` (VSA-H3 metadata with the reference keep rate, the fused +PDD output heads, per-row timesteps), ``forward_rung`` runs the transformer, +and ``step`` applies the scheduler's Euler step. Teacher and student load +the same way, so a student whose quantization is disabled is the teacher. + +The student may carry NVFP4 QAD linears (``qad.linears``) and the sparse-FP4 +attention emulator (``qad.fp4_attention``); see ``fastvideo.layers. +quantization.nvfp4_qad`` and ``fastvideo.attention.backends.fp4_vsa_qat``. +""" +from __future__ import annotations + +import contextlib +import json +import re +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.activation_checkpoint import apply_activation_checkpointing, resolve_checkpointing_type +from fastvideo.train.utils.moduleloader import load_module_from_path + +if TYPE_CHECKING: + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + from fastvideo.train.utils.training_config import TrainingConfig + +logger = init_logger(__name__) + +# Block linears and norms; AdaLN, embeddings, the token refiner and the output heads stay frozen. +DEFAULT_TRAINABLE = (r"transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out|to_gate_compress|norm_q|norm_k)" + r"|norm1|norm2|ff\.(fc_in|fc_out))\.") +INFERENCE_FILE = "fastvideo_inference.json" + + +@dataclass(frozen=True) +class PDDContract: + """The checkpoint's deployed sampling contract (``fastvideo_inference.json``).""" + pdd_step_indices: tuple[int, ...] + video_shift: float + audio_shift: float + vsa_sparsity: float + vsa_tile_size: int + vsa_ref_keep_rate: float | None + + @classmethod + def read(cls, model_path: str) -> PDDContract: + data = json.loads((Path(model_path) / INFERENCE_FILE).read_text()) + if data.get("model_type") != "ref2va" or "pdd_step_indices" not in data: + raise ValueError(f"{model_path} is not a FastH3 Ref2VA PDD export") + return cls(tuple(int(i) for i in data["pdd_step_indices"]), float(data["video_scheduler_shift"]), + float(data["audio_scheduler_shift"]), float(data["vsa_sparsity"]), int(data["vsa_tile_size"]), + data.get("vsa_ref_keep_rate")) + + +class MiniMaxH3OmniRefPDDModel(ModelBase): + """Teacher or QAD student for per-forward regression on the OmniRef PDD ladder.""" + + _transformer_cls_name = "MiniMaxH3Transformer3DModel" + _transformer_module_type = "transformer_ref" + + def __init__(self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + enable_gradient_checkpointing_type: str | None = None, + trainable_patterns: str = DEFAULT_TRAINABLE, + qad: dict[str, Any] | None = None, + data: dict[str, Any] | None = None, + frozen_bf16: bool = False) -> None: + super().__init__(trainable=trainable, attention_backend="VIDEO_SPARSE_ATTN_H3") + # Keep a frozen role's masters in bf16 (the checkpoint's dtype) instead of an fp32 upcast: compute is bf16 either + # way, so outputs are bit-identical, and the role needs half the memory. + self._frozen_bf16 = bool(frozen_bf16) + if training_config.pipeline_config is None: + raise ValueError("MiniMaxH3OmniRefPDDModel requires a resolved MiniMax-H3 pipeline config") + if str(training_config.dit_precision) != "fp32" and trainable: + raise ValueError("QAD needs fp32 master weights (training.dit_precision: fp32); BF16 masters " + "round away updates at QAD learning rates") + self.contract = self._read_contract(init_from) + self._check_tile() + # Inference keeps proj_in/out, audio_proj_in/out and the time embedder in FP32; FSDP mixed precision would + # compute them in BF16 and the harness would no longer be the deployed operating point (E1: -7.6% v_rel at + # forward 0). They are frozen here, so they stay FP32 and replicated outside FSDP. + training_config.pipeline_config.dit_config.uniform_parameter_dtype = False # type: ignore[attr-defined] + training_config.pipeline_config.dit_config.replicate_fp32_modules = True # type: ignore[attr-defined] + self._init_from = str(init_from) + self.training_config = training_config + self.qad_config = dict(qad or {}) + self.data_config = dict(data or {}) + with self._load_precision(trainable): + self.transformer = load_module_from_path(model_path=self._init_from, + module_type=self._transformer_module_type, + training_config=training_config, + override_transformer_cls_name=self._transformer_cls_name, + attention_backend=self.attention_backend) + checkpointing = resolve_checkpointing_type(enable_gradient_checkpointing_type, training_config) + if trainable and checkpointing: + self.transformer = apply_activation_checkpointing(self.transformer, checkpointing_type=checkpointing) + self.num_trainable = self._apply_trainable(trainable_patterns if trainable else None) + self._check_fp32_modules_frozen() + self.patch_size = tuple(int(v) for v in self.transformer.patch_size) + self._plan: Any = None + self._vsa_builder: Any = None + self.qad_linears: dict[str, Any] = {} + self.fp4_numerics: Any = None + self.provisional_rows = 0 + self.amax_sha256 = "" + self.eval_rows: list[Any] = [] + self.train_rows: list[Any] = [] + if self.qad_config.get("enabled", False): + self._install_qad() + self.dataloader: Any = None + self.start_step = 0 + + # ------------------------------------------------------------------ setup + def _load_precision(self, trainable: bool) -> contextlib.AbstractContextManager[None]: + """Context for the transformer load: a frozen role may keep bf16 masters (``frozen_bf16``).""" + if trainable or not self._frozen_bf16: + return contextlib.nullcontext() + return self._bf16_load() + + @contextlib.contextmanager + def _bf16_load(self) -> Iterator[None]: + config = self.training_config + pipeline_config = config.pipeline_config + saved = (config.dit_precision, pipeline_config.dit_precision) + config.dit_precision = pipeline_config.dit_precision = "bf16" + try: + yield + finally: + config.dit_precision, pipeline_config.dit_precision = saved + + def _read_contract(self, model_path: str) -> Any: + return PDDContract.read(model_path) + + def _check_tile(self) -> None: + if self.contract.vsa_tile_size != 128: + raise ValueError(f"the FP4 VSA emulator models tile 128; this export uses {self.contract.vsa_tile_size}") + + def load_row(self, spec: Any) -> PreparedRow: + from fastvideo.train.models.minimax_h3.omniref_data import read_row + return self.prepare(spec, read_row(spec.parquet)) + + def _apply_trainable(self, pattern: str | None) -> int: + regex = re.compile(pattern) if pattern else None + count = 0 + for name, parameter in self.transformer.named_parameters(): + clean = name.replace("_checkpoint_wrapped_module.", "") + wanted = regex is not None and regex.search(clean) is not None + parameter.requires_grad_(wanted) + count += parameter.numel() if wanted else 0 + self.transformer.train(regex is not None) + if regex is not None and count == 0: + raise ValueError(f"trainable_patterns {pattern!r} matched no parameters") + return count + + def _check_fp32_modules_frozen(self) -> None: + """FP32 boundary modules are replicated outside FSDP: training them would desynchronize the ranks.""" + selector = getattr(self.transformer, "_get_parameter_dtype", None) + for name, parameter in self.transformer.named_parameters(): + clean = name.replace("_checkpoint_wrapped_module.", "") + if callable(selector) and selector(clean, torch.bfloat16) == torch.float32: + if parameter.requires_grad: + raise ValueError(f"{clean} is an FP32 replicated module and must stay frozen") + if parameter.dtype != torch.float32: + raise ValueError(f"{clean} should be FP32 (as inference), got {parameter.dtype}") + + def _install_qad(self) -> None: + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table + + linears = dict(self.qad_config.get("linears") or {}) + if linears.get("enabled", True): + amax_path = linears.get("amax_json") + table = load_amax_table(amax_path) if amax_path else None + if amax_path: + self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"), + self.training_config.checkpoint.output_dir) + plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)), + quantize_gate=bool(linears.get("gate", True)), + quantize_ffn=bool(linears.get("ffn", True)), + skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ()))) + # Smoke tests only: without Stage A's table, max-calibrate the FFN scales on a few rows at step 0. + self.provisional_rows = int(linears.get("provisional_calibration_rows", 0)) + if table is None and self.provisional_rows <= 0: + raise ValueError("qad.linears.amax_json (Stage A calibration) is required") + self.qad_linears = install_nvfp4_qad(self.transformer, table, plan, allow_uncalibrated=table is None) + attention = dict(self.qad_config.get("fp4_attention") or {}) + if attention.get("enabled", True): + floor = attention.get("first_block_max_floor") + self.fp4_numerics = FP4AttentionNumerics(quantize=True, + two_level_p=bool(attention.get("two_level_p", False)), + smooth_k=bool(attention.get("smooth_k", False)), + first_block_max_floor=None if floor is None else float(floor)) + elif self._trainable: + # Tile 128 has no other grad-capable kernel: train through the emulator with exact softmax. + self.fp4_numerics = FP4AttentionNumerics(quantize=False) + install_fp4_vsa_attention(self.transformer, self.fp4_numerics) + logger.info("OmniRef QAD student: %d NVFP4 QAD linears, attention %s, %.2fB trainable parameters", + len(self.qad_linears), self.fp4_numerics, self.num_trainable / 1e9) + + @contextlib.contextmanager + def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]: + """Temporarily enable/disable the student's fake quantization (all off = the teacher's numerics).""" + from fastvideo.attention.backends.fp4_vsa_qat import install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled + + set_nvfp4_qad_enabled(self.qad_linears, linears) + if not attention: + install_fp4_vsa_attention(self.transformer, None) + try: + yield + finally: + set_nvfp4_qad_enabled(self.qad_linears, True) + install_fp4_vsa_attention(self.transformer, self.fp4_numerics) + + def init_preprocessors(self, training_config: TrainingConfig) -> None: + from fastvideo.train.models.minimax_h3.omniref_data import (OMNIREF_CASES, RowStream, interleave_by_group, + limit_frames, load_eval_rows, load_manifest_groups, + select_rows, top_up_eval_rows) + cfg = self.data_config + cases = tuple(cfg.get("cases", OMNIREF_CASES)) + resolutions = tuple(cfg.get("eval_resolutions", ("480p", ))) + eval_per_group = cfg.get("eval_per_group", 2) + eval_max_frames = int(cfg.get("eval_max_frames", 0)) + eval_rows = load_eval_rows(cfg["eval_manifest"], cases, resolutions, eval_per_group, eval_max_frames) + groups = load_manifest_groups(list(cfg["manifests"]), cases, int(cfg.get("max_frames", 243))) + missing = sorted(set(cases) - {case for case, _ in groups}) + if missing: + logger.warning("OmniRef cases with no rows in the given manifests: %s", missing) + if isinstance(eval_per_group, dict): + # More held-out rows than the precomputed eval manifest holds: draw them before the training plan, + # under the same eval length cap as the manifest rows. + eval_groups = limit_frames({ + key: entries + for key, entries in groups.items() if key[1] in resolutions + }, eval_max_frames) + eval_rows = top_up_eval_rows(eval_rows, eval_groups, int(cfg.get("eval_seed", 20261008)), + dict(eval_per_group)) + else: + eval_rows = interleave_by_group(eval_rows) + plan = select_rows(groups, int(cfg.get("seed", 20261007)), dict(cfg.get("per_group", {"480p": 400})), + frozenset(row.source for row in eval_rows)) + if int(cfg.get("max_train_rows", 0)) > 0: # smoke tests: overfit a few rows + plan = plan[:int(cfg["max_train_rows"])] + world = get_world_group() + sp_size = int(training_config.distributed.sp_size or 1) + self.eval_rows = eval_rows + self.train_rows = plan + self.dataloader = RowStream(plan, + dp_rank=world.rank // sp_size, + dp_size=world.world_size // sp_size, + seed=int(training_config.data.seed or 0)) + logger.info("OmniRef QAD rows: %d train (%s), %d held-out eval", len(plan), dict(sorted(_count(plan).items())), + len(eval_rows)) + + # ------------------------------------------------------------------ PDD ladder + @property + def plan(self) -> Any: + if self._plan is None: + from fastvideo.layers.pdd import PDDModalitySchedule, build_pdd_sampling_plan + schedules = { + "video": PDDModalitySchedule(shift=self.contract.video_shift), + "audio": PDDModalitySchedule(shift=self.contract.audio_shift) + } + self._plan = build_pdd_sampling_plan(self.contract.pdd_step_indices, schedules, device=self.device) + return self._plan + + @property + def num_rungs(self) -> int: + return int(self.plan.num_steps) + + def schedulers(self) -> tuple[Any, Any]: + """Fresh video/audio schedulers on the PDD node sigmas (as the denoising stage sets them).""" + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + video = MiniMaxH3Scheduler(shift=self.contract.video_shift) + audio = MiniMaxH3Scheduler(shift=self.contract.audio_shift) + video.set_timesteps(sigmas=self.plan.node_sigmas["video"].to(torch.float32), device=self.device) + audio.set_timesteps(sigmas=self.plan.node_sigmas["audio"].to(torch.float32), device=self.device) + return video, audio + + def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow: + from fastvideo.train.models.minimax_h3.omniref_data import prepare_row + start = (float(self.plan.node_sigmas["video"][0]), float(self.plan.node_sigmas["audio"][0])) + return prepare_row(spec, + row, + patch_size=self.patch_size, + latent_channels=24, + audio_channels=32, + pdd_start_sigmas=start, + device=self.device) + + def _metadata(self, prepared: PreparedRow, rung: int) -> Any: + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (_h3_vsa_ref2va_segments, + _h3_vsa_single_region_segments) + if self._vsa_builder is None: + self._vsa_builder = MiniMaxH3VSAMetadataBuilder() + keep = self.contract.vsa_ref_keep_rate + segments = (_h3_vsa_single_region_segments(prepared.layout, self.patch_size) + if keep is None else _h3_vsa_ref2va_segments(prepared.layout, self.patch_size)) + return self._vsa_builder.build(current_timestep=rung, + patch_size=self.patch_size, + VSA_sparsity=self.contract.vsa_sparsity, + packed_segments=segments, + device=self.device, + exempt=True, + dense_layers=(), + tile_size=self.contract.vsa_tile_size, + ref_keep_rate=keep) + + @contextlib.contextmanager + def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]: + """Forward context of PDD forward ``rung``; keep backward inside it (activation recompute reads it).""" + from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps + video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item()) + unique, inverse = build_row_timesteps(prepared.layout, + video_timestep=video_t, + audio_timestep=audio_t, + condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG), + condition_audio_timestep=1.0) + start, end = self.plan.block(rung) + fusion = self.transformer.fuse_pdd_block(start, end, self.plan.integration_weights, torch.float32) + with fusion, set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared, rung)): + yield unique.to(self.device), inverse.to(self.device) + + def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor, + row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Packed velocities ``([Lv, Cv], [La, Ca])`` for one forward; call inside ``rung_scope``.""" + layout = prepared.layout + device = self.device + # No autocast: like inference, the transformer casts its inputs to each weight's dtype (BF16 under + # FSDP mixed precision), so Q/K/V reach attention in BF16 exactly as on the deployed path. Autocast + # would run RMSNorm/RoPE in fp32 and change what the FP4 quantizers see. + video_v, audio_v = self.transformer(hidden_states=video[None], + audio_hidden_states=audio[None], + encoder_hidden_states=prepared.prompt_embeds, + timestep=row_timesteps[0], + timestep_indices=row_timesteps[1], + token_tags=layout.token_tags.to(device), + position_ids=layout.position_ids.to(device), + video_indices=layout.video_indices.to(device), + audio_indices=layout.audio_indices.to(device), + text_indices=layout.text_indices.to(device)) + return video_v[0], audio_v[0] + + # ------------------------------------------------------------------ ModelBase contract (unused paths) + def prepare_batch(self, raw_batch: dict[str, Any], **kwargs: Any) -> Any: # noqa: D102 + raise NotImplementedError("use prepare() + rung_scope() + forward_rung()") + + def add_noise(self, clean_latents: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("PDD rungs start from the generator's noise; see prepare()") + + def predict_noise(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError("use forward_rung()") + + def backward(self, loss: torch.Tensor, ctx: Any, *, grad_accum_rounds: int) -> None: + raise NotImplementedError("the PDD QAD method backpropagates inside rung_scope()") + + def sp_group(self) -> Any: + return get_sp_group() + + +def _pin_amax(path: str, expected_sha256: str | None, output_dir: str) -> str: + """Check the calibration table's sha256 (when pinned) and keep a copy next to the run's checkpoints. + + The static FFN scales are fixed buffers that QAD learns around, so a run must never pick up a table that was + rewritten in place. + """ + import hashlib + import shutil + + digest = hashlib.sha256(Path(path).read_bytes()).hexdigest() + if expected_sha256 and digest != expected_sha256: + raise ValueError(f"{path} has sha256 {digest}, but the config pins {expected_sha256}") + if output_dir and get_world_group().rank == 0: + Path(output_dir).mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, Path(output_dir) / f"amax-{digest[:12]}.json") + logger.info("NVFP4 calibration table %s sha256 %s", path, digest) + return digest + + +def _count(rows: list[Any]) -> dict[str, int]: + counts: dict[str, int] = {} + for row in rows: + key = f"{row.case}/{row.resolution}" + counts[key] = counts.get(key, 0) + 1 + return counts + + +__all__ = ["DEFAULT_TRAINABLE", "MiniMaxH3OmniRefPDDModel", "PDDContract"] diff --git a/fastvideo/train/models/minimax_h3/v2_compression.py b/fastvideo/train/models/minimax_h3/v2_compression.py new file mode 100644 index 000000000..b4ab509ef --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_compression.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Training-side alias of the shared FastH3 token compression (the implementation serves inference too).""" +from fastvideo.models.dits.minimax_h3_token_compression import (MODES, CompressionSpec, TokenCompression, + block_flops_share, compressed_grid, grid_ops, + token_reduction) + +__all__ = [ + "MODES", "CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops", + "token_reduction" +] diff --git a/fastvideo/train/models/minimax_h3/v2_data.py b/fastvideo/train/models/minimax_h3/v2_data.py new file mode 100644 index 000000000..ad645fdcb --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_data.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Data-free text-to-video rows for FastH3 V2 QAD. + +A row is a prompt with a precomputed Qwen3-VL text embedding (``precompute_v2_text.py``) plus a seed. The initial +packed state is pure noise (target video + audio), drawn from one CPU generator so (row, seed) reproduces the same +trajectory everywhere; V2 is a plain 8-step DMD ladder, so there is no PDD start scaling and no conditioning rows. +The result mimics ``omniref_data.PreparedRow`` so the QAD method and metrics run unchanged (``case == "t2v"`` +yields no keyframe metrics). +""" +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +V2_CASE = "t2v" + + +@dataclass(frozen=True) +class V2RowSpec: + id: str + source: str + prompt: str + embedding: str # path of the .pt written by precompute_v2_text.py + seed: int + width: int = 832 + height: int = 480 + num_frames: int = 124 + case: str = V2_CASE + + @property + def resolution(self) -> str: + return f"{self.height}p" + + +def load_row_specs(manifest: str, width: int, height: int, num_frames: int) -> list[V2RowSpec]: + """Rows of a ``build_v2_prompts.py`` manifest (JSONL: id, source, prompt, embedding, seed).""" + specs = [] + for line in Path(manifest).read_text().splitlines(): + if not line.strip(): + continue + item = json.loads(line) + specs.append( + V2RowSpec(item["id"], item.get("source", ""), item["prompt"], item["embedding"], int(item["seed"]), + int(item.get("width", width)), int(item.get("height", height)), + int(item.get("num_frames", num_frames)))) + if not specs: + raise ValueError(f"{manifest} has no rows") + return specs + + +def read_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]: + data = torch.load(path, map_location="cpu", weights_only=True) + embeds, tags = data["embeds"], data["tags"] + if embeds.ndim != 2 or tags.shape != (embeds.shape[0], ): + raise ValueError(f"{path}: bad embedding/tag shapes {tuple(embeds.shape)} / {tuple(tags.shape)}") + return embeds, tags.to(torch.long) + + +def prepare_row(spec: V2RowSpec, *, patch_size: tuple[int, int, int], latent_channels: int, audio_channels: int, + device: torch.device) -> Any: + """Initial packed state of ``spec`` (fp32 noise rows), laid out exactly as ``MiniMaxH3LatentPreparationStage``.""" + from diffusers.utils.torch_utils import randn_tensor + + from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, audio_latent_num_frames, + build_packed_sequence, patchify_video_latents, + video_latent_num_frames) + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + + frames, lat_h, lat_w = video_latent_num_frames(spec.num_frames), spec.height // 16, spec.width // 16 + embeds, tags = read_embedding(spec.embedding) + num_audio = audio_latent_num_frames(spec.num_frames) + layout = build_packed_sequence(tags, frames, lat_h, lat_w, num_audio, patch_size) + generator = torch.Generator("cpu").manual_seed(spec.seed) + video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w), + generator=generator, + device=device, + dtype=torch.float32) + audio = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels), + generator=generator, + device=device, + dtype=torch.float32) + video = patchify_video_latents(video_noise, patch_size) + if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"{spec.id}: packed rows do not match the T2V layout") + return PreparedRow(spec, layout, + embeds.to(device=device, dtype=torch.bfloat16)[None], video, audio, + (latent_channels, frames, lat_h, lat_w), [], torch.empty(0, video.shape[-1])) + + +__all__ = ["V2RowSpec", "V2_CASE", "load_row_specs", "prepare_row", "read_embedding"] diff --git a/fastvideo/train/models/minimax_h3/v2_t2v.py b/fastvideo/train/models/minimax_h3/v2_t2v.py new file mode 100644 index 000000000..98a093fd6 --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_t2v.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastH3 V2 (8-step DMD, text-to-video+audio) as a per-forward model for NVFP4 QAD with optional token compression. + +Same contract as ``omniref_pdd.MiniMaxH3OmniRefPDDModel`` so ``OmniRefPDDQADMethod`` runs unchanged: eight forwards +("rungs", one per DMD step), ``rung_scope`` / ``forward_rung`` / ``schedulers`` / ``prepare`` / ``quantization``. + +Differences from OmniRef: plain ``transformer`` module (no PDD head fusion), V2 sampling contract (video shift 10, +audio shift 3, ``dmd_denoising_steps``), VSA tile 64 with sparsity 0.8 (BF16 Triton kernels by default; +``qad.fp4_attention`` installs the same FP4 emulator as the shared recipe), text-only rows (``v2_data``), and the +optional middle-block spatial token compression (``v2_compression``). + +``quantization(linears=, attention=)``: ``linears`` toggles the NVFP4 QAD linears; ``attention=False`` switches off +the FP4 attention emulator and the token compression, so ``quantization(False, False)`` is the BF16 teacher. +""" +from __future__ import annotations + +import contextlib +import json +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from fastvideo.distributed import get_sp_world_size, get_world_group, model_parallel_is_initialized +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.train.models.minimax_h3.omniref_pdd import DEFAULT_TRAINABLE, MiniMaxH3OmniRefPDDModel, _pin_amax +from fastvideo.train.models.minimax_h3.v2_compression import (CompressionSpec, TokenCompression, block_flops_share, + compressed_grid) + +if TYPE_CHECKING: + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + from fastvideo.train.models.minimax_h3.v2_data import V2RowSpec + from fastvideo.train.utils.training_config import TrainingConfig + +logger = init_logger(__name__) +INFERENCE_FILE = "fastvideo_inference.json" + + +@dataclass(frozen=True) +class V2Contract: + """The checkpoint's deployed sampling contract (``fastvideo_inference.json``).""" + dmd_steps: tuple[int, ...] + video_shift: float + audio_shift: float + vsa_sparsity: float + vsa_tile_size: int + + @classmethod + def read(cls, model_path: str) -> V2Contract: + data = json.loads((Path(model_path) / INFERENCE_FILE).read_text()) + if "dmd_denoising_steps" not in data or data.get("task") != "t2av": + raise ValueError(f"{model_path} is not a FastH3 T2AV DMD export") + return cls(tuple(int(i) for i in data["dmd_denoising_steps"]), float(data["video_scheduler_shift"]), + float(data["audio_scheduler_shift"]), float(data["vsa_sparsity"]), int(data["vsa_tile_size"])) + + +class MiniMaxH3V2T2VModel(MiniMaxH3OmniRefPDDModel): + """Teacher (bf16 V2) or NVFP4 QAD student (+ token compression) on the 8-step V2 ladder.""" + + _transformer_module_type = "transformer" + + def __init__(self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + enable_gradient_checkpointing_type: str | None = None, + trainable_patterns: str = DEFAULT_TRAINABLE, + qad: dict[str, Any] | None = None, + compression: dict[str, Any] | None = None, + frozen_bf16: bool = False, + data: dict[str, Any] | None = None) -> None: + self._compression_cfg = dict(compression or {}) + super().__init__(init_from=init_from, + training_config=training_config, + trainable=trainable, + enable_gradient_checkpointing_type=enable_gradient_checkpointing_type, + trainable_patterns=trainable_patterns, + qad=qad, + data=data, + frozen_bf16=frozen_bf16) + self._vsa_builder: Any = self._vsa_builder + self.compression: TokenCompression | None = None + spec = CompressionSpec.from_config(self._compression_cfg) + if spec is not None: + if model_parallel_is_initialized() and get_sp_world_size() > 1: + raise ValueError("token compression needs sp_size 1") + self.compression = TokenCompression(self.transformer, spec) + logger.info("V2 token compression: blocks [%d, %d) mode %s", spec.start, spec.end, spec.mode) + self._compression_on = True + + # ------------------------------------------------------------------ construction hooks + def _read_contract(self, model_path: str) -> Any: + return V2Contract.read(model_path) + + def _check_tile(self) -> None: + if self.contract.vsa_tile_size != 64: + raise ValueError(f"the V2 QAD harness expects VSA tile 64, got {self.contract.vsa_tile_size}") + + def _install_qad(self) -> None: + """NVFP4 linears plus the optional FP4 VSA attention emulator. + + With ``qad.fp4_attention.enabled`` the student trains against the same FP4 attention numerics + used by the fast inference lane; without it V2 keeps the native bf16 VSA path. + """ + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table + + linears = dict(self.qad_config.get("linears") or {}) + if not linears.get("enabled", True): + return + amax_path = linears.get("amax_json") + if not amax_path: + raise ValueError("qad.linears.amax_json (the Consumer export's static FFN scales) is required") + table = load_amax_table(amax_path) + self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"), self.training_config.checkpoint.output_dir) + plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)), + quantize_gate=bool(linears.get("gate", True)), + quantize_ffn=bool(linears.get("ffn", True)), + skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ()))) + self.qad_linears = install_nvfp4_qad(self.transformer, table, plan) + self.fp4_numerics = None + attention = dict(self.qad_config.get("fp4_attention") or {}) + if attention.get("enabled", False): + floor = attention.get("first_block_max_floor") + self.fp4_numerics = FP4AttentionNumerics(quantize=True, + two_level_p=bool(attention.get("two_level_p", False)), + smooth_k=bool(attention.get("smooth_k", False)), + first_block_max_floor=None if floor is None else float(floor)) + install_fp4_vsa_attention(self.transformer, self.fp4_numerics) + logger.info("V2 QAD student: %d NVFP4 QAD linears, attention %s, %.2fB trainable parameters", + len(self.qad_linears), self.fp4_numerics, self.num_trainable / 1e9) + + @contextlib.contextmanager + def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]: + from fastvideo.attention.backends.fp4_vsa_qat import install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled + + set_nvfp4_qad_enabled(self.qad_linears, linears) + fp4 = getattr(self, "fp4_numerics", None) + if fp4 is not None and not attention: + install_fp4_vsa_attention(self.transformer, None) + previous = self._compression_on + self._compression_on = attention + try: + yield + finally: + set_nvfp4_qad_enabled(self.qad_linears, True) + if fp4 is not None: + install_fp4_vsa_attention(self.transformer, fp4) + self._compression_on = previous + + # ------------------------------------------------------------------ data + def init_preprocessors(self, training_config: TrainingConfig) -> None: + from fastvideo.train.models.minimax_h3.omniref_data import RowStream + from fastvideo.train.models.minimax_h3.v2_data import load_row_specs + + cfg = self.data_config + size = (int(cfg.get("width", 832)), int(cfg.get("height", 480)), int(cfg.get("num_frames", 124))) + eval_rows = load_row_specs(cfg["eval_manifest"], *size) + eval_ids = {row.id for row in eval_rows} + train_rows = [row for row in load_row_specs(cfg["train_manifest"], *size) if row.id not in eval_ids] + if int(cfg.get("max_train_rows", 0)) > 0: + train_rows = train_rows[:int(cfg["max_train_rows"])] + if int(cfg.get("max_eval_rows", 0)) > 0: + eval_rows = eval_rows[:int(cfg["max_eval_rows"])] + world = get_world_group() + sp_size = int(training_config.distributed.sp_size or 1) + self.eval_rows, self.train_rows = eval_rows, train_rows + self.dataloader = RowStream(train_rows, + dp_rank=world.rank // sp_size, + dp_size=world.world_size // sp_size, + seed=int(training_config.data.seed or 0)) + logger.info("V2 QAD rows: %d train, %d held-out eval", len(train_rows), len(eval_rows)) + + def load_row(self, spec: V2RowSpec) -> PreparedRow: + from fastvideo.train.models.minimax_h3.v2_data import prepare_row + return prepare_row(spec, patch_size=self.patch_size, latent_channels=24, audio_channels=32, device=self.device) + + # ------------------------------------------------------------------ DMD ladder + @property + def plan(self) -> Any: # pragma: no cover - PDD-only + raise AttributeError("V2 has no PDD plan") + + @property + def num_rungs(self) -> int: + return len(self.contract.dmd_steps) + + def schedulers(self) -> tuple[Any, Any]: + """Fresh video/audio schedulers on the DMD sigma grid, as ``MiniMaxH3DenoisingStage._set_dmd_schedule``.""" + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + base = torch.tensor([step / 1000.0 for step in self.contract.dmd_steps] + [0.0], dtype=torch.float32) + out = [] + for shift in (self.contract.video_shift, self.contract.audio_shift): + scheduler = MiniMaxH3Scheduler(shift=shift) + scheduler.set_timesteps(sigmas=shift * base / (1 + (shift - 1) * base), device=self.device) + out.append(scheduler) + return out[0], out[1] + + def _token_grid(self, prepared: PreparedRow) -> tuple[int, int, int]: + layout = self.patch_size + return (prepared.layout.num_video_latent_frames // layout[0], prepared.layout.latent_height // layout[1], + prepared.layout.latent_width // layout[2]) + + def _metadata(self, prepared: PreparedRow, rung: int, grid: tuple[int, int, int] | None = None) -> Any: + """VSA-H3 metadata for the full video grid, or for a (compressed) token ``grid``.""" + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import _h3_vsa_single_region_segments + if self._vsa_builder is None: + self._vsa_builder = MiniMaxH3VSAMetadataBuilder() + layout = prepared.layout + if grid is None: + segments = _h3_vsa_single_region_segments(layout, self.patch_size) + else: + p = self.patch_size + segments = (int(layout.text_indices.numel()), int(layout.audio_indices.numel()), + (grid[0] * p[0], grid[1] * p[1], grid[2] * p[2])) + return self._vsa_builder.build(current_timestep=rung, + patch_size=self.patch_size, + VSA_sparsity=self.contract.vsa_sparsity, + packed_segments=segments, + device=self.device, + exempt=True, + dense_layers=(), + tile_size=self.contract.vsa_tile_size) + + @contextlib.contextmanager + def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]: + """Forward context of DMD forward ``rung``; keep backward inside it (activation recompute reads it).""" + from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps + video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item()) + unique, inverse = build_row_timesteps(prepared.layout, + video_timestep=video_t, + audio_timestep=audio_t, + condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG), + condition_audio_timestep=1.0) + with contextlib.ExitStack() as stack: + stack.enter_context(set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared, + rung))) + first_steps = self.compression.spec.first_steps if self.compression is not None else None + if self.compression is not None and self._compression_on and (first_steps is None or rung < first_steps): + grid = self._token_grid(prepared) + sub = compressed_grid(self.compression.spec.mode, grid) + stack.enter_context(self.compression.scope(grid, self._metadata(prepared, rung, sub))) + yield unique.to(self.device), inverse.to(self.device) + + def set_compression(self, spec: CompressionSpec | None) -> None: + """Replace the token-compression spec (probes and sweeps); ``None`` removes it.""" + if self.compression is not None: + self.compression.remove() + self.compression = None if spec is None else TokenCompression(self.transformer, spec) + + def compression_summary(self, prepared: PreparedRow) -> dict[str, float]: + """Token and FLOP accounting of the configured compression for one row.""" + spec = self.compression.spec if self.compression else None + grid = self._token_grid(prepared) + other = prepared.layout.sequence_length - grid[0] * grid[1] * grid[2] + out = block_flops_share(spec, grid, len(self.transformer.transformer_blocks), other) + out["video_tokens"] = float(grid[0] * grid[1] * grid[2]) + out["sequence_tokens"] = float(prepared.layout.sequence_length) + return out + + # ------------------------------------------------------------------ PDD-only paths of the parent + def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor, + row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + return super().forward_rung(prepared, video, audio, row_timesteps) + + +__all__ = ["MiniMaxH3V2T2VModel", "V2Contract"] diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index f4d438feb..f99cb5a65 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -139,6 +139,21 @@ class Trainer: if self.global_rank == 0 and validation_metrics: self.tracker.log(validation_metrics, iteration) + def _consume_method_requests(self, method: TrainingMethod, checkpoint_manager: Any | None, step: int) -> bool: + """Save a method-requested tagged checkpoint and report whether training should stop. + + Runs immediately after every evaluation, including the pre-loop one, so + a request always refers to the weights that were just evaluated. + """ + request_hook = getattr(method, "consume_checkpoint_request", None) + tag = request_hook() if callable(request_hook) else None + if tag and checkpoint_manager is not None: + checkpoint_manager.save_tagged(str(tag), step) + stop = bool(getattr(method, "stop_requested", False)) + if stop and self.global_rank == 0: + self.tracker.log({"stopped_early_at": float(step)}, step) + return stop + def run( self, method: TrainingMethod, @@ -186,6 +201,7 @@ class Trainer: iteration=start_step, ) self._run_method_validation(method, start_step) + stop_early = self._consume_method_requests(method, checkpoint_manager, start_step) method.optimizers_zero_grad(start_step) data_stream = self._iter_dataloader(dataloader, @@ -210,7 +226,10 @@ class Trainer: ) # Allow method-specific optimization flow (e.g. DiffusionNFT). method_manages_optimization = bool(method.manages_optimization()) + last_step = start_step for step in progress: + if stop_early: + break t0 = time.perf_counter() if performance_enabled: self.performance_monitor.reset() @@ -332,14 +351,20 @@ class Trainer: method, iteration=step, ) + last_step = step + + if self._consume_method_requests(method, checkpoint_manager, step): + break + else: + last_step = max_steps self.callbacks.on_train_end( method, - iteration=max_steps, + iteration=last_step, ) if checkpoint_manager is not None: - checkpoint_manager.save_final(max_steps) + checkpoint_manager.save_final(last_step) self.performance_monitor.close() self.tracker.finish() diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 7987c3341..8db7f9155 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -98,6 +98,17 @@ def _rank() -> int: return 0 +_CPU_GROUP: Any = None + + +def _cpu_group() -> Any: + """A gloo group over all ranks (created once, collectively) for checkpoint coordination.""" + global _CPU_GROUP + if _CPU_GROUP is None and dist.is_available() and dist.is_initialized(): + _CPU_GROUP = dist.new_group(backend="gloo") + return _CPU_GROUP + + def _barrier() -> None: if dist.is_available() and dist.is_initialized(): dist.barrier() @@ -773,12 +784,43 @@ class CheckpointManager: shutil.rmtree(path) snapshots = {node: saved for node, saved in snapshots.items() if node != path and path not in node.parents} + def save_tagged(self, tag: str, step: int) -> None: + """Save the trainable roles' model weights only to ``/`` (overwritten). + + Tagged directories are not ``checkpoint-`` and are never pruned; + they hold no optimizer state (resume from regular checkpoints). + """ + # A leading alphanumeric excludes "." and ".." (the run directory and its parent, which rmtree would delete). + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", tag) or _CHECKPOINT_DIR_RE.match(tag): + raise ValueError(f"invalid checkpoint tag {tag!r}") + tagged_dir = Path(self.output_dir) / tag + if tagged_dir.resolve().parent != Path(self.output_dir).resolve(): + raise ValueError(f"checkpoint tag {tag!r} does not name a directory directly inside {self.output_dir}") + dcp_dir = tagged_dir / "dcp" + if _rank() == 0: + if tagged_dir.exists(): + shutil.rmtree(tagged_dir) + os.makedirs(dcp_dir, exist_ok=True) + self._write_metadata(tagged_dir, step, weights_only=True) + logger.info("Saving %s weights (step %d) to %s", tag, step, tagged_dir) + _barrier() + states = {key: value for key, value in self.method.checkpoint_state().items() if key.startswith("roles.")} + # Coordinate over a CPU group: DCP's metadata gather over NCCL needs device memory, and right after an + # evaluation the GPUs can be nearly full (seen as "NCCL Error 1: unhandled cuda error" in gather). + torch.cuda.empty_cache() + dcp.save(states, checkpoint_id=str(dcp_dir), process_group=_cpu_group()) + _barrier() + def _write_metadata( self, checkpoint_dir: Path, step: int, + *, + weights_only: bool = False, ) -> None: metadata: dict[str, Any] = {"step": step} + if weights_only: + metadata["weights_only"] = True if self._raw_config is not None: metadata["config"] = self._raw_config meta_path = checkpoint_dir / "metadata.json" diff --git a/scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py b/scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py new file mode 100644 index 000000000..983f5bef3 --- /dev/null +++ b/scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate FastH3 OmniRef (Ref2VA) DiT latents from precomputed conditioning rows. + +Training and evaluation data for reference-conditioned FastH3 students. Each input +row is one OmniRef Ref2VA training parquet row (``minimax_h3_ref2va`` record +schema): the Qwen3-VL presentation (``text_embedding`` + ``text_token_tags``), +the clean normalized reference latent rows (``reference_video_latent_rows`` / +``reference_audio_latent_rows``) and their ordered geometry. Nothing is +re-encoded: no text encoder, no reference media, no VAE encode. + +Per row this script does what ``MiniMaxH3LatentPreparationStage`` does after its +encoders (0.999 noise augmentation of the reference rows with the request +generator, then target video/audio noise, packed with +``build_ref2va_packed_sequence``), runs the pipeline's own +``MiniMaxH3DenoisingStage`` (PDD fused blocks, VSA-H3 with the checkpoint's +reference keep rate), and keeps the target video rows exactly as +``MiniMaxH3VideoDecodingStage`` would decode them. The saved latent is +``[24, T, H, W]`` float32 in the VAE's normalized space, +``(raw - latents_mean) / latents_std``: the DiT's own space, so this equals +decode-stage ``denormalize_latents`` followed by normalization. + +Output rows match ``generate_h3_latents.py`` (``id``, ``vae_latent_bytes/shape/ +dtype``, ``generator``, ``case``, ``width``, ``height``, ``num_frames``, ``seed``) +and are written atomically, ``FLUSH_EVERY`` rows per file. Re-running skips ids +already written. Rows are chosen deterministically: per (case, resolution) the +manifest entries up to ``--max-frames`` are shuffled with ``--seed`` and the first +``--per-group`` are taken, interleaved across groups so a partial run stays +balanced. + +One process per GPU (each its own single-rank external launcher), e.g.:: + + CUDA_VISIBLE_DEVICES=0 MASTER_PORT=29500 python generate_omniref_latents.py \\ + --model-path /models/fasth3_omniref_composed --manifest /data/omniref/MANIFEST.jsonl \\ + --output-dir /data/gen-latents/omniref/data \\ + --per-group 300 --shard 0 --num-shards 4 +""" +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +import random +import re +import time +from collections import defaultdict +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +FLUSH_EVERY = 8 +GENERATOR_NAME = "omniref_step3500" +CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot") +RESOLUTIONS = ("480p", "768p") +BUCKET_FRAMES = re.compile(r"bucket=\d+x\d+-(\d+)f") +ROW_COLUMNS = [ + "id", + "vae_latent_shape", + "audio_latent_shape", + "text_embedding_bytes", + "text_embedding_shape", + "text_token_tags_bytes", + "text_token_tags_shape", + "reference_video_latent_rows_bytes", + "reference_video_latent_rows_shape", + "reference_audio_latent_rows_bytes", + "reference_audio_latent_rows_shape", + "reference_media_types", + "reference_has_audio", + "reference_num_latent_frames", + "reference_latent_heights", + "reference_latent_widths", + "reference_num_audio_latents", + "width", + "height", + "num_frames", +] + + +# --------------------------------------------------------------------------- plan +def clip_seed(base_seed: int, clip_id: str) -> int: + digest = hashlib.sha256(f"{base_seed}:{clip_id}".encode()).digest() + return int.from_bytes(digest[:4], "little") & 0x7FFFFFFF + + +def load_manifest(path: str, cases: list[str], max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]: + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + with open(path) as handle: + for line in handle: + entry = json.loads(line) + match = BUCKET_FRAMES.search(entry["parquet"]) + if match is None or entry["case"] not in cases or entry["resolution"] not in RESOLUTIONS: + continue + frames = int(match.group(1)) + if frames <= max_frames: + groups[(entry["case"], entry["resolution"])].append({**entry, "frames": frames}) + return groups + + +def select_clips(groups: dict[tuple[str, str], list[dict[str, Any]]], + seed: int, + per_group: int, + exclude_sources: frozenset[str] = frozenset()) -> list[dict[str, Any]]: + """Per group: seeded shuffle, drop excluded source clips, take ``per_group``; interleave the groups.""" + picked = [] + for key in sorted(groups): + entries = sorted(groups[key], key=lambda entry: entry["id"]) + random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries) + entries = [entry for entry in entries if entry["id"].split(":", 1)[-1] not in exclude_sources] + picked.append(entries[:per_group]) + plan = [] + for index in range(max((len(entries) for entries in picked), default=0)): + for entries in picked: + if index < len(entries): + entry = entries[index] + source_id = entry["id"].split(":", 1)[-1] + clip_id = f"omniref-{entry['case']}-{source_id}-{entry['resolution']}" + plan.append({ + "id": clip_id, + "source": source_id, + "case": entry["case"], + "resolution": entry["resolution"], + "parquet": entry["parquet"], + "seed": clip_seed(seed, clip_id) + }) + return plan + + +def clip_plan(args: argparse.Namespace) -> list[dict[str, Any]]: + """The full deterministic plan; shard ``s`` takes every ``num_shards``-th entry from ``s``.""" + groups = load_manifest(args.manifest, args.cases, args.max_frames) + if args.eval_dir is None: + return select_clips(groups, args.seed, args.per_group)[args.shard::args.num_shards] + # A held-out set: no source clip (at either resolution) of the main plan. + main_plan = select_clips(groups, args.exclude_seed, args.exclude_per_group) + excluded = frozenset(clip["source"] for clip in main_plan) + return select_clips(groups, args.seed, args.per_group, excluded)[args.shard::args.num_shards] + + +def existing_ids(data_dir: Path) -> set[str]: + import pyarrow.parquet as pq + + done: set[str] = set() + for path in data_dir.glob("*.parquet"): + done.update(pq.read_table(path, columns=["id"]).column("id").to_pylist()) + return done + + +def write_part(data_dir: Path, shard: int, rows: list[dict[str, Any]]) -> Path: + import pyarrow as pa + import pyarrow.parquet as pq + + path = data_dir / f"omniref-shard{shard:03d}-{int(time.time() * 1000)}.parquet" + tmp = path.with_suffix(".tmp") + pq.write_table(pa.Table.from_pylist(rows), tmp) + tmp.replace(path) + return path + + +# --------------------------------------------------------------------------- rows +def _tensor(row: dict[str, Any], name: str) -> torch.Tensor: + shape = [int(value) for value in row[f"{name}_shape"]] + array = np.frombuffer(row[f"{name}_bytes"], dtype=np.float32) + if array.size != int(np.prod(shape, dtype=np.int64)): + raise ValueError(f"{name}: {array.size} values do not fill shape {shape}") + return torch.from_numpy(array.reshape(shape).copy()) + + +def read_row(parquet: str) -> dict[str, Any]: + import pyarrow.parquet as pq + + rows = pq.read_table(parquet, columns=ROW_COLUMNS).to_pylist() + if len(rows) != 1: + raise ValueError(f"{parquet}: expected one row, got {len(rows)}") + return rows[0] + + +def prepared_references(row: dict[str, Any]) -> list[Any]: + from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference + + fields = ("reference_media_types", "reference_has_audio", "reference_num_latent_frames", + "reference_latent_heights", "reference_latent_widths", "reference_num_audio_latents") + return [ + MiniMaxH3PreparedReference(media_type=media_type, + has_audio=bool(has_audio), + num_latent_frames=int(frames), + latent_height=int(height), + latent_width=int(width), + num_audio_latents=int(audio)) + for media_type, has_audio, frames, height, width, audio in zip(*(row[name] for name in fields), strict=True) + ] + + +# --------------------------------------------------------------------------- generation +class OmniRefLatentGenerator: + """Single-rank (or SPMD) pipeline driver that injects precomputed Ref2VA conditioning.""" + + def __init__(self, args: argparse.Namespace) -> None: + from fastvideo import VideoGenerator + from fastvideo.api import (ComponentConfig, EngineConfig, GeneratorConfig, OffloadConfig, ParallelismConfig, + PipelineSelection) + + world_size = int(os.environ.setdefault("WORLD_SIZE", "1")) + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", str(args.master_port)) + config = GeneratorConfig( + model_path=args.model_path, + engine=EngineConfig( + num_gpus=world_size, + execution_backend="external_launcher", + use_fsdp_inference=world_size > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=world_size), + offload=OffloadConfig(dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + lazy_module_load=True), + ), + pipeline=PipelineSelection( + workload_type="i2v", + components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"), + ), + ) + self.generator = VideoGenerator.from_config(config) + worker = self.generator.executor.worker.worker + self.pipeline = worker.pipeline + self.fastvideo_args = worker.fastvideo_args + self.is_output_rank = bool(self.fastvideo_args.is_output_rank) + if not self.pipeline._denoise_stages_ready: + self.pipeline._add_denoise_stages(ref2va=True) + self.denoise = self.pipeline._stage_name_mapping["denoising_stage"] + self.scheduler = self.pipeline.get_module("scheduler") + + def _geometry(self, row: dict[str, Any]) -> tuple[int, int, int, int]: + from fastvideo.pipelines.basic.minimax_h3.packing import audio_latent_num_frames, video_latent_num_frames + + arch = self.fastvideo_args.pipeline_config.vae_config.arch_config + ratio = int(arch.spatial_compression_ratio) + num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"]) + geometry = (video_latent_num_frames(num_frames), height // ratio, width // ratio) + expected = tuple(int(value) for value in row["vae_latent_shape"][1:]) + if geometry != expected or int(row["vae_latent_shape"][0]) != int(arch.latent_channels): + raise ValueError(f"target latent geometry {geometry} disagrees with the row's {row['vae_latent_shape']}") + num_audio = audio_latent_num_frames(num_frames) + if num_audio != int(row["audio_latent_shape"][-1]): + raise ValueError(f"{num_audio} audio latents disagree with the row's {row['audio_latent_shape']}") + return (*geometry, num_audio) + + def _batch(self, row: dict[str, Any], seed: int) -> Any: + """The ForwardBatch ``MiniMaxH3LatentPreparationStage`` would hand to denoising.""" + from diffusers.utils.torch_utils import randn_tensor + + from fastvideo.distributed import get_local_torch_device + from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_KEYFRAME_NOISE_AUG, + build_ref2va_packed_sequence, h3_dit_patch_size, + keyframe_condition_noise, patchify_video_latents) + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + device = get_local_torch_device() + args = self.fastvideo_args + patch = h3_dit_patch_size(args) + channels = int(args.pipeline_config.vae_config.arch_config.latent_channels) + audio_channels = int(args.pipeline_config.audio_vae_config.arch_config.latent_channels) + frames, height, width, num_audio = self._geometry(row) + references = prepared_references(row) + tags = _tensor(row, "text_token_tags").round().to(torch.long) + embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None] + layout = build_ref2va_packed_sequence(tags, references, frames, height, width, num_audio, patch) + + generator = torch.Generator("cpu").manual_seed(seed) + shapes = tuple((ref.num_latent_frames, ref.latent_height, ref.latent_width) for ref in references + if ref.media_type != "audio") + noise = keyframe_condition_noise(shapes, patch, channels, generator=generator, device=device) + clean = _tensor(row, "reference_video_latent_rows").to(device) + condition_video = self.scheduler.scale_noise(clean, MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) + video_noise = randn_tensor((1, channels, frames, height, width), + generator=generator, + device=device, + dtype=torch.float32) + audio_rows = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels), + generator=generator, + device=device, + dtype=torch.float32) + condition_audio = _tensor(row, "reference_audio_latent_rows").to(device) + video_rows = torch.cat((condition_video, patchify_video_latents(video_noise, patch))) + audio_rows = torch.cat((condition_audio, audio_rows)) + if video_rows.shape[0] != layout.video_indices.numel() or audio_rows.shape[0] != layout.audio_indices.numel(): + raise ValueError("packed rows do not match the Ref2VA layout") + + batch = ForwardBatch(data_type="video", + prompt="", + seed=seed, + num_frames=int(row["num_frames"]), + height=int(row["height"]), + width=int(row["width"]), + num_inference_steps=len(args.pipeline_config.pdd_step_indices) - 1, + guidance_scale=1.0, + VSA_sparsity=float(args.VSA_sparsity)) + batch.generator = generator + batch.prompt_embeds = [embeds] + batch.latents = video_rows + batch.audio_latents = audio_rows + batch.raw_latent_shape = (1, channels, frames, height, width) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + return batch + + @torch.no_grad() + def generate(self, row: dict[str, Any], seed: int) -> torch.Tensor: + """Normalized target-video latent ``[24, T, H, W]`` (float32, CPU).""" + from fastvideo.pipelines.basic.minimax_h3.packing import h3_dit_patch_size, unpatchify_video_tokens + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY + + batch = self.denoise.forward(self._batch(row, seed), self.fastvideo_args) + layout = batch.extra[MINIMAX_H3_LAYOUT_KEY] + _, channels, frames, height, width = batch.raw_latent_shape + latents = unpatchify_video_tokens(batch.latents[layout.num_condition_video_rows:], frames, height, width, + channels, h3_dit_patch_size(self.fastvideo_args)) + return latents[0].float().cpu().contiguous() + + @torch.no_grad() + def decode(self, normalized: torch.Tensor) -> np.ndarray: + """Decode one normalized ``[24, T, H, W]`` latent to uint8 ``[T, H, W, 3]`` with the original VAE.""" + from fastvideo.distributed import get_local_torch_device + + from fastvideo.models import pinned_offload + + device = get_local_torch_device() + vae = self.pipeline.get_module("vae") + pinned_offload.load(vae, device, pin=False) + try: + latents = vae.denormalize_latents(normalized[None].to(device=device, dtype=torch.float32)) + output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + with torch.autocast(device_type=device.type, dtype=torch.float16): + vae.decode_to_pixels(latents, output) + finally: + pinned_offload.unload(vae) + return output[0].mul(255).clamp(0, 255).to(torch.uint8).permute(1, 2, 3, 0).numpy() + + def shutdown(self) -> None: + self.generator.shutdown() + + +# --------------------------------------------------------------------------- preview +def first_reference_latent(row: dict[str, Any], patch: tuple[int, int, int]) -> torch.Tensor | None: + from fastvideo.pipelines.basic.minimax_h3.packing import unpatchify_video_tokens + + references = [ref for ref in prepared_references(row) if ref.media_type != "audio"] + if not references: + return None + ref = references[0] + count = ref.num_latent_frames * (ref.latent_height // patch[1]) * (ref.latent_width // patch[2]) + rows = _tensor(row, "reference_video_latent_rows")[:count] + return unpatchify_video_tokens(rows, ref.num_latent_frames, ref.latent_height, ref.latent_width, 24, patch)[0] + + +def write_preview(driver: OmniRefLatentGenerator, clip: dict[str, Any], row: dict[str, Any], saved: bytes, + shape: list[int], preview_dir: Path) -> None: + """Decode the serialized latent, the dataset target and the first reference for a visual check.""" + import imageio.v2 as imageio + import pyarrow.parquet as pq + + from fastvideo.pipelines.basic.minimax_h3.packing import h3_dit_patch_size + + preview_dir.mkdir(parents=True, exist_ok=True) + generated = driver.decode(torch.from_numpy(np.frombuffer(saved, dtype=np.float32).reshape(shape).copy())) + target_row = pq.read_table(clip["parquet"], columns=["vae_latent_bytes", "vae_latent_shape"]).to_pylist()[0] + target_latent = _tensor(target_row, "vae_latent") + target = driver.decode(target_latent) + stem = preview_dir / clip["id"] + imageio.mimsave(f"{stem}.mp4", list(np.concatenate((generated, target), axis=2)), fps=24, format="mp4") + reference = first_reference_latent(row, h3_dit_patch_size(driver.fastvideo_args)) + if reference is not None: + frames = driver.decode(reference) + imageio.mimsave(f"{stem}.reference.mp4", list(frames), fps=24, format="mp4") + stats = { + "id": clip["id"], + "generated_mean": float(np.frombuffer(saved, dtype=np.float32).mean()), + "generated_std": float(np.frombuffer(saved, dtype=np.float32).std()), + "target_mean": float(target_latent.mean()), + "target_std": float(target_latent.std()), + } + print(json.dumps({"preview": str(stem), **stats}), flush=True) + + +# --------------------------------------------------------------------------- main +def append_eval_manifest(path: Path, entry: dict[str, Any]) -> None: + """Append one row under an exclusive lock, rereading so shard processes don't drop each other's rows.""" + path.parent.mkdir(parents=True, exist_ok=True) + with open(path.with_suffix(".lock"), "a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + current = json.loads(path.read_text()) if path.is_file() else [] + if entry["id"] not in {item["id"] for item in current}: + current.append(entry) + tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp") + tmp.write_text(json.dumps(current, indent=1)) + tmp.replace(path) + finally: + fcntl.flock(lock, fcntl.LOCK_UN) + + +def run_eval(args: argparse.Namespace, plan: list[dict[str, Any]]) -> None: + """Write raw (denormalized) ``[1, 24, T, H, W]`` latents, as the decode stage decodes them, plus a manifest.""" + eval_dir = Path(args.eval_dir) + eval_dir.mkdir(parents=True, exist_ok=True) + manifest_path = eval_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text()) if manifest_path.is_file() else [] + done = {item["id"] for item in manifest} + # The model's own vae/config.json: the pipeline config's VAE arch in this process keeps its dataclass + # defaults (zeros/ones) until the VAE itself loads. + vae_config = json.loads((Path(args.model_path) / "vae" / "config.json").read_text()) + mean = torch.tensor(vae_config["latents_mean"], dtype=torch.float32).view(1, -1, 1, 1, 1) + std = torch.tensor(vae_config["latents_std"], dtype=torch.float32).view(1, -1, 1, 1, 1) + driver = OmniRefLatentGenerator(args) + for clip in plan: + if clip["id"] in done: + continue + start = time.perf_counter() + row = read_row(clip["parquet"]) + raw = (driver.generate(row, clip["seed"])[None] * std + mean).contiguous() + if not bool(torch.isfinite(raw).all()): + print(json.dumps({"id": clip["id"], "error": "non-finite latent"}), flush=True) + continue + name = f"{clip['case']}_{clip['resolution']}_{clip['id']}.pt" + if driver.is_output_rank: + tmp = eval_dir / f"{name}.tmp" + torch.save(raw, tmp) + tmp.replace(eval_dir / name) + append_eval_manifest( + manifest_path, { + "file": name, + **{key: clip[key] for key in ("id", "source", "case", "resolution", "parquet", "seed")}, + "width": int(row["width"]), + "height": int(row["height"]), + "num_frames": int(row["num_frames"]), + "shape": list(raw.shape), + "dtype": "float32", + "space": "raw (denormalized; decode-stage input)", + "generator": args.generator_name, + }) + print(json.dumps({"id": clip["id"], "shape": list(raw.shape), "seconds": round(time.perf_counter() - start, 1)}), + flush=True) + driver.shutdown() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", required=True, help="composed FastH3 OmniRef PDD model directory") + parser.add_argument("--manifest", required=True, help="OmniRef MANIFEST.jsonl (case, id, parquet, resolution)") + parser.add_argument("--output-dir", default=None, help="directory receiving the parquet parts") + parser.add_argument("--cases", nargs="+", default=list(CASES)) + parser.add_argument("--per-group", type=int, default=300, help="clips per (case, resolution)") + parser.add_argument("--max-frames", type=int, default=243) + parser.add_argument("--seed", type=int, default=20261007) + parser.add_argument("--shard", type=int, default=0) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--limit", type=int, default=0, help="stop after this many new clips (0: no limit)") + parser.add_argument("--generator-name", default=GENERATOR_NAME) + parser.add_argument("--master-port", type=int, default=29500) + parser.add_argument("--preview-dir", default=None, help="decode the first --preview clips here as MP4") + parser.add_argument("--preview", type=int, default=0) + parser.add_argument("--eval-dir", + default=None, + help="held-out mode: write raw [1, 24, T, H, W] .pt latents and manifest.json here, from " + "source clips outside the main plan (--exclude-seed / --exclude-per-group)") + parser.add_argument("--exclude-seed", type=int, default=20261007) + parser.add_argument("--exclude-per-group", type=int, default=300) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + plan = clip_plan(args) + if args.eval_dir is not None: + print(json.dumps({"eval_clips": len(plan)}), flush=True) + run_eval(args, plan) + return + if args.output_dir is None: + raise SystemExit("--output-dir is required unless --eval-dir is given") + data_dir = Path(args.output_dir) + data_dir.mkdir(parents=True, exist_ok=True) + done = existing_ids(data_dir) + todo = [clip for clip in plan if clip["id"] not in done] + if args.limit: + todo = todo[:args.limit] + print(json.dumps({"shard": args.shard, "planned": len(plan), "todo": len(todo)}), flush=True) + if not todo: + return + driver = OmniRefLatentGenerator(args) + pending: list[dict[str, Any]] = [] + previews = 0 + for clip in todo: + start = time.perf_counter() + try: + row = read_row(clip["parquet"]) + latent = driver.generate(row, clip["seed"]) + if not bool(torch.isfinite(latent).all()): + raise ValueError("non-finite latent") + except Exception as error: # noqa: BLE001 - record and continue; one bad row must not stop the shard + print(json.dumps({"id": clip["id"], "error": repr(error)[:500]}), flush=True) + continue + record = { + "id": clip["id"], + "vae_latent_bytes": latent.numpy().tobytes(), + "vae_latent_shape": list(latent.shape), + "vae_latent_dtype": "float32", + "generator": args.generator_name, + "case": clip["case"], + "width": int(row["width"]), + "height": int(row["height"]), + "num_frames": int(row["num_frames"]), + "seed": int(clip["seed"]), + } + print(json.dumps({ + "id": clip["id"], + "shape": record["vae_latent_shape"], + "seconds": round(time.perf_counter() - start, 1) + }), + flush=True) + if not driver.is_output_rank: + continue + pending.append(record) + if args.preview_dir and previews < args.preview: + write_preview(driver, clip, row, record["vae_latent_bytes"], record["vae_latent_shape"], + Path(args.preview_dir)) + previews += 1 + if len(pending) >= FLUSH_EVERY: + write_part(data_dir, args.shard, pending) + pending = [] + if pending and driver.is_output_rank: + write_part(data_dir, args.shard, pending) + driver.shutdown() + + +if __name__ == "__main__": + main() diff --git a/scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py b/scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py new file mode 100644 index 000000000..3a1ad9197 --- /dev/null +++ b/scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Build the V2 QAD prompt manifests (train + held-out eval) from prompt sources. + +Sources: parquet shards with a ``prompt`` column (the decoder latent generator's outputs; only that column is read) +and JSONL files with a ``prompt`` field (e.g. ``source.jsonl``, OmniRef captions). Prompts are deduplicated, then a +seeded shuffle assigns ``--eval-count`` held-out prompts and the rest to training. Each row gets a deterministic seed. +Embeddings are filled in by ``precompute_v2_text.py`` (rows carry the future embedding path). +""" +from __future__ import annotations + +import argparse +import glob +import hashlib +import json +import random +from pathlib import Path + + +def read_prompts(parquet_globs: list[str], jsonl_globs: list[str]) -> list[str]: + import pyarrow.parquet as pq + + found: dict[str, str] = {} + for pattern in parquet_globs: + for path in sorted(glob.glob(pattern)): + if "prompt" in pq.read_schema(path).names: + for prompt in pq.read_table(path, columns=["prompt"]).column("prompt").to_pylist(): + found.setdefault(" ".join(prompt.split()), path) + for pattern in jsonl_globs: + for path in sorted(glob.glob(pattern)): + with open(path) as handle: + for line in handle: + if line.strip(): + prompt = json.loads(line).get("prompt") + if isinstance(prompt, str) and prompt.strip(): + found.setdefault(" ".join(prompt.split()), path) + return sorted(found) + + +def make_rows(prompts: list[str], embed_dir: str, seed: int) -> list[dict]: + rows = [] + for prompt in prompts: + digest = hashlib.sha256(prompt.encode()).hexdigest()[:16] + rows.append({ + "id": f"v2-{digest}", + "source": "prompt", + "prompt": prompt, + "embedding": str(Path(embed_dir) / f"{digest}.pt"), + "seed": int(hashlib.sha256(f"{seed}:{digest}".encode()).hexdigest()[:8], 16) & 0x7FFFFFFF, + }) + return rows + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--parquet", nargs="*", default=[], help="globs of parquet shards with a prompt column") + parser.add_argument("--jsonl", nargs="*", default=[], help="globs of JSONL files with a prompt field") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--eval-count", type=int, default=30) + parser.add_argument("--max-train", type=int, default=0, help="0 = all") + parser.add_argument("--seed", type=int, default=20261016) + args = parser.parse_args() + prompts = read_prompts(args.parquet, args.jsonl) + random.Random(args.seed).shuffle(prompts) + if len(prompts) <= args.eval_count: + raise SystemExit(f"only {len(prompts)} prompts") + out = Path(args.out_dir) + embed_dir = out / "embeddings" + embed_dir.mkdir(parents=True, exist_ok=True) + eval_prompts, train_prompts = prompts[:args.eval_count], prompts[args.eval_count:] + if args.max_train: + train_prompts = train_prompts[:args.max_train] + for name, subset in (("eval", eval_prompts), ("train", train_prompts)): + rows = make_rows(subset, str(embed_dir), args.seed) + (out / f"{name}.jsonl").write_text("".join(json.dumps(row) + "\n" for row in rows)) + print(json.dumps({"split": name, "rows": len(rows)})) + + +if __name__ == "__main__": + main() diff --git a/scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py b/scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py new file mode 100644 index 000000000..9a9db65e0 --- /dev/null +++ b/scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Precompute FastH3 text embeddings (Qwen3-VL conditioner, bf16 V2 model) for the V2 QAD manifests. + +Builds the V2 inference pipeline once (single GPU), calls its conditioning stage on each prompt and writes +``{"embeds": [L, D] bf16, "tags": [L] long}`` to every row's ``embedding`` path. Existing files are kept, so reruns and +sharding (``--shard i --num-shards n``) are safe. +""" +from __future__ import annotations + +import argparse +import json +import os +import types +from pathlib import Path + +import torch + + +def _build(config_path: str, model_path: str, shard: int): + """One-rank external-launcher generator, configured for latent output only.""" + from fastvideo.api.compat import generator_config_to_fastvideo_args, normalize_generator_config + from fastvideo.api.parser import load_raw_config + from fastvideo.entrypoints.video_generator import VideoGenerator + + for name, value in (("WORLD_SIZE", "1"), ("RANK", "0"), ("LOCAL_RANK", "0"), ("MASTER_ADDR", "127.0.0.1")): + os.environ.setdefault(name, value) + raw = load_raw_config(config_path) + generator_raw = dict(raw["generator"], model_path=model_path) + generator_raw["engine"] = dict(generator_raw["engine"], execution_backend="external_launcher") + fastvideo_args = generator_config_to_fastvideo_args(normalize_generator_config(generator_raw)) + fastvideo_args.output_type = "latent" + generator = VideoGenerator.from_fastvideo_args(fastvideo_args) + return generator, generator.executor.worker.worker.pipeline + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, help="FastH3 run config (generator + request)") + parser.add_argument("--model-path", required=True) + parser.add_argument("--manifests", nargs="+", required=True) + parser.add_argument("--shard", type=int, default=0) + parser.add_argument("--num-shards", type=int, default=1) + args = parser.parse_args() + + rows = [] + for path in args.manifests: + rows += [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()] + todo = [r for r in rows[args.shard::args.num_shards] if not Path(r["embedding"]).exists()] + print(json.dumps({"rows": len(rows), "todo": len(todo)}), flush=True) + if not todo: + return + os.environ.setdefault("MASTER_PORT", str(29800 + args.shard)) + generator, pipeline = _build(args.config, args.model_path, args.shard) + pipeline.post_init() # idempotent: builds the stages + stage = next(s for s in pipeline._stage_name_mapping.values() if type(s).__name__ == "MiniMaxH3ConditioningStage") + device = torch.device("cuda") + with torch.no_grad(): + for done, row in enumerate(todo, 1): + embeds, tags = stage._encode_fl2va(types.SimpleNamespace(prompt=row["prompt"], extra={}), device) + target = Path(row["embedding"]) + tmp = target.with_suffix(".tmp") + torch.save({"embeds": embeds[0].to(torch.bfloat16).cpu(), "tags": tags.cpu()}, tmp) + tmp.replace(target) + if done % 50 == 0: + print(json.dumps({"done": done, "of": len(todo)}), flush=True) + generator.shutdown() + + +if __name__ == "__main__": + main()