Resolve the FP4 attention kernel and backend conflicts with main's #1980 versions (re-adding install_fp4_vsa_attention), keep main's FP32 FSDP grouping, and keep both save_inference and save_tagged in the checkpointer.
This commit is contained in:
Aryan Kumar
@@ -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"
|
||||
|
||||
@@ -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-<step>/` DCP checkpoints plus a `best/`
|
||||
directory:
|
||||
|
||||
- `best/`: the best eligible state (weights only), refreshed at every eval and
|
||||
never pruned.
|
||||
- `checkpoint-<step>/`: 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.
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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: {}
|
||||
@@ -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=<linears-only 300-step 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
|
||||
@@ -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"]
|
||||
|
||||
@@ -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<block>.<sub>`` 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"
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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": {
|
||||
|
||||
@@ -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] = []
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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"
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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"
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
|
||||
@@ -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 ``<output_dir>/<tag>`` (overwritten).
|
||||
|
||||
Tagged directories are not ``checkpoint-<step>`` 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"
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user