From fed3ce18b63168dec54e3345fd010fd4ece5f958 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 19 Oct 2026 12:38:35 -0700 Subject: [PATCH] [fix]: align the QAD configs and docs with the released provenance - V2 config now matches the canonical 16-GPU stage 1: warmup 10, min LR ratio 0.3, HSDP 16. - State that the published V2.1 checkpoint is linears-only (BF16 VSA attention); the FP4 adaptation is a separate, non-public stage. - Name the released checkpoints and starting models. - Vendor and pin both calibrated amax tables. - Add the V2 prompt/embedding and OmniRef eval data-prep scripts. --- docs/training/minimax_h3_fp4_qad.md | 145 ++++- .../calibration/fasth3/amax-consumer.json | 102 ++++ .../fasth3/amax-omniref-mse-bf16-top8.json | 102 ++++ .../minimax_h3/omniref_nvfp4_qad.yaml | 12 +- .../minimax_h3/v2_nvfp4_qad.yaml | 44 +- .../tests/train/utils/test_checkpoint.py | 9 + .../generate_omniref_latents.py | 529 ++++++++++++++++++ .../minimax_h3_v2_qad/build_v2_prompts.py | 79 +++ .../minimax_h3_v2_qad/precompute_v2_text.py | 70 +++ 9 files changed, 1042 insertions(+), 50 deletions(-) create mode 100644 examples/train/calibration/fasth3/amax-consumer.json create mode 100644 examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json create mode 100644 scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py create mode 100644 scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py create mode 100644 scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py diff --git a/docs/training/minimax_h3_fp4_qad.md b/docs/training/minimax_h3_fp4_qad.md index 3f2657521..0ec3f18c0 100644 --- a/docs/training/minimax_h3_fp4_qad.md +++ b/docs/training/minimax_h3_fp4_qad.md @@ -8,20 +8,45 @@ the FastH3 (MiniMax-H3) family. One recipe covers two model profiles: | OmniRef | `omniref_pdd.MiniMaxH3OmniRefPDDModel` | Ref2VA | 8-forward PDD ladder | 128 | | V2 / V2.1 Consumer | `v2_t2v.MiniMaxH3V2T2VModel` | T2VA | 8-step DMD | 64 | -Both profiles ship the same deployed numerics: NVFP4 W4A4 linears in all 50 -blocks plus FP4 sparse attention. 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. +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. OmniRef's published v2 checkpoint includes FP4 attention; the +published V2.1 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 v2 result | `FastVideo/FastH3-Omni-NVFP4-QAD-v2` | +| V2.1 result | `FastVideo/FastVideo-FastH3-8-Step-V2.1-NVFP4-Consumer` | +| 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 step-3500 OmniRef checkpoint +used as QAD step 0**, not from the released `-v2` result; point `init_from` at +that checkpoint. The V2 config starts from the bf16 V2 snapshot. + +The published V2.1 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 V2.1 checkpoint bit for bit. Set +`models.student.qad.fp4_attention.enabled: false` to get the published +linears-only variant. + ## Target numerics -The student runs exactly the deployed inference numerics from step 0, so the -starting point is the released NVFP4 model rather than a bf16 approximation. +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. @@ -40,9 +65,16 @@ starting point is the released NVFP4 model rather than a bf16 approximation. `proj_out` and the audio projections stay fp32 and frozen, matching inference. -The calibrated `amax` table is produced once per released checkpoint and -shipped with it. `qad.linears.amax_sha256` pins the table; the run copies it -into the output directory and refuses to start if the bytes differ. +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 @@ -80,12 +112,49 @@ in the sampled output. - **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/V2.1 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, `amax` -table, latent manifests, output directory), then launch with `torchrun`. +Fill in the placeholder paths in the config (starting checkpoint, data +manifests, output directory), then launch with `torchrun`. -OmniRef, 4 nodes x 4 GPUs (SP 4): +OmniRef, 4 nodes x 4 GPUs (sequence parallel 4): ```bash torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \ @@ -94,30 +163,34 @@ torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \ --config examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml ``` -V2/V2.1, 4 GPUs: - -```bash -torchrun --nproc_per_node=4 \ - -m fastvideo.train.entrypoint.train \ - --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml -``` - -The V2/V2.1 recipe has an optional lower-learning-rate continuation stage. It -keeps the same numerics and objective and resumes from the first stage's `best` -DCP. The knobs are listed at the bottom of the config: +V2/V2.1, 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 \ - --method.init_student_dcp outputs/v2_1_nvfp4_qad/best \ - --training.optimizer.learning_rate 1.5e-6 \ - --training.distributed.num_gpus 16 \ - --training.distributed.hsdp_shard_dim 16 \ - --training.checkpoint.output_dir outputs/v2_1_nvfp4_qad_continuation + --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml ``` +To reproduce the published V2.1 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 first stage'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 V2.1 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/` @@ -155,6 +228,10 @@ python -m fastvideo.train.entrypoint.dcp_to_diffusers \ | `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 @@ -165,3 +242,13 @@ 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 | +|---|---|---|---| +| V2.1 (linears-only) | mean per-forward x0 rel-L2 | 0.431 | 0.351 (video 0.382, audio 0.319) | +| OmniRef QAD v2 | teacher-forced score | 0.362 | 0.303 at step 120 (video 0.198, audio 0.409) | + +The V2.1 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 index 28f829cbd..d026f8618 100644 --- a/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml +++ b/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml @@ -10,19 +10,21 @@ # See docs/training/minimax_h3_fp4_qad.md for the recipe shared with the V2 # (V2.1 Consumer) profile, calibration, and the checkpoint contract. # -# Paths are placeholders. Fill in the calibrated starting checkpoint, the -# pinned amax table shipped with it, the latent manifests and the output dir. +# Released result: FastVideo/FastH3-Omni-NVFP4-QAD-v2 (step 120 of 200). +# `init_from` is the calibrated step-3500 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 + init_from: data/models/omniref-nvfp4-ptq trainable: true enable_gradient_checkpointing_type: full qad: enabled: true linears: enabled: true - amax_json: data/calibration/amax-omniref-mse-bf16-top8.json + 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 @@ -63,7 +65,7 @@ models: 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 + init_from: data/models/omniref-nvfp4-ptq trainable: false frozen_bf16: true 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 index 8c0a24d1d..046a59927 100644 --- a/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml +++ b/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml @@ -12,24 +12,28 @@ # # 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/FastVideo-FastH3-8-Step-V2.1-NVFP4-Consumer; 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-train + init_from: data/models/FastH3-V2 trainable: true enable_gradient_checkpointing_type: full qad: enabled: true linears: enabled: true - amax_json: data/calibration/amax-consumer.json + 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: null + amax_sha256: 0800516caec04ec829d21c45ca8b18eede6b81008c040997c0b84d48e0d39e52 fp4_attention: enabled: true two_level_p: true @@ -51,7 +55,7 @@ models: 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-train + init_from: data/models/FastH3-V2 trainable: false method: @@ -83,11 +87,11 @@ method: training: distributed: - num_gpus: 4 + num_gpus: 16 sp_size: 1 tp_size: 1 hsdp_replicate_dim: 1 - hsdp_shard_dim: 4 + hsdp_shard_dim: 16 pin_cpu_memory: false data: train_batch_size: 1 @@ -98,8 +102,8 @@ training: eps: 1.0e-08 weight_decay: 0.0 lr_scheduler: cosine_with_min_lr - lr_warmup_steps: 0 - min_lr_ratio: 1.0 + lr_warmup_steps: 10 + min_lr_ratio: 0.3 loop: max_train_steps: 300 gradient_accumulation_steps: 1 @@ -124,14 +128,22 @@ callbacks: pipeline: {} -# Continuation stage (optional). If the released V2.1 checkpoint was refined -# past the 300-step stage above, it was resumed from that stage's `best` DCP -# with the same numerics and a lower learning rate. The stage keeps the full -# recipe above and changes only these knobs: +# The published FastVideo/FastVideo-FastH3-8-Step-V2.1-NVFP4-Consumer +# 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. # -# method.init_student_dcp=outputs/v2_1_nvfp4_qad/best +# 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.distributed.num_gpus=16 -# training.distributed.hsdp_shard_dim=16 -# training.checkpoint.output_dir=outputs/v2_1_nvfp4_qad_continuation +# training.checkpoint.output_dir=outputs/v2_1_nvfp4_qad_fp4_adapt diff --git a/fastvideo/tests/train/utils/test_checkpoint.py b/fastvideo/tests/train/utils/test_checkpoint.py index 1bbab529c..e90dc8721 100644 --- a/fastvideo/tests/train/utils/test_checkpoint.py +++ b/fastvideo/tests/train/utils/test_checkpoint.py @@ -233,6 +233,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/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..5389b0857 --- /dev/null +++ b/scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py @@ -0,0 +1,529 @@ +# 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 \\ + --vae-dir /models/MiniMax-H3/vae --output-dir /data/gen-latents/omniref/data \\ + --per-group 300 --shard 0 --num-shards 4 +""" +from __future__ import annotations + +import argparse +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 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) + manifest.append({ + "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, + }) + tmp_manifest = manifest_path.with_suffix(".tmp") + tmp_manifest.write_text(json.dumps(manifest, indent=1)) + tmp_manifest.replace(manifest_path) + 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()