From 212f558aab211ef8148d744f454cde2b548a25ab Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 19 Oct 2026 12:00:33 -0700 Subject: [PATCH] [feat]: add the shared FastH3 NVFP4 FP4-attention QAD recipe Add the OmniRef (Ref2VA, VSA tile 128) and V2/V2.1 Consumer (T2VA, VSA tile 64) profiles of one quantization-aware distillation recipe: NVFP4 W4A4 q/k/v/out and FFN linears in all 50 blocks, the FP4 sparse-attention emulator, frozen VSA gates and norms, output plus layer-reconstruction regression against a frozen bf16 teacher, and best-by-score checkpoints with audio eligibility guards. Training code: the PDD-QAD method, the two H3 model plugins, the NVFP4 emulated linear, and the tile-64/128 Triton emulator with its PyTorch numeric reference. The VSA attention impl gains a fine attention override hook and lets tile 128 run grad-capable forwards through it. The trainer gains opt-in tagged checkpoints and early stop hooks for the QAD guards. Also add two runnable configs, the training guide, and CPU/GPU tests. --- docs-site/navigation.json | 4 + docs/training/minimax_h3_fp4_qad.md | 166 ++++ .../minimax_h3/omniref_nvfp4_qad.yaml | 133 ++++ .../minimax_h3/v2_nvfp4_qad.yaml | 137 ++++ .../triton_kernels/attn_qat_vsa_reference.py | 317 ++++++++ .../triton_kernels/attn_qat_vsa_train.py | 393 ++++++++++ .../tests/test_attn_qat_vsa_train.py | 304 ++++++++ fastvideo/attention/backends/fp4_vsa_qat.py | 120 +++ .../backends/video_sparse_attn_h3.py | 12 +- fastvideo/layers/quantization/nvfp4_qad.py | 258 ++++++ .../dits/minimax_h3_token_compression.py | 275 +++++++ .../models/test_minimax_h3_omniref_qad.py | 629 +++++++++++++++ .../methods/knowledge_distillation/pdd_qad.py | 736 ++++++++++++++++++ .../pdd_qad_audio_stats.py | 83 ++ .../knowledge_distillation/pdd_qad_metrics.py | 193 +++++ .../knowledge_distillation/pdd_qad_recon.py | 249 ++++++ .../knowledge_distillation/pdd_qad_step0.py | 134 ++++ .../train/models/minimax_h3/omniref_data.py | 295 +++++++ .../train/models/minimax_h3/omniref_pdd.py | 405 ++++++++++ .../train/models/minimax_h3/v2_compression.py | 10 + fastvideo/train/models/minimax_h3/v2_data.py | 94 +++ fastvideo/train/models/minimax_h3/v2_t2v.py | 271 +++++++ fastvideo/train/trainer.py | 18 +- fastvideo/train/utils/checkpoint.py | 38 + 24 files changed, 5270 insertions(+), 4 deletions(-) create mode 100644 docs/training/minimax_h3_fp4_qad.md create mode 100644 examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml create mode 100644 examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml create mode 100644 fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py create mode 100644 fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py create mode 100644 fastvideo-kernel/tests/test_attn_qat_vsa_train.py create mode 100644 fastvideo/attention/backends/fp4_vsa_qat.py create mode 100644 fastvideo/layers/quantization/nvfp4_qad.py create mode 100644 fastvideo/models/dits/minimax_h3_token_compression.py create mode 100644 fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py create mode 100644 fastvideo/train/methods/knowledge_distillation/pdd_qad.py create mode 100644 fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py create mode 100644 fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py create mode 100644 fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py create mode 100644 fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py create mode 100644 fastvideo/train/models/minimax_h3/omniref_data.py create mode 100644 fastvideo/train/models/minimax_h3/omniref_pdd.py create mode 100644 fastvideo/train/models/minimax_h3/v2_compression.py create mode 100644 fastvideo/train/models/minimax_h3/v2_data.py create mode 100644 fastvideo/train/models/minimax_h3/v2_t2v.py diff --git a/docs-site/navigation.json b/docs-site/navigation.json index 06ba8f7cc..83169fb66 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -278,6 +278,10 @@ "label": "Attn-QAT Training", "href": "/training/attn_qat" }, + { + "label": "FastH3 NVFP4 QAD", + "href": "/training/minimax_h3_fp4_qad" + }, { "label": "HunyuanVideo 1.5 T2V", "href": "/training/hunyuan15_overfit" diff --git a/docs/training/minimax_h3_fp4_qad.md b/docs/training/minimax_h3_fp4_qad.md new file mode 100644 index 000000000..1727faebd --- /dev/null +++ b/docs/training/minimax_h3_fp4_qad.md @@ -0,0 +1,166 @@ +# 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 / V2.1 Consumer | `v2_t2v.MiniMaxH3V2T2VModel` | T2VA | 8-step DMD | 64 | + +Both profiles ship the same deployed numerics: NVFP4 W4A4 linears in all 50 +blocks plus FP4 sparse attention. QAD trains the student against a frozen bf16 +teacher so that quantization error is absorbed by the weights instead of +appearing as drift, ghosting or audio loss at inference. + +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` + +## Target numerics + +The student runs exactly the deployed inference numerics from step 0, so the +starting point is the released NVFP4 model rather than a bf16 approximation. + +- **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/V2.1). 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` table is produced once per released checkpoint and +shipped with it. `qad.linears.amax_sha256` pins the table; the run copies it +into the output directory and refuses to start if the bytes differ. + +## 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. + +## Running + +Fill in the placeholder paths in the config (starting checkpoint, `amax` +table, latent manifests, output directory), then launch with `torchrun`. + +OmniRef, 4 nodes x 4 GPUs (SP 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/V2.1, 4 GPUs: + +```bash +torchrun --nproc_per_node=4 \ + -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml +``` + +The V2/V2.1 recipe has an optional lower-learning-rate continuation stage. It +keeps the same numerics and objective and resumes from the first stage's `best` +DCP. The knobs are listed at the bottom of the config: + +```bash +torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \ + --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \ + -m fastvideo.train.entrypoint.train \ + --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml \ + --method.init_student_dcp outputs/v2_1_nvfp4_qad/best \ + --training.optimizer.learning_rate 1.5e-6 \ + --training.distributed.num_gpus 16 \ + --training.distributed.hsdp_shard_dim 16 \ + --training.checkpoint.output_dir outputs/v2_1_nvfp4_qad_continuation +``` + +## Checkpoints + +Runs write regular `checkpoint-/` DCP checkpoints plus a `best/` +directory: + +- `best/`: the best eligible state (weights only), refreshed at every eval and + never pruned. +- `checkpoint-/`: weights, optimizer, RNG state and the method baseline; + the last step is always saved. + +The saved state is the bf16 master plus the quantization plan, not packed NVFP4 +weights; inference applies the same converter used for the released model. Use +`fastvideo/train/entrypoint/dcp_to_diffusers.py` to export a checkpoint if your +deployment path expects Diffusers weights: + +```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/methods/knowledge_distillation/pdd_qad_audio_stats.py` | Audio guard statistics | +| `fastvideo/train/models/minimax_h3/omniref_pdd.py` | OmniRef Ref2VA profile | +| `fastvideo/train/models/minimax_h3/v2_t2v.py` | V2/V2.1 T2VA 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 | + +## 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. diff --git a/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml new file mode 100644 index 000000000..28f829cbd --- /dev/null +++ b/examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml @@ -0,0 +1,133 @@ +# 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 +# (V2.1 Consumer) profile, calibration, and the checkpoint contract. +# +# Paths are placeholders. Fill in the calibrated starting checkpoint, the +# pinned amax table shipped with it, the latent manifests and the output dir. +models: + student: + _target_: fastvideo.train.models.minimax_h3.omniref_pdd.MiniMaxH3OmniRefPDDModel + init_from: data/models/omniref + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: data/calibration/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 + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + rungs_per_row: 4 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/omniref_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/omniref_nvfp4_qad/step0_cache.json + +training: + distributed: + num_gpus: 16 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 16 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 5.0e-06 + betas: [0.9, 0.999] + eps: 1.0e-08 + weight_decay: 0.0 + lr_scheduler: cosine_with_min_lr + lr_warmup_steps: 10 + min_lr_ratio: 0.3 + loop: + max_train_steps: 200 + gradient_accumulation_steps: 1 + checkpoint: + output_dir: outputs/omniref_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: fasth3-omniref-nvfp4 + run_name: omniref-nvfp4-qad + model: + enable_gradient_checkpointing_type: full + dit_precision: fp32 + vsa_sparsity: 0.0 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: {} diff --git a/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml new file mode 100644 index 000000000..8c0a24d1d --- /dev/null +++ b/examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml @@ -0,0 +1,137 @@ +# FastH3 V2 (V2.1 Consumer) 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. +models: + student: + _target_: fastvideo.train.models.minimax_h3.v2_t2v.MiniMaxH3V2T2VModel + init_from: data/models/FastH3-V2-train + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: data/calibration/amax-consumer.json + # Attention and gate linears use the unit activation scale; only the + # FFN linears read their calibrated amax from this table. + attention: true + gate: true + ffn: true + skip_blocks: [] + amax_sha256: null + 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-train + 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: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + 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: 0 + min_lr_ratio: 1.0 + 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: {} + +# Continuation stage (optional). If the released V2.1 checkpoint was refined +# past the 300-step stage above, it was resumed from that stage's `best` DCP +# with the same numerics and a lower learning rate. The stage keeps the full +# recipe above and changes only these knobs: +# +# method.init_student_dcp=outputs/v2_1_nvfp4_qad/best +# training.optimizer.learning_rate=1.5e-06 +# training.loop.max_train_steps=300 +# training.distributed.num_gpus=16 +# training.distributed.hsdp_shard_dim=16 +# training.checkpoint.output_dir=outputs/v2_1_nvfp4_qad_continuation diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py new file mode 100644 index 000000000..a6b6b6f44 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Slow PyTorch reference for block-sparse FP4 attention as ``attn_qat_infer`` computes it. + +This is the numeric specification that ``attn_qat_vsa_train`` (the Triton +fake-quant training kernel) must reproduce. It follows the deployed sm_120/121 +kernel (``sageattn_blackwell_sparse_bshd`` on lists from +``vsa_tile_mask_to_fp4_blocks``) step by step: + +* Q and K: NVFP4 E2M1 with one E4M3 scale per 16 channels along D, no global + scale (``scaled_fp4_quant(x, ..., 1)``): ``s = e4m3(amax / 6)``, value + ``e2m1(x / s) * s``, round-to-nearest-even, saturating. +* V: the same, but grouped along tokens (``scale_and_quant_fp4_transpose``). +* Each 128-row query block visits its listed 128-key blocks in ascending block + order (the lists are stored descending and read from the last entry). +* P: the online softmax of ``softmax_fused.h``. For each visited block the row + max is updated first, ``P~ = exp2(S * log2e * scale - m_scaled)``; per 16 keys + ``s_P = max(P~) / 6`` is computed in fp32, ``P~ / s_P`` is rounded to E2M1 and + the MMA multiplies by ``e4m3(s_P)``. Single-level P has ``m_scaled = m``; + two-level P adds ``log2(1 / (448 * 6))`` so ``P~`` spans ``[0, 2688]``. +* The row sum uses the unquantized ``P~``; the output is ``acc / l``. +* Masked keys (``kv_valid`` halves) score ``-inf``; a fully masked 16-key group + has ``s_P = 0`` and contributes exactly zero. + +``first_block_max_floor`` models an uninitialized-register hazard in the +deployed kernel: on the first visited block, ``AbsMaxP`` is read before it is +written (``fmaxf(AbsMaxP, s)``). ``None`` assumes the compiler folds that to +the true max; a number (e.g. ``0.0``) floors the first block's group maxima +and hence the row max. The cross-device fidelity check decides which applies. +""" +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +QUANT_GROUP = 16 +BLOCK = 128 +HALF = BLOCK // 2 +E4M3_MAX = 448.0 +LOG2E = 1.4426950408889634 +LOG2_INV_FP4_MAX = -2.584962500721156 # log2(1 / 6) +LOG2_INV_TWO_LEVEL = -11.392317422778762 # log2(1 / (448 * 6)) + + +def e4m3_rn(x: torch.Tensor) -> torch.Tensor: + """Round non-negative fp32 scales to E4M3 (RNE, saturating at 448) and back to fp32.""" + return x.clamp(max=E4M3_MAX).to(torch.float8_e4m3fn).to(torch.float32) + + +def e2m1_rn(y: torch.Tensor) -> torch.Tensor: + """Round fp32 values to the E2M1 grid {0, .5, 1, 1.5, 2, 3, 4, 6} (RNE, saturating).""" + a = y.abs() + q = torch.where( + a <= 0.25, 0.0, + torch.where( + a < 0.75, 0.5, + torch.where( + a <= 1.25, 1.0, + torch.where(a < 1.75, 1.5, + torch.where(a <= 2.5, 2.0, torch.where(a < 3.5, 3.0, torch.where(a <= 5.0, 4.0, + 6.0))))))) + return torch.copysign(q, y) + + +def nvfp4_fake_quant_lastdim(x: torch.Tensor) -> torch.Tensor: + """Fake-quantize along the last dim in groups of 16 (no global scale); returns fp32.""" + if x.shape[-1] % QUANT_GROUP: + raise ValueError(f"last dim {x.shape[-1]} is not a multiple of {QUANT_GROUP}") + xf = x.float().unflatten(-1, (-1, QUANT_GROUP)) + scale = e4m3_rn(xf.abs().amax(dim=-1, keepdim=True) / 6.0) + inv = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) + return (e2m1_rn(xf * inv) * scale).flatten(-2) + + +def fake_quant_qkv(q: torch.Tensor, k: torch.Tensor, + v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``[B, H, L, D]`` inputs -> dequantized Q (along D), K (along D), V (along L); fp32.""" + v_hat = nvfp4_fake_quant_lastdim(v.transpose(-1, -2)).transpose(-1, -2) + return nvfp4_fake_quant_lastdim(q), nvfp4_fake_quant_lastdim(k), v_hat + + +def smooth_keys(k: torch.Tensor) -> torch.Tensor: + """K smoothing as ``preprocess_qkv`` (``k -= k.mean(dim=-2)``): per-head token mean, in the input dtype. + + The dtype matters: FP4 rounding is so input-sensitive that smoothing in + fp32 instead of bf16 alone moves the attention output by ~8% rel-L2. The + mean runs over every row given (tile-ordered, pad rows included); a + deployed K-smoothing path must use the same rows. + """ + return k - k.mean(dim=-2, keepdim=True) + + +def key_valid_mask(kv_valid: torch.Tensor | None, n_blocks: int, device: torch.device) -> torch.Tensor: + """``[n_blocks, 128]`` bool: which key columns of each block are real tokens.""" + if kv_valid is None: + return torch.ones(n_blocks, BLOCK, dtype=torch.bool, device=device) + col = torch.arange(BLOCK, device=device) + limit = torch.where(col < HALF, kv_valid[0::2, None], kv_valid[1::2, None]) + return (col % HALF)[None, :] < limit + + +def tile_mask_to_fp4_blocks(tile_mask: torch.Tensor, + tile_tokens: int, + tile_valid: torch.Tensor | None = None) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Pure-PyTorch ``attn_qat_infer.api.vsa_tile_mask_to_fp4_blocks`` for tile 64 or 128. + + The inference API imports sm_12x extensions at module load, so data center + GPUs cannot call it. The kernel works on 128x128 blocks. With 64-token VSA + tiles two query tiles and two key tiles are paired into one block; a block + is listed when any of its four 64x64 quadrants is selected and ``q2k_quad`` + carries the selected quadrants (bit ``2 * query_half + key_half``). With + 128-token tiles the mapping is 1:1 and no quadrant mask is needed. + + Returns ``(q2k_idx [B, H, M, N], q2k_num [B, H, M], kv_valid [2N], + q2k_quad [B, H, M, N] | None)``; queries cover ``M * 128`` padded rows. + """ + if tile_tokens % HALF or (tile_tokens > HALF and tile_tokens % BLOCK): + raise ValueError(f"tile_tokens={tile_tokens} must be {HALF} or a multiple of {BLOCK}") + batch, heads, n_tiles, _ = tile_mask.shape + device = tile_mask.device + halves_per_tile = tile_tokens // HALF + half_mask = tile_mask + if halves_per_tile > 1: + half_mask = half_mask.repeat_interleave(halves_per_tile, dim=2).repeat_interleave(halves_per_tile, dim=3) + n_halves = n_tiles * halves_per_tile + if tile_valid is None: + half_valid = torch.full((n_halves, ), HALF, device=device, dtype=torch.int32) + else: + offsets = torch.arange(halves_per_tile, device=device, dtype=torch.int32) * HALF + half_valid = (tile_valid.to(torch.int32)[:, None] - offsets[None, :]).clamp(0, HALF).reshape(-1) + if n_halves % 2: + # Pad to whole 128-blocks: the extra key half is empty, the extra query + # half (discarded output) attends block 0 only. + half_mask = F.pad(half_mask, (0, 1, 0, 1), value=False) + half_mask[:, :, -1, 0] = True + half_valid = F.pad(half_valid, (0, 1), value=0) + n_halves += 1 + half_mask = half_mask & (half_valid > 0)[None, None, None, :] + n_blocks = n_halves // 2 + quads = half_mask.view(batch, heads, n_blocks, 2, n_blocks, 2) + weights = torch.tensor([[1, 2], [4, 8]], device=device, dtype=torch.uint8) # [row_half, col_half] + quad = (quads.to(torch.uint8) * weights[None, None, None, :, None, :]).sum(dim=(3, 5), dtype=torch.uint8) + block_mask = quad != 0 + # Descending compaction without a sort: position = running count from the right. + rev = block_mask.flip(-1) + pos = rev.cumsum(-1, dtype=torch.int32) - 1 + q2k_num = (pos[..., -1] + 1).contiguous() + cols = torch.arange(n_blocks - 1, -1, -1, device=device, dtype=torch.int32).expand_as(pos) + slot = torch.where(rev, pos, torch.full_like(pos, n_blocks)).long() + q2k_idx = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.int32) + q2k_idx.scatter_(-1, slot, cols) + q2k_idx = q2k_idx[..., :n_blocks].contiguous() + q2k_quad = None + if halves_per_tile == 1: + q2k_quad = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.uint8) + q2k_quad.scatter_(-1, slot, quad.flip(-1)) + q2k_quad = q2k_quad[..., :n_blocks].contiguous() + return q2k_idx, q2k_num, half_valid.contiguous(), q2k_quad + + +def tile128_mask_to_fp4_blocks(tile_mask: torch.Tensor, + tile_valid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward-compatible 128-token-tile wrapper.""" + idx, num, valid, quad = tile_mask_to_fp4_blocks(tile_mask, BLOCK, tile_valid) + if quad is not None: + raise AssertionError("128-token tiles must not produce a quadrant mask") + return idx, num, valid + + +def invert_block_lists(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, + n_kblocks: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Query->key lists to key->query lists, keeping each entry's slot in its query list. + + Returns ``(k2q_idx [B, H, NK, M], k2q_slot [B, H, NK, M], k2q_num [B, H, NK])`` (int32). + """ + batch, heads, n_qblocks, width = q2k_idx.shape + device = q2k_idx.device + slots = torch.arange(width, device=device, dtype=torch.int32).expand_as(q2k_idx) + listed = slots < q2k_num[..., None] + slot_of = torch.full((batch, heads, n_qblocks, n_kblocks + 1), -1, device=device, dtype=torch.int32) + slot_of.scatter_(-1, torch.where(listed, q2k_idx.long(), n_kblocks), torch.where(listed, slots, -1)) + by_key = slot_of[..., :n_kblocks].transpose(-1, -2) # [B, H, NK, M] + present = by_key >= 0 + k2q_num = present.sum(-1, dtype=torch.int32) + position = torch.where(present, present.cumsum(-1) - 1, n_qblocks).long() + queries = torch.arange(n_qblocks, device=device, dtype=torch.int32).expand_as(by_key) + k2q_idx = torch.zeros((batch, heads, n_kblocks, n_qblocks + 1), device=device, dtype=torch.int32) + k2q_slot = torch.zeros_like(k2q_idx) + k2q_idx.scatter_(-1, position, queries) + k2q_slot.scatter_(-1, position, by_key) + return k2q_idx[..., :n_qblocks].contiguous(), k2q_slot[..., :n_qblocks].contiguous(), k2q_num.contiguous() + + +def visit_order(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, b: int, h: int, m: int) -> list[int]: + """Blocks query block ``m`` visits, in kernel order (last list entry first).""" + count = int(q2k_num[b, h, m]) + return [int(block) for block in q2k_idx[b, h, m, :count].flip(0)] + + +@torch.no_grad() +def fp4_vsa_attention_reference(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + q2k_quad: torch.Tensor | None = None, + *, + two_level_p: bool = False, + smooth_k: bool = False, + quantize: bool = True, + first_block_max_floor: float | None = None, + sm_scale: float | None = None, + return_dv_weights: bool = False): + """Emulate the deployed block-sparse FP4 attention on ``[B, H, L, D]`` inputs (L a multiple of 128). + + ``quantize=False`` runs the same visit order with exact softmax numerics + (BF16-input reference). Returns ``[B, H, L, D]`` fp32, and with + ``return_dv_weights`` also the dense ``[B, H, L, L]`` matrix ``W`` with + ``out = W @ v_hat`` (the forward-consistent dV weights). + """ + batch, heads, seq_len, dim = q.shape + if seq_len % BLOCK: + raise ValueError(f"sequence length {seq_len} must be a multiple of {BLOCK}") + scale_log2 = (sm_scale if sm_scale is not None else 1.0 / math.sqrt(dim)) * LOG2E + if smooth_k: + k = smooth_keys(k) + if quantize: + q_hat, k_hat, v_hat = fake_quant_qkv(q, k, v) + else: + q_hat, k_hat, v_hat = q.float(), k.float(), v.float() + offset = LOG2_INV_TWO_LEVEL if two_level_p else 0.0 + n_blocks = seq_len // BLOCK + valid = key_valid_mask(kv_valid, n_blocks, q.device) + out = torch.zeros(batch, heads, seq_len, dim, dtype=torch.float32, device=q.device) + weights = torch.zeros(batch, heads, seq_len, seq_len, device=q.device) if return_dv_weights else None + for b in range(batch): + for h in range(heads): + for m in range(n_blocks): + rows = slice(m * BLOCK, (m + 1) * BLOCK) + row_max = torch.full((BLOCK, ), -math.inf, device=q.device) + row_sum = torch.zeros(BLOCK, device=q.device) + acc = torch.zeros(BLOCK, dim, device=q.device) + visited: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + count = int(q2k_num[b, h, m]) + entries = list(range(count))[::-1] + for position, entry in enumerate(entries): + n = int(q2k_idx[b, h, m, entry]) + cols = slice(n * BLOCK, (n + 1) * BLOCK) + scores = q_hat[b, h, rows] @ k_hat[b, h, cols].T + scores = scores.masked_fill(~valid[n][None, :], -math.inf) + if q2k_quad is not None: + quad = int(q2k_quad[b, h, m, entry]) + half = torch.arange(BLOCK, device=q.device) // HALF + qvalid = ((quad >> (half[:, None] * 2 + half[None, :])) & 1).bool() + scores = scores.masked_fill(~qvalid, -math.inf) + group_max = scores.unflatten(-1, (-1, QUANT_GROUP)).amax(dim=-1) + if position == 0 and first_block_max_floor is not None: + group_max = group_max.clamp(min=first_block_max_floor) + new_max = torch.maximum(row_max, group_max.amax(dim=-1)) + m_scaled = new_max * scale_log2 + offset + p_tilde = torch.exp2(scores * scale_log2 - m_scaled[:, None]) + alpha = torch.exp2((row_max - new_max) * scale_log2) + alpha = torch.where(torch.isinf(row_max), torch.zeros_like(alpha), alpha) + row_sum = row_sum * alpha + p_tilde.sum(dim=-1) + if quantize: + s_p = torch.exp2(group_max * scale_log2 - m_scaled[:, None] + LOG2_INV_FP4_MAX) + s_full = s_p.repeat_interleave(QUANT_GROUP, dim=-1) + ratio = torch.where(s_full == 0, torch.zeros_like(p_tilde), p_tilde / s_full) + p_hat = e2m1_rn(ratio) * e4m3_rn(s_full) + else: + p_hat = p_tilde + acc = acc * alpha[:, None] + p_hat @ v_hat[b, h, cols] + visited.append((n, p_hat, new_max)) + row_max = new_max + inv = torch.where((row_sum == 0) | torch.isnan(row_sum), torch.zeros_like(row_sum), 1.0 / row_sum) + out[b, h, rows] = acc * inv[:, None] + if weights is not None: + for n, p_hat, block_max in visited: + decay = torch.exp2((block_max - row_max) * scale_log2) * inv + weights[b, h, rows, n * BLOCK:(n + 1) * BLOCK] = p_hat * decay[:, None] + if weights is not None: + return out, weights + return out + + +def dense_block_mask(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, seq_len: int, + kv_valid: torch.Tensor | None = None) -> torch.Tensor: + """Token-level ``[B, H, L, L]`` bool mask of the listed blocks (and valid keys).""" + batch, heads, n_blocks, width = q2k_idx.shape + listed = torch.arange(width, device=q2k_idx.device) < q2k_num[..., None] + block_mask = torch.zeros(batch, heads, n_blocks, n_blocks + 1, dtype=torch.bool, device=q2k_idx.device) + block_mask.scatter_(-1, torch.where(listed, q2k_idx.long(), n_blocks), True) + block_mask = block_mask[..., :n_blocks] + token = block_mask.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :seq_len, :seq_len] + valid = key_valid_mask(kv_valid, n_blocks, q2k_idx.device).reshape(-1)[:seq_len] + return token & valid + + +def ste_gradients_reference(q_hat: torch.Tensor, k_hat: torch.Tensor, v_hat: torch.Tensor, token_mask: torch.Tensor, + dv_weights: torch.Tensor, grad_out: torch.Tensor, + sm_scale: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The training kernel's STE gradients, computed densely in fp32. + + dQ and dK are the exact gradients of the high-precision surrogate + ``softmax(q_hat k_hat^T) v_hat`` on the sparse mask; dV uses the forward's + quantized, online-softmax weights (forward-consistent dV). + """ + q_hat, k_hat, v_hat = (t.detach().float().requires_grad_(True) for t in (q_hat, k_hat, v_hat)) + scores = (q_hat @ k_hat.transpose(-1, -2)) * sm_scale + probs = torch.softmax(scores.masked_fill(~token_mask, -math.inf), dim=-1) + surrogate = probs @ v_hat + dq, dk = torch.autograd.grad(surrogate, (q_hat, k_hat), grad_out.float()) + dv = dv_weights.transpose(-1, -2) @ grad_out.float() + return dq, dk, dv diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py new file mode 100644 index 000000000..f540b3091 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse FP4 Attn-QAT for VSA tile-128 training (forward + backward). + +The forward reproduces the deployed sm_120/121 sparse FP4 kernel +(``attn_qat_infer.api.sageattn_blackwell_sparse_bshd`` on lists from +``vsa_tile_mask_to_fp4_blocks(mask, 128, vbs)``) on any CUDA GPU with Triton, +so a GB200 can train against Spark/5090 numerics. The numeric contract is +spelled out in ``attn_qat_vsa_reference`` (the slow PyTorch specification the +unit tests compare against): + +* Q/K/V are fake-quantized once per call (NVFP4, per-16 E4M3 scales, no global + scale; V grouped along tokens) and stored in BF16, which holds every + dequantized value exactly. +* Every query row visits its block's listed 128-key blocks in ascending order; + the running max is updated per 128-key block before P is formed, P is + quantized per 16 keys (single- or two-level), the row sum is unquantized. + +Backward is a straight-through estimator with one deliberate choice per +operand (the lessons of the GB10 Attn-QAT dV fix): + +* dV is forward-consistent: it is built from the forward's quantized + online-softmax weights, recomputed from the saved per-(row, visited block) + running maxima, never by re-quantizing the normalized probabilities (whose + E4M3 group scales underflow on long sequences and silence whole heads). +* dQ/dK are the exact gradients of the high-precision surrogate + ``softmax(q_hat k_hat^T) v_hat`` on the same sparse mask; delta uses the + high-precision output. +* The Q/K/V quantizers pass gradients through unchanged. + +Rows are independent, so programs cover 64 query rows (the forward keeps whole +128-key blocks, where the running max is defined); the backward splits keys +into 64-column halves, which the per-16 P groups allow. +""" +from __future__ import annotations + +import math + +import os + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from .attn_qat_vsa_reference import BLOCK, LOG2_INV_TWO_LEVEL, LOG2E, invert_block_lists + +_ROWS = 64 # query rows per program +_HALF = BLOCK // 2 +_QUANT_ROWS = 64 # fake-quant rows per program (a multiple of 16; L is a multiple of 128) +_WARPS = int(os.environ.get("FASTVIDEO_QAT_VSA_WARPS", "8")) # 8 beats 4 by 15-20% on GB200 at 35k-73k tokens + + +@triton.jit +def _e2m1_rn(y): + """Round to the E2M1 grid {0, .5, 1, 1.5, 2, 3, 4, 6}, ties to even, saturating at 6. + + The grid is 0.5-spaced below 2, 1-spaced in [2, 4) and 2-spaced above, so + round-half-even on the scaled value is exact (ties land on even codes). + """ + a = tl.minimum(tl.abs(y), 6.0) + q = tl.where(a < 2.0, + libdevice.rint(a * 2.0) * 0.5, + tl.where(a < 4.0, libdevice.rint(a), + libdevice.rint(a * 0.5) * 2.0)) + return tl.where(y < 0, -q, q) + + +@triton.jit +def _e4m3_rn(x): + """Round non-negative fp32 to E4M3 (RN, saturating at 448) and back.""" + return tl.minimum(x, 448.0).to(tl.float8e4nv).to(tl.float32) + + +@triton.jit +def _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, + ROWS: tl.constexpr, COLS: tl.constexpr): + """``P~`` and its fake-quantized value for raw scores ``s`` (masked = -inf) under row maxima ``m_scaled``.""" + p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None]) + group_max = tl.max(tl.reshape(s, (ROWS, COLS // 16, 16)), axis=2) + if FLOOR: + group_max = tl.where(first_visit, tl.maximum(group_max, FLOOR_VALUE), group_max) + s_p = tl.math.exp2(group_max * scale_log2 - m_scaled[:, None] - 2.584962500721156) + # Per-group values (fp32 divisor, E4M3 multiplier) broadcast back over each group's 16 keys. + s_full = tl.reshape(tl.broadcast_to(s_p[:, :, None], (ROWS, COLS // 16, 16)), (ROWS, COLS)) + s_e4m3 = tl.reshape(tl.broadcast_to(_e4m3_rn(s_p)[:, :, None], (ROWS, COLS // 16, 16)), (ROWS, COLS)) + ratio = tl.where(s_full == 0.0, 0.0, p_tilde / s_full) + return p_tilde, _e2m1_rn(ratio) * s_e4m3 + + +@triton.jit +def _fake_quant_kernel(X, Y, n_rows, D: tl.constexpr, ROWS: tl.constexpr, ALONG_ROWS: tl.constexpr): + """NVFP4 fake quant of a ``[n_rows, D]`` matrix, groups of 16 along D or (``ALONG_ROWS``) along rows.""" + pid = tl.program_id(0).to(tl.int64) + offs_r = pid * ROWS + tl.arange(0, ROWS) + offs_d = tl.arange(0, D) + x = tl.load(X + offs_r[:, None] * D + offs_d[None, :]).to(tl.float32) + if ALONG_ROWS: + x3 = tl.reshape(x, (ROWS // 16, 16, D)) + amax = tl.max(tl.abs(x3), axis=1) + scale = _e4m3_rn(amax / 6.0) + inv = tl.where(scale == 0.0, 0.0, 1.0 / scale) + y3 = _e2m1_rn(x3 * inv[:, None, :]) * scale[:, None, :] + y = tl.reshape(y3, (ROWS, D)) + else: + x3 = tl.reshape(x, (ROWS, D // 16, 16)) + amax = tl.max(tl.abs(x3), axis=2) + scale = _e4m3_rn(amax / 6.0) + inv = tl.where(scale == 0.0, 0.0, 1.0 / scale) + y3 = _e2m1_rn(x3 * inv[:, :, None]) * scale[:, :, None] + y = tl.reshape(y3, (ROWS, D)) + tl.store(Y + offs_r[:, None] * D + offs_d[None, :], y.to(Y.dtype.element_ty)) + + +def fake_quant_qkv_triton(q: torch.Tensor, k: torch.Tensor, + v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fake_quant_qkv`` (Q/K along D, V along tokens) as one Triton pass each; BF16 ``[B, H, L, D]`` out.""" + outs = [] + for x, along_rows in ((q, False), (k, False), (v, True)): + x = x.contiguous() + y = torch.empty(x.shape, device=x.device, dtype=torch.bfloat16) + rows = x.numel() // x.shape[-1] + _fake_quant_kernel[(rows // _QUANT_ROWS, )](x, y, rows, D=x.shape[-1], ROWS=_QUANT_ROWS, + ALONG_ROWS=along_rows, num_warps=4) + outs.append(y) + return outs[0], outs[1], outs[2] + + +@triton.jit +def _key_ok(KV_VALID, n_block, cols, HAS_VALID: tl.constexpr, HALF: tl.constexpr): + """Valid key columns (block-local ``cols`` in [0, 128)) of 128-key block ``n_block``.""" + if HAS_VALID: + valid0 = tl.load(KV_VALID + 2 * n_block) + valid1 = tl.load(KV_VALID + 2 * n_block + 1) + limit = tl.where(cols < HALF, valid0, valid1) + return (cols % HALF) < limit + return cols >= 0 + + +@triton.jit +def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD, seq_len, n_qblocks, list_width, + scale_log2, OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr, + HAS_QUAD: tl.constexpr, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, HIGH_PREC_O: tl.constexpr, + D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + m_block = pid // halves + row0 = (pid % halves) * ROWS + offs_r = tl.arange(0, ROWS) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + base = bh * seq_len * D + rows = m_block * BLOCK_N + row0 + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + meta = bh * n_qblocks + m_block + count = tl.load(Q2K_NUM + meta) + m_i = tl.full((ROWS, ), float("-inf"), tl.float32) + l_i = tl.zeros((ROWS, ), tl.float32) + acc = tl.zeros((ROWS, D), tl.float32) + acc_hp = tl.zeros((ROWS, D), tl.float32) + for it in range(0, count): + slot = count - 1 - it + n_block = tl.load(Q2K_IDX + meta * list_width + slot) + keys = n_block * BLOCK_N + offs_n + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + s = tl.dot(q, tl.trans(k)) + ok = _key_ok(KV_VALID, n_block, offs_n, HAS_VALID, BLOCK_N // 2) + if HAS_QUAD: + quad = tl.load(Q2K_QUAD + meta * list_width + slot).to(tl.int32) + qh = pid % halves + bit = tl.where(offs_n < ROWS, (quad >> (qh * 2)) & 1, (quad >> (qh * 2 + 1)) & 1) + ok = ok & (bit != 0) + s = tl.where(ok[None, :], s, float("-inf")) + group_max = tl.max(tl.reshape(s, (ROWS, BLOCK_N // 16, 16)), axis=2) + if FLOOR: + group_max = tl.where(it == 0, tl.maximum(group_max, FLOOR_VALUE), group_max) + m_new = tl.maximum(m_i, tl.max(group_max, axis=1)) + m_scaled = m_new * scale_log2 + OFFSET + alpha = tl.math.exp2((m_i - m_new) * scale_log2) + if QUANT: + p_tilde, p_hat = _quantized_p(s, m_scaled, it == 0, scale_log2, FLOOR, FLOOR_VALUE, ROWS, BLOCK_N) + else: + p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None]) + p_hat = p_tilde + l_i = l_i * alpha + tl.sum(p_tilde, axis=1) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + acc = acc * alpha[:, None] + tl.dot(p_hat.to(tl.bfloat16), v) + if HIGH_PREC_O: + acc_hp = acc_hp * alpha[:, None] + tl.dot(p_tilde.to(tl.bfloat16), v) + tl.store(MRUN + (meta * list_width + slot) * BLOCK_N + row0 + offs_r, m_new) + m_i = m_new + inv = tl.where((l_i == 0.0) | (l_i != l_i), 0.0, 1.0 / l_i) + tl.store(O + base + rows[:, None] * D + offs_d[None, :], (acc * inv[:, None]).to(O.dtype.element_ty)) + if HIGH_PREC_O: + tl.store(OHP + base + rows[:, None] * D + offs_d[None, :], (acc_hp * inv[:, None]).to(OHP.dtype.element_ty)) + tl.store(LSE + bh * seq_len + rows, m_i * scale_log2 + OFFSET + tl.math.log2(l_i)) + + +@triton.jit +def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD, seq_len, n_qblocks, + list_width, scale_log2, sm_scale, HAS_VALID: tl.constexpr, HAS_QUAD: tl.constexpr, + D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + m_block = pid // halves + offs_r = tl.arange(0, ROWS) + offs_c = tl.arange(0, ROWS) # key columns per step: one 64-key half + offs_d = tl.arange(0, D) + base = bh * seq_len * D + rows = m_block * BLOCK_N + (pid % halves) * ROWS + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + do = tl.load(DO + base + rows[:, None] * D + offs_d[None, :]) + lse = tl.load(LSE + bh * seq_len + rows) + delta = tl.load(DELTA + bh * seq_len + rows) + meta = bh * n_qblocks + m_block + count = tl.load(Q2K_NUM + meta) + dq = tl.zeros((ROWS, D), tl.float32) + for it in range(0, count * halves): + n_block = tl.load(Q2K_IDX + meta * list_width + it // halves) + cols = (it % halves) * ROWS + offs_c + keys = n_block * BLOCK_N + cols + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + s = tl.dot(q, tl.trans(k)) + ok = _key_ok(KV_VALID, n_block, cols, HAS_VALID, BLOCK_N // 2) + if HAS_QUAD: + quad = tl.load(Q2K_QUAD + meta * list_width + (it // halves)).to(tl.int32) + bit = (quad >> ((pid % halves) * 2 + (it % halves))) & 1 + ok = ok & (bit != 0) + p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse[:, None]), 0.0) + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) + dq += tl.dot(ds.to(tl.bfloat16), k) + tl.store(DQ + base + rows[:, None] * D + offs_d[None, :], (dq * sm_scale).to(DQ.dtype.element_ty)) + + +@triton.jit +def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K2Q_NUM, Q2K_NUM, KV_VALID, + Q2K_QUAD, seq_len, n_qblocks, n_kblocks, list_width, scale_log2, sm_scale, + OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr, HAS_QUAD: tl.constexpr, + FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr, + BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + n_block = pid // halves + offs_r = tl.arange(0, ROWS) + offs_d = tl.arange(0, D) + cols = (pid % halves) * ROWS + tl.arange(0, ROWS) + base = bh * seq_len * D + keys = n_block * BLOCK_N + cols + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + ok = _key_ok(KV_VALID, n_block, cols, HAS_VALID, BLOCK_N // 2) + kmeta = bh * n_kblocks + n_block + count = tl.load(K2Q_NUM + kmeta) + dk = tl.zeros((ROWS, D), tl.float32) + dv = tl.zeros((ROWS, D), tl.float32) + for it in range(0, count * halves): + entry = kmeta * n_qblocks + it // halves + m_block = tl.load(K2Q_IDX + entry) + slot = tl.load(K2Q_SLOT + entry) + row0 = (it % halves) * ROWS + rows = m_block * BLOCK_N + row0 + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + do = tl.load(DO + base + rows[:, None] * D + offs_d[None, :]) + lse = tl.load(LSE + bh * seq_len + rows) + delta = tl.load(DELTA + bh * seq_len + rows) + s = tl.dot(q, tl.trans(k)) + if HAS_QUAD: + quad = tl.load(Q2K_QUAD + (bh * n_qblocks + m_block) * list_width + slot).to(tl.int32) + bit = (quad >> ((it % halves) * 2 + (pid % halves))) & 1 + s = tl.where(ok[None, :] & (bit != 0), s, float("-inf")) + else: + s = tl.where(ok[None, :], s, float("-inf")) + p = tl.math.exp2(s * scale_log2 - lse[:, None]) + if QUANT: + qmeta = bh * n_qblocks + m_block + first_visit = slot == tl.load(Q2K_NUM + qmeta) - 1 + m_run = tl.load(MRUN + (qmeta * list_width + slot) * BLOCK_N + row0 + offs_r) + m_scaled = m_run * scale_log2 + OFFSET + _, p_hat = _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR, FLOOR_VALUE, ROWS, ROWS) + weight = p_hat * tl.math.exp2(m_scaled - lse)[:, None] + else: + weight = p + dv += tl.dot(tl.trans(weight.to(tl.bfloat16)), do) + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) + dk += tl.dot(tl.trans(ds.to(tl.bfloat16)), q) + tl.store(DK + base + keys[:, None] * D + offs_d[None, :], (dk * sm_scale).to(DK.dtype.element_ty)) + tl.store(DV + base + keys[:, None] * D + offs_d[None, :], dv.to(DV.dtype.element_ty)) + + +class _FP4VSAAttnQAT(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize, two_level_p, first_block_max_floor, + sm_scale, high_prec_o): # type: ignore[override] + batch, heads, seq_len, dim = q.shape + if seq_len % BLOCK or dim not in (64, 128): + raise ValueError(f"need L % {BLOCK} == 0 and D in (64, 128), got L={seq_len}, D={dim}") + n_blocks = seq_len // BLOCK + if q2k_idx.shape[:3] != (batch, heads, n_blocks): + raise ValueError(f"q2k_idx {tuple(q2k_idx.shape)} does not match {(batch, heads, n_blocks)}") + if quantize: + q_hat, k_hat, v_hat = fake_quant_qkv_triton(q, k, v) + else: + q_hat, k_hat, v_hat = (t.to(torch.bfloat16).contiguous() for t in (q, k, v)) + q2k_idx = q2k_idx.to(torch.int32).contiguous() + q2k_num = q2k_num.to(torch.int32).contiguous() + has_valid = kv_valid is not None + kv_valid_arg = kv_valid.to(torch.int32).contiguous() if has_valid else q2k_num + has_quad = q2k_quad is not None + q2k_quad_arg = q2k_quad.to(torch.uint8).contiguous() if has_quad else q2k_idx + width = q2k_idx.shape[-1] + out = torch.empty_like(q_hat) + out_hp = torch.empty_like(q_hat) if high_prec_o else out + lse = torch.empty((batch, heads, seq_len), device=q.device, dtype=torch.float32) + m_run = torch.empty((batch, heads, n_blocks, width, BLOCK), device=q.device, dtype=torch.float32) + scale_log2 = sm_scale * LOG2E + offset = LOG2_INV_TWO_LEVEL if two_level_p else 0.0 + floor = first_block_max_floor is not None + grid = (n_blocks * (BLOCK // _ROWS), batch * heads) + _fwd_kernel[grid](q_hat, k_hat, v_hat, out, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg, + q2k_quad_arg, seq_len, n_blocks, width, scale_log2, OFFSET=offset, QUANT=bool(quantize), + HAS_VALID=has_valid, HAS_QUAD=has_quad, FLOOR=floor, + FLOOR_VALUE=float(first_block_max_floor or 0.0), + HIGH_PREC_O=bool(high_prec_o), D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS, + num_stages=2) + ctx.save_for_backward(q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg, + q2k_quad_arg) + ctx.options = (bool(quantize), has_valid, has_quad, floor, float(first_block_max_floor or 0.0), offset, + sm_scale) + return out + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid, q2k_quad = ctx.saved_tensors + quantize, has_valid, has_quad, floor, floor_value, offset, sm_scale = ctx.options + batch, heads, seq_len, dim = q_hat.shape + n_blocks = seq_len // BLOCK + width = q2k_idx.shape[-1] + grad_out = grad_out.to(torch.bfloat16).contiguous() + delta = (grad_out.float() * out_hp.float()).sum(-1).contiguous() + k2q_idx, k2q_slot, k2q_num = invert_block_lists(q2k_idx, q2k_num, n_blocks) + dq = torch.empty_like(q_hat) + dk = torch.empty_like(k_hat) + dv = torch.empty_like(v_hat) + scale_log2 = sm_scale * LOG2E + grid = (n_blocks * (BLOCK // _ROWS), batch * heads) + _bwd_dq_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, dq, q2k_idx, q2k_num, kv_valid, q2k_quad, + seq_len, n_blocks, width, scale_log2, sm_scale, HAS_VALID=has_valid, + HAS_QUAD=has_quad, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS, num_stages=2) + _bwd_dkdv_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, m_run, dk, dv, k2q_idx, k2q_slot, k2q_num, + q2k_num, kv_valid, q2k_quad, seq_len, n_blocks, n_blocks, width, scale_log2, + sm_scale, OFFSET=offset, QUANT=quantize, HAS_VALID=has_valid, HAS_QUAD=has_quad, + FLOOR=floor, FLOOR_VALUE=floor_value, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, + num_warps=_WARPS, num_stages=2) + return dq, dk, dv, None, None, None, None, None, None, None, None, None + + +def fp4_vsa_attn_qat(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + q2k_quad: torch.Tensor | None = None, + *, + quantize: bool = True, + two_level_p: bool = False, + smooth_k: bool = False, + first_block_max_floor: float | None = None, + sm_scale: float | None = None, + high_prec_o: bool = True) -> torch.Tensor: + """Block-sparse attention with the deployed FP4 numerics in the forward and STE gradients. + + ``q``/``k``/``v`` are ``[B, H, L, D]`` (L a multiple of 128, tile-ordered, + pad rows zero); the lists are ``vsa_tile_mask_to_fp4_blocks(mask, 128, vbs)`` + outputs (``q2k_quad`` is None for 128-token tiles, which map 1:1 onto + blocks, and carries the four 64x64 quadrant bits for 64-token tiles). + ``smooth_k`` subtracts each head's token mean from K first (exact for + softmax, differentiable through autograd). Returns ``[B, H, L, D]`` BF16. + """ + if smooth_k: + k = k - k.mean(dim=-2, keepdim=True) + scale = sm_scale if sm_scale is not None else 1.0 / math.sqrt(q.shape[-1]) + return _FP4VSAAttnQAT.apply(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize, two_level_p, + first_block_max_floor, scale, high_prec_o) + + +__all__ = ["fake_quant_qkv_triton", "fp4_vsa_attn_qat"] diff --git a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py new file mode 100644 index 000000000..4d13bfa91 --- /dev/null +++ b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse FP4 Attn-QAT (VSA tiles 64 and 128): the PyTorch specification and the Triton kernel. + +CPU tests pin the reference emulator (E2M1/E4M3 rounding, list builder, quadrant +masks, online softmax, dV weights). GPU tests (any CUDA GPU with Triton; run on +GB200) compare the Triton forward and STE backward with the reference. +""" +from __future__ import annotations + +import importlib +import importlib.util +import math +import sys +from pathlib import Path + +import pytest +import torch + +_DIR = Path(__file__).resolve().parents[1] / "python" / "fastvideo_kernel" / "triton_kernels" +_PKG = "_test_qat_vsa_kernels" + + +def _module(name: str): + if _PKG not in sys.modules: + spec = importlib.util.spec_from_loader(_PKG, loader=None, is_package=True) + spec.submodule_search_locations.append(str(_DIR)) + sys.modules[_PKG] = importlib.util.module_from_spec(spec) + return importlib.import_module(f"{_PKG}.{name}") + + +ref = _module("attn_qat_vsa_reference") +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs a CUDA GPU") + + +def _problem(n_tiles: int, heads: int = 2, keep: float = 0.4, seed: int = 0, device: str = "cpu", + partial_tail: bool = True): + """Random tile-ordered Q/K/V with a VSA-like mask: block 0 always kept, tail tile partially valid.""" + gen = torch.Generator().manual_seed(seed) + seq_len = n_tiles * ref.BLOCK + tile_valid = torch.full((n_tiles, ), ref.BLOCK, dtype=torch.int32) + if partial_tail: + tile_valid[-1] = 37 + tile_valid[n_tiles // 2] = 100 + token_valid = (torch.arange(ref.BLOCK)[None, :] < tile_valid[:, None]).reshape(-1) + q, k, v = (torch.randn(1, heads, seq_len, 128, generator=gen) * token_valid[None, None, :, None] + for _ in range(3)) + # RMS-normed Q/K have per-channel structure; add a channel bias so smoothing matters. + k = k + 0.5 * torch.randn(1, heads, 1, 128, generator=gen) * token_valid[None, None, :, None] + mask = torch.rand(1, heads, n_tiles, n_tiles, generator=gen) < keep + mask[..., 0] = True + q2k_idx, q2k_num, kv_valid = ref.tile128_mask_to_fp4_blocks(mask, tile_valid) + move = lambda t: t.to(device) # noqa: E731 + return (move(q.bfloat16()), move(k.bfloat16()), move(v.bfloat16()), move(q2k_idx), move(q2k_num), + move(kv_valid), move(mask), move(tile_valid)) + + +def test_e2m1_grid_and_ties_round_to_even(): + values = torch.tensor([0.0, 0.25, 0.26, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 5.1, 7.0, -0.75, -5.0]) + expected = torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0, 6.0, 6.0, -1.0, -4.0]) + assert torch.equal(ref.e2m1_rn(values), expected) + + +def test_e4m3_rounding_saturates_and_flushes(): + values = torch.tensor([448.0, 1000.0, 1.0625, 2.0**-10, 2.0**-9, 0.0]) + out = ref.e4m3_rn(values) + assert out[0] == 448 and out[1] == 448 + assert out[2] == 1.0 # tie between 1.0 and 1.125 rounds to even + assert out[3] == 0.0 and out[4] == 2.0**-9 and out[5] == 0.0 + + +def test_fake_quant_matches_nvfp4_definition(): + x = torch.tensor([[0.0] * 15 + [6.0], [1.0] * 16]).float() + out = ref.nvfp4_fake_quant_lastdim(x) + assert torch.equal(out[0], x[0]) # amax 6 -> scale exactly 1 + # amax 1 -> scale e4m3(1/6) = 0.171875, 1 / 0.171875 = 5.82 -> E2M1 6 -> 1.03125 + assert torch.equal(out[1], torch.full((16, ), 1.03125)) + assert ref.nvfp4_fake_quant_lastdim(torch.zeros(1, 16)).abs().sum() == 0 + + +def test_list_builder_is_descending_and_counts_valid_tiles(): + _, _, _, q2k_idx, q2k_num, kv_valid, mask, tile_valid = _problem(6) + for h in range(mask.shape[1]): + for m in range(6): + count = int(q2k_num[0, h, m]) + listed = q2k_idx[0, h, m, :count].tolist() + assert listed == sorted(listed, reverse=True) + assert listed == sorted(torch.nonzero(mask[0, h, m]).flatten().tolist(), reverse=True) + assert listed[-1] == 0 + assert kv_valid.tolist()[-2:] == [37, 0] + assert kv_valid.numel() == 2 * tile_valid.numel() + + +def test_invert_block_lists_round_trips(): + _, _, _, q2k_idx, q2k_num, _, mask, _ = _problem(5, heads=3, seed=3) + k2q_idx, k2q_slot, k2q_num = ref.invert_block_lists(q2k_idx, q2k_num, 5) + for h in range(3): + for n in range(5): + entries = [(int(k2q_idx[0, h, n, i]), int(k2q_slot[0, h, n, i])) for i in range(int(k2q_num[0, h, n]))] + assert [m for m, _ in entries] == torch.nonzero(mask[0, h, :, n]).flatten().tolist() + for m, slot in entries: + assert int(q2k_idx[0, h, m, slot]) == n + + +def test_reference_without_quantization_is_masked_softmax(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(4) + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + token_mask = ref.dense_block_mask(q2k_idx, q2k_num, q.shape[2], kv_valid) + scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(128) + expected = torch.softmax(scores.masked_fill(~token_mask, -math.inf), dim=-1) @ v.float() + torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("two_level_p", [False, True]) +def test_reference_dv_weights_reproduce_output(two_level_p): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(3, seed=1) + out, weights = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + return_dv_weights=True) + _, _, v_hat = ref.fake_quant_qkv(q, k, v) + torch.testing.assert_close(weights @ v_hat, out, rtol=1e-4, atol=1e-5) + + +def test_two_level_p_is_closer_to_exact_than_single_level(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(6, seed=2, keep=0.8) + exact = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + errors = {} + for two_level in (False, True): + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level) + errors[two_level] = float((out - exact).norm() / exact.norm()) + assert errors[True] < errors[False] + + +# ----------------------------------------------------------------------------- tile 64: quadrant masks +def _problem64(n_tiles: int, heads: int = 2, keep: float = 0.4, seed: int = 0, device: str = "cpu"): + """A 64-token VSA problem: two tiles share one 128-token kernel block.""" + gen = torch.Generator().manual_seed(seed) + tile_valid = torch.full((n_tiles, ), 64, dtype=torch.int32) + if n_tiles > 2: + tile_valid[-1] = 21 + tile_valid[n_tiles // 2] = 40 + tile_mask = torch.rand(1, heads, n_tiles, n_tiles, generator=gen) < keep + tile_mask[..., 0] = True + q2k_idx, q2k_num, kv_valid, q2k_quad = ref.tile_mask_to_fp4_blocks(tile_mask, 64, tile_valid) + n_blocks = q2k_idx.shape[-1] + seq_len = n_blocks * ref.BLOCK + token_valid = torch.zeros(seq_len, dtype=torch.bool) + for tile in range(n_tiles): + token_valid[tile * 64:(tile + 1) * 64] = torch.arange(64) < tile_valid[tile] + q, k, v = (torch.randn(1, heads, seq_len, 128, generator=gen) * token_valid[None, None, :, None] + for _ in range(3)) + k = k + 0.5 * torch.randn(1, heads, 1, 128, generator=gen) * token_valid[None, None, :, None] + move = lambda t: t.to(device) # noqa: E731 + return (move(q.bfloat16()), move(k.bfloat16()), move(v.bfloat16()), move(q2k_idx), move(q2k_num), + move(kv_valid), move(q2k_quad), move(tile_mask), move(tile_valid)) + + +def _dense_quad_mask(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, q2k_quad: torch.Tensor, seq_len: int) -> torch.Tensor: + """``[B, H, S, S]`` token mask: listed blocks intersected with their selected 64x64 quadrants.""" + batch, heads, n_blocks, _ = q2k_idx.shape + mask = torch.zeros(batch, heads, seq_len, seq_len, dtype=torch.bool, device=q2k_idx.device) + half = torch.arange(ref.BLOCK, device=q2k_idx.device) // ref.HALF + for b in range(batch): + for h in range(heads): + for m in range(n_blocks): + for slot in range(int(q2k_num[b, h, m])): + n = int(q2k_idx[b, h, m, slot]) + quad = int(q2k_quad[b, h, m, slot]) + qvalid = ((quad >> (half[:, None] * 2 + half[None, :])) & 1).bool() + mask[b, h, m * ref.BLOCK:(m + 1) * ref.BLOCK, n * ref.BLOCK:(n + 1) * ref.BLOCK] = qvalid + return mask + + +def test_tile64_lists_a_block_for_any_selected_quadrant(): + _, _, _, q2k_idx, q2k_num, kv_valid, q2k_quad, tile_mask, tile_valid = _problem64(6, heads=2, seed=2) + n_blocks = q2k_idx.shape[-1] + half_mask = tile_mask & (tile_valid > 0)[None, None, None, :] + for h in range(2): + for m in range(n_blocks): + listed = q2k_idx[0, h, m, :int(q2k_num[0, h, m])].tolist() + assert listed == sorted(listed, reverse=True) + expected = torch.nonzero(half_mask[0, h, 2 * m:2 * m + 2, :].any(dim=0).view(-1, 2).any(dim=1)).flatten() + assert listed == sorted(expected.tolist(), reverse=True) + assert kv_valid.numel() == 2 * n_blocks + + +def test_tile64_quadrant_bits_match_the_half_tile_mask(): + _, _, _, q2k_idx, q2k_num, _, q2k_quad, tile_mask, tile_valid = _problem64(6, heads=2, seed=3) + half_mask = tile_mask & (tile_valid > 0)[None, None, None, :] + weights = torch.tensor([[1, 2], [4, 8]]) + for h in range(2): + for m in range(q2k_idx.shape[2]): + for slot in range(int(q2k_num[0, h, m])): + n = int(q2k_idx[0, h, m, slot]) + quads = half_mask[0, h, 2 * m:2 * m + 2, 2 * n:2 * n + 2] + expected = int((quads.to(torch.int32) * weights).sum()) + assert int(q2k_quad[0, h, m, slot]) == expected + + +def test_tile128_route_is_the_general_route_without_quadrants(): + gen = torch.Generator().manual_seed(4) + tile_mask = torch.rand(1, 2, 5, 5, generator=gen) < 0.5 + tile_valid = torch.tensor([128, 128, 77, 0, 128], dtype=torch.int32) + q2k_idx, q2k_num, kv_valid, q2k_quad = ref.tile_mask_to_fp4_blocks(tile_mask, 128, tile_valid) + old_idx, old_num, old_valid = ref.tile128_mask_to_fp4_blocks(tile_mask, tile_valid) + assert q2k_quad is None + assert torch.equal(q2k_idx, old_idx) and torch.equal(q2k_num, old_num) and torch.equal(kv_valid, old_valid) + + +def test_tile64_quadrant_mask_is_exact_masked_softmax(): + q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, _, _ = _problem64(5, seed=5) + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize=False) + mask = _dense_quad_mask(q2k_idx, q2k_num, q2k_quad, q.shape[2]) + valid = ref.key_valid_mask(kv_valid, q2k_idx.shape[-1], q.device).reshape(-1) + mask = mask & valid[None, None, None, :] + scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(128) + expected = torch.softmax(scores.masked_fill(~mask, -math.inf), dim=-1) @ v.float() + torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-5) + + +# ----------------------------------------------------------------------------- Triton (GPU) +def _triton(): + return _module("attn_qat_vsa_train") + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)) + + +@needs_gpu +@pytest.mark.parametrize("two_level_p,smooth_k", [(False, False), (True, True)]) +def test_triton_tile64_forward_matches_reference(two_level_p, smooth_k): + q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, _, _ = _problem64(6, heads=2, seed=10, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, + two_level_p=two_level_p, smooth_k=smooth_k) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad=q2k_quad, + two_level_p=two_level_p, smooth_k=smooth_k) + exact = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize=False) + assert _rel(out, expected) < 2e-3 < 0.1 * _rel(expected, exact) + + +@needs_gpu +@pytest.mark.parametrize("two_level_p,smooth_k", [(False, False), (True, False), (False, True), (True, True)]) +def test_triton_forward_matches_reference(two_level_p, smooth_k): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(6, heads=2, seed=4, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + smooth_k=smooth_k) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + smooth_k=smooth_k) + exact = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + # The emulator must sit far closer to the reference than FP4 sits to exact attention. + assert _rel(out, expected) < 2e-3 < 0.1 * _rel(expected, exact) + + +@needs_gpu +def test_triton_forward_unquantized_matches_softmax(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, seed=5, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + assert _rel(out, expected) < 5e-3 + + +@needs_gpu +def test_triton_first_block_floor_matches_reference(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(4, seed=6, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, first_block_max_floor=0.0) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, first_block_max_floor=0.0) + assert _rel(out, expected) < 2e-3 + + +@needs_gpu +@pytest.mark.parametrize("two_level_p", [False, True]) +def test_triton_backward_matches_ste_reference(two_level_p): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, heads=2, seed=7, device="cuda") + q, k, v = (t.clone().requires_grad_(True) for t in (q, k, v)) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p) + grad_out = torch.randn_like(out) + dq, dk, dv = torch.autograd.grad(out, (q, k, v), grad_out) + _, weights = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + return_dv_weights=True) + q_hat, k_hat, v_hat = ref.fake_quant_qkv(q.detach(), k.detach(), v.detach()) + token_mask = ref.dense_block_mask(q2k_idx, q2k_num, q.shape[2], kv_valid) + expected = ref.ste_gradients_reference(q_hat, k_hat, v_hat, token_mask, weights, grad_out, 1 / math.sqrt(128)) + for name, got, want in zip("qkv", (dq, dk, dv), expected, strict=True): + assert _rel(got, want) < 2e-2, name + + +@needs_gpu +def test_triton_dv_reaches_every_head_on_long_sequences(): + """Regression for the GB10 Attn-QAT bug: re-quantized normalized P underflowed and zeroed dV per head.""" + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(64, heads=4, keep=0.9, seed=8, device="cuda") + v = v.clone().requires_grad_(True) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid) + (dv, ) = torch.autograd.grad(out, (v, ), torch.ones_like(out)) + per_head = dv.float().norm(dim=(-1, -2))[0] + assert bool((per_head > 0.1 * per_head.mean()).all()), per_head + + +@needs_gpu +def test_triton_fake_quant_is_bit_identical_to_reference(): + q, k, v, *_ = _problem(4, heads=3, seed=9, device="cuda") + q = q * 37.0 # exercise large and subnormal E4M3 scales + got = _triton().fake_quant_qkv_triton(q, k, v) + want = ref.fake_quant_qkv(q, k, v) + for name, a, b in zip("qkv", got, want, strict=True): + assert torch.equal(a.float(), b.to(torch.bfloat16).float()), name diff --git a/fastvideo/attention/backends/fp4_vsa_qat.py b/fastvideo/attention/backends/fp4_vsa_qat.py new file mode 100644 index 000000000..123a1e523 --- /dev/null +++ b/fastvideo/attention/backends/fp4_vsa_qat.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VSA-H3 tile-128 fine attention with the deployed sparse-FP4 numerics, for quantization-aware training. + +``FP4VSAFineAttention`` plugs into ``MiniMaxH3VSAImpl.fine_attention_override``: +VSA-H3's own selection (tile pooling, top-k mask, gated compression branch) +runs unchanged and in BF16, and only the block-sparse attention is replaced by +``fastvideo_kernel.triton_kernels.attn_qat_vsa_train.fp4_vsa_attn_qat``, which +emulates ``attn_qat_infer``'s sparse FP4 kernel on any CUDA GPU (forward) and +provides straight-through gradients (backward). + +Training images may carry a prebuilt ``fastvideo_kernel`` that predates these +modules; ``load_qat_vsa_modules`` then loads them from this checkout's +``fastvideo-kernel/python`` tree under a private package name. +""" +from __future__ import annotations + +import importlib +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Any + +import torch + +_PRIVATE_PACKAGE = "_fastvideo_qat_vsa_kernels" +_KERNEL_DIR = Path( + __file__).resolve().parents[3] / "fastvideo-kernel" / "python" / "fastvideo_kernel" / "triton_kernels" + + +def _load(name: str) -> ModuleType: + try: + return importlib.import_module(f"fastvideo_kernel.triton_kernels.{name}") + except ImportError: + pass + if _PRIVATE_PACKAGE not in sys.modules: + if not (_KERNEL_DIR / f"{name}.py").is_file(): + raise ImportError(f"{name}.py is neither in the installed fastvideo_kernel nor in {_KERNEL_DIR}") + spec = importlib.util.spec_from_loader(_PRIVATE_PACKAGE, loader=None, is_package=True) + assert spec is not None and spec.submodule_search_locations is not None + spec.submodule_search_locations.append(str(_KERNEL_DIR)) + sys.modules[_PRIVATE_PACKAGE] = importlib.util.module_from_spec(spec) + return importlib.import_module(f"{_PRIVATE_PACKAGE}.{name}") + + +def load_qat_vsa_modules() -> tuple[ModuleType, ModuleType]: + """``(attn_qat_vsa_reference, attn_qat_vsa_train)``; the second needs Triton.""" + return _load("attn_qat_vsa_reference"), _load("attn_qat_vsa_train") + + +@dataclass(frozen=True) +class FP4AttentionNumerics: + """Numeric mode of the deployed kernel; mirror whatever inference ships. + + ``quantize=False`` keeps the emulator's sparse structure with exact + softmax (a BF16, grad-capable tile-128 route). + """ + quantize: bool = True + two_level_p: bool = False + smooth_k: bool = False + first_block_max_floor: float | None = None + + +class FP4VSAFineAttention: + """``fine_attention_override`` callable: ``(q, k, v, mask, metadata) -> [B, S, H, D]``.""" + + def __init__(self, numerics: FP4AttentionNumerics) -> None: + self.numerics = numerics + self._reference, self._train = load_qat_vsa_modules() + + def __call__(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor, + metadata: Any) -> torch.Tensor: + tile_elems = int(metadata.tile_elems) + if tile_elems not in (64, 128): + raise ValueError(f"the FP4 VSA emulator models 64- or 128-token tiles, got {tile_elems}") + q2k_idx, q2k_num, kv_valid, q2k_quad = self._reference.tile_mask_to_fp4_blocks( + mask, tile_elems, metadata.variable_block_sizes) + q, k, v = (t.transpose(1, 2) for t in (query, key, value)) + seq_len = q.shape[2] + padded = (seq_len + 127) // 128 * 128 + if padded != seq_len: + pad = padded - seq_len + q, k, v = (torch.nn.functional.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) + out = self._train.fp4_vsa_attn_qat(q, + k, + v, + q2k_idx, + q2k_num, + kv_valid, + q2k_quad=q2k_quad, + quantize=self.numerics.quantize, + two_level_p=self.numerics.two_level_p, + smooth_k=self.numerics.smooth_k, + first_block_max_floor=self.numerics.first_block_max_floor) + if padded != seq_len: + out = out[:, :, :seq_len] + return out.transpose(1, 2) + + +def vsa_impls(transformer: torch.nn.Module) -> list[Any]: + """Every VSA-H3 attention implementation in a MiniMax-H3 transformer.""" + impls = [] + for module in transformer.modules(): + impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None) + if impl is not None and hasattr(impl, "fine_attention_override"): + impls.append(impl) + return impls + + +def install_fp4_vsa_attention(transformer: torch.nn.Module, numerics: FP4AttentionNumerics | None) -> int: + """Route every VSA-H3 layer's fine attention through the emulator (``None`` restores the kernels).""" + override = FP4VSAFineAttention(numerics) if numerics is not None else None + impls = vsa_impls(transformer) + for impl in impls: + impl.fine_attention_override = override + return len(impls) + + +__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "install_fp4_vsa_attention", "load_qat_vsa_modules"] diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index e369cf44f..929ff1995 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -589,6 +589,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): # 128-token blocks, an even tile count with the partner tile, integer # tile sizes), which every call meets, so it is evaluated once. self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {} + # Optional replacement of the fine block-sparse kernel only (selection, + # mask and the gated compression branch are unchanged). Called as + # ``fn(query, key, value, mask, metadata) -> [B, S, H, D]`` on the + # logical tile-ordered rows. Quantization-aware training installs its + # FP4 emulator here; it is also the only grad-capable tile-128 route. + self.fine_attention_override: Any = None def prepare_for_compile(self, device: torch.device) -> None: """Tensorize per-layer state shared by every torch.compile route.""" @@ -758,7 +764,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): raise RuntimeError("VSA-H3 regional fullgraph compile requires 64-token tiles; disable " "inference_torch_compile for tile-128/256 runs.") if tile_elems == 64: - if block_sparse_attn_64_bhsd is None: + if block_sparse_attn_64_bhsd is None and self.fine_attention_override is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn is not installed") elif tile_elems == 256 and block_sparse_attn_256_bshd is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed") @@ -838,7 +844,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): # preserves the eager dense-layer contract without a Python branch. mask = mask | force_dense - if tile_elems in _PTX_BLACKWELL_TILE_ELEMS: + if self.fine_attention_override is not None and not compiling: + out = self.fine_attention_override(logical_query, logical_key, logical_value, mask, attn_metadata) + elif tile_elems in _PTX_BLACKWELL_TILE_ELEMS: # Native 64/128-token path: the block map is already at the # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose diff --git a/fastvideo/layers/quantization/nvfp4_qad.py b/fastvideo/layers/quantization/nvfp4_qad.py new file mode 100644 index 000000000..01f13ad26 --- /dev/null +++ b/fastvideo/layers/quantization/nvfp4_qad.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 quantization-aware distillation (QAD) for the MiniMax-H3 DiT block linears. + +``NVFP4QADLinearMethod`` replaces a ``ReplicatedLinear``'s ``quant_method`` +in place, so parameter names, FSDP sharding and checkpoints are unchanged. +Its forward *is* the deployed NVFP4 GEMM, computed from the layer's current +(BF16-cast) master weight exactly as the packed export and +``NVFP4QuantizeMethod.apply`` compute it: + +* weight: RTN NVFP4 with global scale ``448 * 6 / amax(W)`` (the converter's + ``quantize_dense_linear``; ModelOpt max calibration gives the same codes); +* activation: per-16 E4M3 block scales under a global scale that is + ``static`` (``448 * 6 / calibrated amax``, the FFN linears) or ``unit`` + (1.0, attention projections and VSA gates); +* GEMM: FlashInfer ``mm_fp4`` with ``alpha = (1 / weight_sf) / input_sf``. + +With gradients enabled the forward runs inside a straight-through estimator +whose backward is full precision. ``enabled = +False`` delegates to the original method, so a disabled student is the +teacher bit for bit. +""" +from __future__ import annotations + +import contextlib +import json +import re +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +FP4_RANGE = 448.0 * 6.0 +H3_BLOCK_LINEAR = re.compile(r"transformer_blocks\.(\d+)\.(attn\.to_q|attn\.to_k|attn\.to_v|attn\.to_out" + r"|attn\.to_gate_compress|ff\.fc_in|ff\.fc_out)$") +STATIC_SUFFIXES = ("ff.fc_in", "ff.fc_out") + + +def weight_global_sf(weight: torch.Tensor) -> torch.Tensor: + """``448 * 6 / amax(|W|)`` in fp32, as the packed-export converter computes it. + + The max of a BF16 tensor is exact in BF16, so one BF16 reduction gives the converter's value + (``weight.float().abs().nan_to_num().max()``) for finite weights without fp32 temporaries. + """ + return FP4_RANGE / torch.nan_to_num(weight.abs().amax().float()).clamp(min=1e-12) + + +def nvfp4_gemm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, + input_global_sf: torch.Tensor) -> torch.Tensor: + """The deployed NVFP4 linear on a BF16 weight: quantize both operands, FP4 GEMM.""" + from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend, _nvfp4_quantize + from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer + + sf_layout, _, _ = _require_flashinfer() + weight = weight.to(torch.bfloat16).contiguous() + w_sf = weight_global_sf(weight) + w_fp4, w_scale = _nvfp4_quantize(weight, w_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + alpha = (1.0 / w_sf) / input_global_sf + shape = x.shape + x2d = x.to(torch.bfloat16).reshape(-1, shape[-1]).contiguous() + x_fp4, x_scale = _nvfp4_quantize(x2d, input_global_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + out = _mm_fp4(x_fp4, w_fp4.T, x_scale, w_scale.T, alpha, torch.bfloat16, None, backend=_mm_fp4_backend()) + if bias is not None: + out = out + bias.to(out.dtype) + return out.view(*shape[:-1], weight.shape[0]) + + +class _NVFP4STE(torch.autograd.Function): + """Forward = ``nvfp4_gemm``; backward = full-precision linear gradients (straight-through).""" + + @staticmethod + def forward(ctx, x, weight, bias, input_global_sf): # type: ignore[override] + ctx.save_for_backward(x, weight) + ctx.has_bias = bias is not None + return nvfp4_gemm(x, weight, bias, input_global_sf) + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + x, weight = ctx.saved_tensors + grad_2d = grad_out.reshape(-1, grad_out.shape[-1]) + grad_x = (grad_2d @ weight.to(grad_2d.dtype)).reshape(x.shape).to(x.dtype) + grad_w = (grad_2d.t() @ x.reshape(-1, x.shape[-1]).to(grad_2d.dtype)).to(weight.dtype) + grad_b = grad_2d.sum(dim=0) if ctx.has_bias else None + return grad_x, grad_w, grad_b, None + + +class NVFP4QADLinearMethod: + """Drop-in ``quant_method`` for one H3 block linear (see module docstring).""" + + def __init__(self, + original: Any, + prefix: str, + act_scale: str, + input_amax: float | None = None, + allow_uncalibrated: bool = False) -> None: + if act_scale not in ("static", "unit"): + raise ValueError(f"act_scale must be 'static' or 'unit', got {act_scale!r}") + if act_scale == "static" and not allow_uncalibrated and (input_amax is None or not input_amax > 0): + raise ValueError(f"{prefix}: a static activation scale needs a positive calibrated amax") + self.original = original + self.prefix = prefix + self.act_scale = act_scale + self.input_amax = float(input_amax) if input_amax is not None else None + self.enabled = True + # Calibration: run the original (dense) path and fold each input's abs-max into ``observed_amax``. + self.calibrating = False + self.observed_amax: torch.Tensor | None = None + self._sf: dict[torch.device, torch.Tensor] = {} + + def set_input_amax(self, amax: float) -> None: + if not amax > 0: + raise ValueError(f"{self.prefix}: calibrated amax must be positive, got {amax}") + self.input_amax = float(amax) + self._sf.clear() + + def input_global_sf(self, device: torch.device) -> torch.Tensor: + if self.act_scale == "static" and self.input_amax is None: + raise RuntimeError(f"{self.prefix}: static activation scale used before calibration") + if device not in self._sf: + value = 1.0 if self.act_scale == "unit" else FP4_RANGE / max(self.input_amax or 0.0, 1e-12) + self._sf[device] = torch.tensor(value, dtype=torch.float32, device=device) + return self._sf[device] + + def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + if self.calibrating: + amax = x.detach().abs().amax().float() + self.observed_amax = amax if self.observed_amax is None else torch.maximum(self.observed_amax, amax) + return self.original.apply(layer, x, bias) + if not self.enabled: + return self.original.apply(layer, x, bias) + input_sf = self.input_global_sf(x.device) + if torch.is_grad_enabled() and (layer.weight.requires_grad or x.requires_grad): + return _NVFP4STE.apply(x, layer.weight, bias, input_sf) + return nvfp4_gemm(x, layer.weight, bias, input_sf) + + def __getattr__(self, name: str) -> Any: + # Callers probe optional quant-method hooks (e.g. wants_prequantized_input); fall back to the original. + # Before __init__ (copy.deepcopy, unpickling) there is no original: AttributeError keeps hasattr working. + try: + original = self.__dict__["original"] + except KeyError: + raise AttributeError(name) from None + return getattr(original, name) + + +@dataclass(frozen=True) +class NVFP4QADPlan: + """Which block linears are quantized and how (``h3_dit_vsa``: attention, gates and FFN).""" + quantize_attention: bool = True + quantize_gate: bool = True + quantize_ffn: bool = True + skip_blocks: tuple[int, ...] = () + + +def load_amax_table(path: str | Path) -> dict[str, float]: + """``FASTVIDEO_NVFP4_ACT_AMAX``-format JSON (``b.`` or full prefix -> amax or {"all": amax}).""" + raw = json.loads(Path(path).read_text()) + return {key: float(value["all"] if isinstance(value, dict) else value) for key, value in raw.items()} + + +def _amax_for(table: dict[str, float], name: str, block: str, sub: str) -> float | None: + for key in (name, f"b{block}.{sub}", f"transformer_blocks.{block}.{sub}"): + if key in table: + return table[key] + return None + + +def install_nvfp4_qad(transformer: torch.nn.Module, + amax_table: dict[str, float] | None, + plan: NVFP4QADPlan, + allow_uncalibrated: bool = False) -> dict[str, NVFP4QADLinearMethod]: + """Install QAD methods on the selected block linears; returns them by module name. + + FFN linears take ``static`` scales from ``amax_table`` (required unless + ``allow_uncalibrated``, which leaves them for ``calibrate``); attention + projections and gates use the unit scale, as the deployed export does. + """ + installed: dict[str, NVFP4QADLinearMethod] = {} + for raw_name, module in transformer.named_modules(): + # Activation checkpointing wraps blocks; key everything by the checkpoint (unwrapped) name. + name = raw_name.replace("_checkpoint_wrapped_module.", "") + match = H3_BLOCK_LINEAR.search(name) + if match is None or not hasattr(module, "quant_method") or int(match.group(1)) in plan.skip_blocks: + continue + block, sub = match.group(1), match.group(2) + if sub.startswith("ff.") and not plan.quantize_ffn: + continue + if sub == "attn.to_gate_compress" and not plan.quantize_gate: + continue + if sub.startswith("attn.to_") and sub != "attn.to_gate_compress" and not plan.quantize_attention: + continue + original = module.quant_method + if isinstance(original, NVFP4QADLinearMethod): + original = original.original + if sub in STATIC_SUFFIXES: + amax = _amax_for(amax_table or {}, name, block, sub) + if amax is None and not allow_uncalibrated: + raise KeyError(f"no calibrated input amax for {name} (key b{block}.{sub}); run Stage A calibration") + method = NVFP4QADLinearMethod(original, name, "static", amax, allow_uncalibrated=allow_uncalibrated) + else: + method = NVFP4QADLinearMethod(original, name, "unit") + module.quant_method = method + installed[name] = method + if not installed: + raise ValueError("no MiniMax-H3 block linears matched; is this a MiniMax-H3 transformer?") + return installed + + +@contextlib.contextmanager +def calibrating(methods: dict[str, NVFP4QADLinearMethod]) -> Iterator[None]: + """Max-calibrate the static linears that have no amax yet (dense forwards while active). + + On exit the observed maxima are all-reduced (MAX) over the default process + group, so every rank, including each sequence-parallel shard, agrees. + """ + pending = {name: m for name, m in methods.items() if m.act_scale == "static" and m.input_amax is None} + for method in pending.values(): + method.calibrating, method.observed_amax = True, None + try: + yield + finally: + for method in pending.values(): + method.calibrating = False + if not pending: + return + names = sorted(pending) + device = next((m.observed_amax.device for m in pending.values() if m.observed_amax is not None), + torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else None) + values = torch.stack([(pending[n].observed_amax if pending[n].observed_amax is not None else torch.zeros( + (), device=device)).float() for n in names]) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.MAX) + for name, value in zip(names, values.tolist(), strict=True): + if not value > 0: + raise RuntimeError(f"{name} saw no activations during calibration") + pending[name].set_input_amax(value) + + +def set_nvfp4_qad_enabled(methods: dict[str, NVFP4QADLinearMethod], enabled: bool) -> None: + for method in methods.values(): + method.enabled = enabled + + +def nvfp4_qad_export_scales(methods: dict[str, NVFP4QADLinearMethod]) -> dict[str, float]: + """The static FFN amax values the export must carry (``--act-amax`` input of the packed converter).""" + table = {} + for name, method in methods.items(): + match = H3_BLOCK_LINEAR.search(name) + if method.act_scale == "static" and match is not None: + table[f"b{match.group(1)}.{match.group(2)}"] = float(method.input_amax or 0.0) + return table + + +__all__ = [ + "NVFP4QADLinearMethod", "NVFP4QADPlan", "calibrating", "install_nvfp4_qad", "load_amax_table", "nvfp4_gemm", + "nvfp4_qad_export_scales", "set_nvfp4_qad_enabled", "weight_global_sf" +] diff --git a/fastvideo/models/dits/minimax_h3_token_compression.py b/fastvideo/models/dits/minimax_h3_token_compression.py new file mode 100644 index 000000000..fe74fd0d8 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_token_compression.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Middle-block spatial token compression for FastH3 (QAD training and inference; hook based). + +Blocks ``[start, end)`` run on fewer video tokens. Text and audio rows are untouched; a kept video token keeps its +original RoPE coordinates and AdaLN row; after the last compressed block every video token receives the update of +its nearest kept token: ``H_out = H + expand(F(P H) - P H)`` where ``P`` is the compression operator. With +compression off nothing changes (no hook acts), so the student equals the teacher bit for bit. + +Unlike ``scripts/benchmarks/minimax_h3_realtime/h3_levers.py`` (which replaces the block list), this installs +forward hooks, so module names, FSDP sharding, activation checkpointing and QAD linears are untouched: + +* block ``start`` pre-hook compresses the hidden states and swaps AdaLN rows / RoPE / sequence length; +* every block in the range gets the compressed AdaLN rows / RoPE; +* the inner attention module of each compressed block swaps the VSA metadata for the compressed video grid (a hook + on the *inner* module, so activation-checkpoint recompute during backward swaps it too); +* block ``end - 1`` post-hook expands and restores. + +Single sequence-parallel rank only (the compression indexes the whole packed sequence). +""" +from __future__ import annotations + +import contextlib +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn + +from fastvideo.forward_context import get_forward_context + +MODES = ("subsample", "pair_x", "pool2x2") + + +@dataclass(frozen=True) +class CompressionSpec: + start: int + end: int + mode: str = "subsample" + # Compress only the first ``first_steps`` of the sampling forwards (None = all); the low-noise steps that set + # fine detail then run on the full token grid. + first_steps: int | None = None + + def __post_init__(self) -> None: + if self.mode not in MODES: + raise ValueError(f"compression mode must be one of {MODES}, got {self.mode!r}") + if not 0 <= self.start < self.end: + raise ValueError(f"need 0 <= start < end, got [{self.start}, {self.end})") + + @classmethod + def parse(cls, text: str) -> CompressionSpec | None: + """``"off"``/empty -> None; ``"mode:start:end"`` (e.g. ``subsample:3:41``).""" + text = text.strip().lower() + if text in ("", "0", "off", "none"): + return None + parts = text.split(":") + if len(parts) not in (3, 4): + raise ValueError(f"token compression must be 'off' or 'mode:start:end[:first_steps]', got {text!r}") + return cls(int(parts[1]), int(parts[2]), parts[0], int(parts[3]) if len(parts) == 4 else None) + + @classmethod + def from_config(cls, cfg: dict[str, Any] | None) -> CompressionSpec | None: + if not cfg or not cfg.get("enabled", True): + return None + first = cfg.get("first_steps") + return cls(int(cfg["start"]), int(cfg["end"]), str(cfg.get("mode", "subsample")), + None if first is None else int(first)) + + +def compressed_grid(mode: str, grid: tuple[int, int, int]) -> tuple[int, int, int]: + frames, height, width = grid + if mode == "pair_x": + return frames, height, (width + 1) // 2 + return frames, (height + 1) // 2, (width + 1) // 2 + + +def token_reduction(mode: str, grid: tuple[int, int, int]) -> float: + sub = compressed_grid(mode, grid) + return (grid[0] * grid[1] * grid[2]) / (sub[0] * sub[1] * sub[2]) + + +def _pad_to_even(x: torch.Tensor, dims: tuple[int, ...]) -> torch.Tensor: + for dim in dims: + if x.shape[dim] % 2: + x = torch.cat((x, x.narrow(dim, x.shape[dim] - 1, 1)), dim=dim) + return x + + +def grid_ops(mode: str, grid: tuple[int, int, int]) -> tuple[Callable, Callable, torch.Tensor]: + """``(compress, expand, keep_index)`` for a video token grid ``(T, H, W)``. + + ``compress``: rows ``[B, T*H*W, C]`` -> ``[B, T*h*w, C]``; ``expand``: delta ``[B, T*h*w, C]`` -> ``[B, T*H*W, C]``; + ``keep_index``: for every compressed token the video row whose RoPE / AdaLN row it keeps (CPU long). + """ + frames, height, width = grid + sub = compressed_grid(mode, grid) + index = torch.arange(frames * height * width).reshape(frames, height, width) + + def as_grid(rows: torch.Tensor) -> torch.Tensor: + return rows.reshape(rows.shape[0], frames, height, width, rows.shape[-1]) + + if mode == "subsample": + keep = index[:, ::2, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + return as_grid(rows)[:, :, ::2, ::2].reshape(rows.shape[0], -1, rows.shape[-1]) + + elif mode == "pair_x": + keep = index[:, :, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + g = _pad_to_even(as_grid(rows), (3, )) + pair = g.float().unflatten(3, (-1, 2)).mean(dim=4) + return pair.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1]) + + else: # pool2x2 + keep = index[:, ::2, ::2].reshape(-1) + + def compress(rows: torch.Tensor) -> torch.Tensor: + g = _pad_to_even(as_grid(rows), (2, 3)) + pooled = g.float().unflatten(3, (-1, 2)).unflatten(2, (-1, 2)).mean(dim=(3, 5)) + return pooled.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1]) + + repeat_h = 1 if mode == "pair_x" else 2 + + def expand(delta: torch.Tensor) -> torch.Tensor: + d = delta.reshape(delta.shape[0], *sub, delta.shape[-1]) + if repeat_h == 2: + d = d.repeat_interleave(2, dim=2)[:, :, :height] + d = d.repeat_interleave(2, dim=3)[:, :, :, :width] + return d.reshape(delta.shape[0], -1, delta.shape[-1]) + + return compress, expand, keep + + +def _unwrap(module: nn.Module) -> nn.Module: + return getattr(module, "_checkpoint_wrapped_module", module) + + +@dataclass +class _Forward: + """Per-forward state (one transformer call).""" + full: torch.Tensor | None = None + compressed_in: torch.Tensor | None = None + adaln: torch.Tensor | None = None + rotary: tuple[torch.Tensor, torch.Tensor] | None = None + seq_len: int = 0 + num_prefix: int = 0 + + +class TokenCompression: + """Installs the hooks; ``scope(grid, compressed_metadata)`` arms them for one forward (and its backward).""" + + def __init__(self, transformer: nn.Module, spec: CompressionSpec) -> None: + self.spec = spec + self.blocks = list(transformer.transformer_blocks) + if spec.end > len(self.blocks): + raise ValueError(f"compression range [{spec.start}, {spec.end}) exceeds {len(self.blocks)} blocks") + self._handles: list[Any] = [] + self._grid: tuple[int, int, int] | None = None + self._ops: tuple[Callable, Callable, torch.Tensor] | None = None + self._metadata: Any = None + self._saved_metadata: Any = None + self._state = _Forward() + for index in range(spec.start, spec.end): + block = self.blocks[index] + self._handles.append(block.register_forward_pre_hook(self._pre(index))) + if index == spec.end - 1: + self._handles.append(block.register_forward_hook(self._post)) + attn = _unwrap(block).attn + self._handles.append(attn.register_forward_pre_hook(self._attn_pre)) + self._handles.append(attn.register_forward_hook(self._attn_post)) + + @property + def active(self) -> bool: + return self._grid is not None + + @contextlib.contextmanager + def scope(self, grid: tuple[int, int, int], compressed_metadata: Any) -> Iterator[None]: + """Arm the hooks for forwards (and their backward recompute) inside the block.""" + self._grid, self._metadata = grid, compressed_metadata + self._ops = grid_ops(self.spec.mode, grid) + try: + yield + finally: + self._grid = self._ops = self._metadata = None + self._state = _Forward() + + def remove(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + # -- hooks + def _pre(self, index: int): + + def hook(module: nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...] | None: + if not self.active: + return None + hidden, temb, adaln, rotary, _ = args + if index == self.spec.start: + self._enter(hidden, adaln, rotary) + hidden = self._state.compressed_in + state = self._state + return hidden, temb, state.adaln, state.rotary, state.seq_len + + return hook + + def _enter(self, hidden: torch.Tensor, adaln: torch.Tensor, rotary: tuple[torch.Tensor, torch.Tensor]) -> None: + assert self._grid is not None and self._ops is not None + compress, _, keep = self._ops + num_video = self._grid[0] * self._grid[1] * self._grid[2] + num_prefix = hidden.shape[1] - num_video + if num_prefix < 0 or adaln.shape[0] != hidden.shape[1]: + raise ValueError("token compression needs the whole packed sequence on one rank " + f"(hidden {tuple(hidden.shape)}, adaln {tuple(adaln.shape)}, video {num_video})") + rows = torch.cat((torch.arange(num_prefix), keep + num_prefix)).to(hidden.device) + self._state = _Forward(full=hidden, + compressed_in=torch.cat((hidden[:, :num_prefix], compress(hidden[:, num_prefix:])), dim=1), + adaln=adaln.index_select(0, rows), + rotary=(rotary[0].index_select(0, rows), rotary[1].index_select(0, rows)), + seq_len=int(rows.numel()), + num_prefix=num_prefix) + + def _post(self, module: nn.Module, args: tuple[Any, ...], output: torch.Tensor) -> torch.Tensor | None: + if not self.active: + return None + assert self._ops is not None + state = self._state + n = state.num_prefix + delta = self._ops[1](output[:, n:] - state.compressed_in[:, n:]) + restored = torch.cat((output[:, :n], state.full[:, n:] + delta.to(state.full.dtype)), dim=1) + self._state = _Forward() + return restored + + def _attn_pre(self, module: nn.Module, args: tuple[Any, ...], kwargs: Any = None) -> None: + if not self.active: + return + context = get_forward_context() + self._saved_metadata = context.attn_metadata + context.attn_metadata = self._metadata + + def _attn_post(self, module: nn.Module, args: tuple[Any, ...], output: Any) -> None: + if not self.active: + return + get_forward_context().attn_metadata = self._saved_metadata + + +def block_flops_share(spec: CompressionSpec | None, grid: tuple[int, int, int], num_blocks: int, num_other_tokens: int, + attn_keep: float = 0.2) -> dict[str, float]: + """Relative DiT cost of the compressed model (1.0 = uncompressed), by FLOPs. + + Per block, per token: linear layers ~ ``12 d^2`` MACs (qkv+out+gate ~ 5 d^2, FFN SwiGLU 3 * d * 2.67 d), attention + ~ ``2 * attended_keys * d``; VSA keeps ``attn_keep`` of the video keys. ``d`` cancels in the ratio when expressed in + ``d`` units, so only the token counts matter; ``d = 5376`` is used for the attention/linear balance. + """ + d = 5376.0 + video = grid[0] * grid[1] * grid[2] + + def cost(tokens_video: float) -> float: + n = tokens_video + num_other_tokens + linear = 13.4 * d * d * n + attn = 2.0 * 2.0 * d * n * (attn_keep * tokens_video + num_other_tokens) + return linear + attn + + full = cost(video) + if spec is None: + return {"relative_flops": 1.0, "speedup_flops": 1.0} + reduced = cost(video / token_reduction(spec.mode, grid)) + inside = spec.end - spec.start + total = (num_blocks - inside) * full + inside * reduced + return {"relative_flops": total / (num_blocks * full), "speedup_flops": num_blocks * full / total} + + +__all__ = ["CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops", "token_reduction"] diff --git a/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py new file mode 100644 index 000000000..aca155a81 --- /dev/null +++ b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py @@ -0,0 +1,629 @@ +# 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): + plan.append(entries[index]) + 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_audio_stats_and_seed_floor_gate(): + from fastvideo.train.methods.knowledge_distillation.pdd_qad_audio_stats import (audio_stats, log_mel_distance, + seed_floor_gate) + rate = 32000 + t = torch.arange(rate) / rate + low, high = torch.sin(2 * math.pi * 1000 * t), torch.sin(2 * math.pi * 10000 * t) + stats_low, stats_high = audio_stats(low, rate), audio_stats(high[:, None].repeat(1, 2), rate) + assert stats_low["spectral_centroid_hz"] == pytest.approx(1000, rel=0.05) + assert stats_low["share_above_8khz"] < 0.01 and stats_high["share_above_8khz"] > 0.99 + assert log_mel_distance(low, low, rate) == 0.0 and log_mel_distance(low, high, rate) > 1.0 + teacher = [{"spectral_centroid_hz": 950.0, "share_above_8khz": 0.1, "dynamic_range_db": 30.0, + "loudness_iqr_db": 6.0}] * 3 + other = [{**row, "spectral_centroid_hz": 900.0, "dynamic_range_db": 28.0, "loudness_iqr_db": 5.0, + "share_above_8khz": 0.09} for row in teacher] + good = [{**row, "spectral_centroid_hz": 930.0, "share_above_8khz": 0.095} for row in teacher] + dull = [{**row, "spectral_centroid_hz": 530.0} for row in teacher] + assert seed_floor_gate(good, teacher, other)["pass"] + result = seed_floor_gate(dull, teacher, other) + assert not result["pass"] and not result["spectral_centroid_hz/pass"] + + +def test_fp32_boundary_modules_match_inference_and_stay_frozen(cpu_loader, single_process_group): # noqa: F811 + """The QAD loader keeps H3's FP32 boundary modules FP32 (as inference) and refuses to train them.""" + from types import SimpleNamespace + from fastvideo.train.models.minimax_h3.omniref_pdd import MiniMaxH3OmniRefPDDModel + model = _tiny_model().to(torch.float32) + assert not model.config.uniform_parameter_dtype + fp32 = {n for n, _ in model.named_parameters() if model._get_parameter_dtype(n, torch.bfloat16) == torch.float32} + assert fp32 and {n.split(".", 1)[0] for n in fp32} <= set(model._keep_in_fp32_modules) + model.requires_grad_(False) + fake = SimpleNamespace(transformer=model) + MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake) + model.proj_in.weight.requires_grad_(True) + with pytest.raises(ValueError, match="must stay frozen"): + MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake) + + +def test_block_first_call_keeps_the_fsdp_root(single_process_group): # noqa: F811 + """Stage 2 from a checkpoint calibrates reconstruction (student blocks called directly) before any full student + forward; without pinning the root first, FSDP2 makes block 0 the root and the next full forward fails.""" + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import ensure_fsdp_root_initialized + device = "cuda" if torch.cuda.is_available() else "cpu" + + class Toy(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = torch.nn.ModuleList(torch.nn.Linear(4, 4) for _ in range(2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for block in self.transformer_blocks: + x = block(x) + return x + + def sharded() -> Toy: + mesh = init_device_mesh(device, (1, )) + model = Toy().to(device) + for block in model.transformer_blocks: + fully_shard(block, mesh=mesh) + return fully_shard(model, mesh=mesh) + + x = torch.randn(2, 4, device=device) + broken = sharded() + broken.transformer_blocks[0](x) + with pytest.raises(RuntimeError, match="already been lazily initialized"): + broken(x) + fixed = sharded() + ensure_fsdp_root_initialized(fixed) + ensure_fsdp_root_initialized(fixed) # idempotent + fixed.transformer_blocks[0](x) + assert fixed(x).shape == (2, 4) + ensure_fsdp_root_initialized(torch.nn.Linear(2, 2)) # no FSDP: no-op + + +def test_step0_cache_round_trip_and_fingerprint(tmp_path): + from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0, + step0_fingerprint_inputs) + method = {"init_student_dcp": "/runs/e3/best", "noise_floor_eval": True, "eligibility_keys": ["endpoint/audio"], + "recon_group_weights": {"video": 1.0, "audio": 1.0}, "eval_every": 20} + data = {"eval_manifest": "/m.json", "eval_per_group": {"480p": 4, "768p": 2}, "cases": ["first_frame"]} + qad = {"linears": {"skip_blocks": [42, 43]}, "fp4_attention": {"two_level_p": True}} + models = dict(student_init_from="/ckpt/student", teacher_init_from="/ckpt/teacher") + inputs = step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **models) + metrics = {"eval/score": 0.31, "eval/x0_mse/video/rung0": 0.2, "eval/seconds": 6000.0, + "init/recon_step0/rung0": 0.04, "init/recon_step0/rung1": 0.03, + "init/noise_sigma/endpoint/audio": 0.05, "init/noise_sigma/score": 0.001} + path = tmp_path / "step0.json" + assert load_step0(path, inputs) is None + save_step0(path, inputs, metrics) + # Settings that cannot change step 0 (eval cadence) do not invalidate the cache. + same = step0_fingerprint_inputs({**method, "eval_every": 40}, data, qad, "ab" * 32, 30, **models) + cached = load_step0(path, same) + assert cached is not None and "eval/seconds" not in cached and cached["eval/score"] == 0.31 + summary, recon_norm, tolerances = restore_step0(cached, 2, recon=True, noise_floor=True, min_tolerance=0.02) + assert summary == {"score": 0.31, "x0_mse/video/rung0": 0.2} + assert recon_norm == [0.04, 0.03] + assert tolerances == {"endpoint/audio": pytest.approx(0.1), "score": 0.02} + # Training-plan keys only matter with provisional calibration (no amax_json). + assert load_step0(path, step0_fingerprint_inputs(method, {**data, "seed": 5}, qad, "ab" * 32, 30, **models)) + provisional = {**qad, "linears": {"skip_blocks": [42, 43], "provisional_calibration_rows": 2}} + for changed in (step0_fingerprint_inputs({**method, "init_student_dcp": "/runs/e3b/best"}, data, qad, "ab" * 32, 30, + **models), + step0_fingerprint_inputs(method, {**data, "eval_per_group": 1}, qad, "ab" * 32, 30, **models), + step0_fingerprint_inputs(method, data, {**qad, "linears": {"skip_blocks": []}}, "ab" * 32, 30, + **models), + step0_fingerprint_inputs(method, data, qad, "cd" * 32, 30, **models), + step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "student_init_from": "/x"}), + step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "teacher_init_from": "/x"}), + step0_fingerprint_inputs({**method, "audio_weight": 2.0}, data, qad, "ab" * 32, 30, **models), + step0_fingerprint_inputs(method, data, provisional, "ab" * 32, 30, **models)): + with pytest.raises(ValueError, match="other inputs"): + load_step0(path, changed) diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py new file mode 100644 index 000000000..d605792a2 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-forward quantization-aware distillation (QAD) of a FastH3 OmniRef PDD student. + +Student = the teacher's own weights + NVFP4 fake-quantized block linears +(Stage A calibrated static FFN scales, unit attention/gate scales) + the +deployed sparse-FP4 attention numerics. Loss = per-forward regression of the +denoiser output (x0) to the frozen bf16 teacher, at identical inputs on the +eight-forward PDD ladder; no DMD, no adversarial term. + +Each optimizer micro-step takes one OmniRef row: the teacher walks its PDD +trajectory (teacher forcing), or with probability ``rollout_fraction`` the +student walks its own (DAgger-style exposure-bias correction); at every +forward both models see the same state and the student's x0 error is +backpropagated at once, inside that forward's scope. + +Guardrails: step-0 checks (quantization disabled -> bit-identical to the +teacher; quantized error vs the Stage A PTQ number), per-(rung, modality) +loss normalization by the step-0 PTQ error, held-out evaluation every +``eval_every`` steps (teacher-forced x0 rel-L2 per forward, student-rollout +endpoint, latent keyframe fidelity), best-checkpoint requests, and a stop +request after ``patience`` consecutive evaluations worse than step 0 (teacher +alignment counts full evaluations only). The step-0 baseline and the guard +state are checkpointed (``method.qad_baseline``), so a resumed run keeps them. +""" +from __future__ import annotations + +import hashlib +import json +import math +import time +from pathlib import Path +from typing import Any + +import torch + +from fastvideo.distributed import get_world_group +from fastvideo.logger import init_logger +from fastvideo.train.methods.base import LogScalar, TrainingMethod +from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction +from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0, + step0_fingerprint_inputs) +from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import (MODALITIES, RungStats, keyframe_latents, + rel_l2, step_packed, summarize, + target_slices, x0_pair) +from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler + +logger = init_logger(__name__) + +# Teacher alignment: the student's own sample (same seed) and its keyframes vs the teacher's; full evaluations only. +ALIGNMENT_KEYS = ("endpoint/video", "endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last") +BASELINE_FORMAT = 1 + + +class _BaselineState: + """DCP-stateful step-0 baseline and guard state of an ``OmniRefPDDQADMethod`` (one JSON string).""" + + def __init__(self, method: OmniRefPDDQADMethod) -> None: + self.method = method + + def state_dict(self) -> dict[str, Any]: + return {"json": json.dumps(self.method._baseline_state(), sort_keys=True)} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.method._load_baseline_state(json.loads(state_dict["json"])) + + +def _unit_hash(*parts: Any) -> float: + digest = hashlib.sha256(":".join(str(p) for p in parts).encode()).digest() + return int.from_bytes(digest[:8], "little") / 2.0**64 + + +class OmniRefPDDQADMethod(TrainingMethod): + """Teacher-forced (+ optional student-rollout) per-forward x0 regression; see module docstring.""" + + def __init__(self, *, cfg: Any, role_models: dict[str, Any]) -> None: + super().__init__(cfg=cfg, role_models=role_models) + if "teacher" not in role_models: + raise ValueError("OmniRefPDDQADMethod needs a frozen 'teacher' role") + self.teacher = role_models["teacher"] + if getattr(self.teacher, "_trainable", False) or self.teacher.qad_linears: + raise ValueError("the teacher must be frozen and unquantized") + mc = self.method_config + self.weights = {"video": float(mc.get("video_weight", 1.0)), "audio": float(mc.get("audio_weight", 1.0))} + self.rollout_fraction = float(mc.get("rollout_fraction", 0.0)) + self.rollout_start_step = int(mc.get("rollout_start_step", 0)) + # Student forward+backward on this many of the 8 forwards per row (the teacher walks all of them). + self.rungs_per_row = int(mc.get("rungs_per_row", 0)) + # "output": per-forward x0 regression. "recon": layer-wise reconstruction only. + # "output+recon": x0 regression + lambda * reconstruction. + self.objective = str(mc.get("objective", "output")) + if self.objective not in ("output", "recon", "output+recon"): + raise ValueError(f"objective must be output, recon or output+recon, got {self.objective!r}") + recon_lambda = float(mc.get("recon_lambda", 1.0)) + self.output_weight = {"output": 1.0, "recon": 0.0, "output+recon": 1.0 / (1.0 + recon_lambda)}[self.objective] + self.recon_weight = { + "output": 0.0, + "recon": 1.0, + "output+recon": recon_lambda / (1.0 + recon_lambda) + }[self.objective] + if self.objective == "recon" and float(mc.get("rollout_fraction", 0.0)) > 0: + raise ValueError("reconstruction trains on teacher states only; set rollout_fraction: 0") + # Best checkpoint by this eval key (lower is better): "score" (teacher-forced x0 rel-L2) or + # "endpoint/video" (the student's own sample vs the teacher's: teacher alignment). + self.select_metric = str(mc.get("select_metric", "score")) + # Save each eval's final latents (student's own sample, teacher's) for side-by-side decoding. + self.eval_save_dir = mc.get("eval_save_dir") + # Full evaluation (with the student's own 8-step sample) every this many steps; teacher-forced only otherwise. + self.full_eval_every = int(mc.get("full_eval_every", 0)) or int(mc.get("eval_every", 50)) + # Start the student from another run's tagged weights (e.g. Stage 1 best) instead of the teacher's. + self.init_student_dcp = mc.get("init_student_dcp") + self._eval_iteration = 0 + self.eval_every = int(mc.get("eval_every", 50)) + self.patience = int(mc.get("patience", 3)) + self.fidelity_tolerance = float(mc.get("fidelity_tolerance", 0.05)) + self.ptq_reference_json = mc.get("ptq_reference_json") + self.ptq_tolerance = float(mc.get("ptq_tolerance", 0.05)) + self.require_bitwise_init = bool(mc.get("require_bitwise_init", True)) + self.max_eval_rows = int(mc.get("max_eval_rows", 0)) + self.save_best = bool(mc.get("save_best", True)) + # Reconstruction per row group (target video, target audio, conditioning): over the whole packed sequence + # audio is ~1.4% of rows and text tokens dominate the norm, so an ungrouped ratio barely trains audio. + self.recon_group_weights = dict( + mc.get("recon_group_weights") or { + "video": 1.0, + "audio": 1.0, + "condition": 0.25 + }) + # "step0": divide each forward's reconstruction loss by its step-0 value (as the x0 terms are), so + # recon_lambda means what it says; raw reconstruction gradients are ~1e3-1e4x smaller than x0's. + self.recon_normalize = str(mc.get("recon_normalize", "none")) + if self.recon_normalize not in ("none", "step0"): + raise ValueError(f"recon_normalize must be none or step0, got {self.recon_normalize!r}") + self.recon_norm_rows = int(mc.get("recon_norm_rows", 4)) + self.recon_norm: list[float] | None = None + # Full evaluations roll out (8 more student forwards) on at most this many rows; others are teacher-forced. + self.max_rollout_rows = int(mc.get("max_rollout_rows", 0)) + # A checkpoint is "best" only if these full-eval metrics stay <= step0 * (1 + tolerance). With + # noise_floor_eval the tolerance is 2 sigma of a seed+1 re-evaluation at step 0 (>= min_eligibility_tolerance); + # otherwise fidelity_tolerance. + self.eligibility_keys = tuple( + mc.get("eligibility_keys") + or ("endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last", "x0_rel_l2/audio/mean")) + self.noise_floor_eval = bool(mc.get("noise_floor_eval", False)) + self.min_eligibility_tolerance = float(mc.get("min_eligibility_tolerance", 0.02)) + self.tolerances: dict[str, float] = {} + # Stage A's per-rung reference metric for the step-0 PTQ match (v_rel is sign-independent). + self.ptq_reference_metric = str(mc.get("ptq_reference_metric", "v_rel_l2")) + self._recon_criteria: dict[str, Any] = {} + # Reuse (or write) this run's step-0 evaluation, keyed by a fingerprint of everything it depends on. + self.step0_cache = mc.get("step0_cache") + self._method_config_snapshot = mc + self.student.init_preprocessors(cfg.training) + tc = self.training_config + params = [p for p in self.student.transformer.parameters() if p.requires_grad] + self.optimizer, self.lr_scheduler = build_optimizer_and_scheduler(params=params, + optimizer_config=tc.optimizer, + loop_config=tc.loop, + learning_rate=float( + tc.optimizer.learning_rate), + betas=tuple(tc.optimizer.betas), + scheduler_name=str(tc.optimizer.lr_scheduler)) + self.normalizers: dict[str, list[float]] | None = None + self.step0: dict[str, float] = {} + self.best_score = math.inf + # Consecutive evaluations with a worse teacher-forced score (every eval) and with worse teacher alignment + # (full evals only); ``bad_evals`` is the larger. + self.bad_score_evals = 0 + self.bad_alignment_evals = 0 + self.bad_evals = 0 + self.stop_requested = False + self._checkpoint_request: str | None = None + if self.init_student_dcp: + import torch.distributed.checkpoint as dcp + + # The same DCP wrapper the checkpoint manager saves tagged/regular student weights with. + states = { + key: value + for key, value in self.checkpoint_state().items() if key == "roles.student.transformer" + } + if not states: + raise ValueError("init_student_dcp needs a trainable 'student' role with a transformer") + dcp.load(states, checkpoint_id=str(Path(self.init_student_dcp) / "dcp")) + logger.info("QAD student initialized from %s", self.init_student_dcp) + world = get_world_group() + self.sp_size = int(tc.distributed.sp_size or 1) + self.dp_rank, self.dp_size = world.rank // self.sp_size, world.world_size // self.sp_size + self.is_sp_leader = world.rank % self.sp_size == 0 + self.sampling_seed = int(tc.data.seed or 0) + + # ------------------------------------------------------------------ optimizer plumbing + def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]: + return [self.optimizer] + + def get_lr_schedulers(self, iteration: int) -> list[Any]: + return [self.lr_scheduler] + + @property + def _optimizer_dict(self) -> dict[str, Any]: + return {"student": self.optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self.lr_scheduler} + + def backward(self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int = 1) -> None: + """No-op: each forward's loss was backpropagated inside its own rung scope.""" + + def consume_checkpoint_request(self) -> str | None: + request, self._checkpoint_request = self._checkpoint_request, None + return request + + # ------------------------------------------------------------------ resume + def checkpoint_state(self) -> dict[str, Any]: + """Roles, optimizer and scheduler, plus the step-0 baseline and guard state. + + Without the baseline a resumed run would re-measure step 0 on its trained weights (wrong loss scale, + stop history and best-checkpoint reference); a checkpoint without ``method.qad_baseline`` cannot be + resumed (DCP reports the missing key). + """ + return {**super().checkpoint_state(), "method.qad_baseline": _BaselineState(self)} + + def _baseline_state(self) -> dict[str, Any]: + provisional = {} + if getattr(self.student, "provisional_rows", 0) > 0: + provisional = { + name: m.input_amax + for name, m in self.student.qad_linears.items() if m.act_scale == "static" and m.input_amax is not None + } + return { + "format": BASELINE_FORMAT, + "normalizers": self.normalizers, + "step0": self.step0, + "recon_norm": self.recon_norm, + "tolerances": self.tolerances, + "best_score": self.best_score, + "bad_score_evals": self.bad_score_evals, + "bad_alignment_evals": self.bad_alignment_evals, + "provisional_amax": provisional, + } + + def _load_baseline_state(self, state: dict[str, Any]) -> None: + if state.get("format") != BASELINE_FORMAT: + raise ValueError(f"unsupported QAD baseline state format {state.get('format')!r}") + if state["normalizers"] is None: + raise ValueError("the checkpoint has no step-0 baseline (saved before the step-0 evaluation)") + self.normalizers = {m: [float(v) for v in values] for m, values in state["normalizers"].items()} + self.step0 = {k: float(v) for k, v in state["step0"].items()} + self.recon_norm = None if state["recon_norm"] is None else [float(v) for v in state["recon_norm"]] + self.tolerances = {k: float(v) for k, v in state["tolerances"].items()} + self.best_score = float(state["best_score"]) + self.bad_score_evals = int(state["bad_score_evals"]) + self.bad_alignment_evals = int(state["bad_alignment_evals"]) + self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals) + for name, amax in state.get("provisional_amax", {}).items(): + method = self.student.qad_linears.get(name) + if method is None: + raise KeyError(f"checkpoint has a provisional scale for {name}, which is not a QAD linear here") + if method.input_amax is None: + method.set_input_amax(float(amax)) + if get_world_group().rank == 0: + logger.info("QAD step-0 baseline restored from the checkpoint (score %.5f, bad evals %d/%d)", + self.step0.get("score", float("nan")), self.bad_score_evals, self.bad_alignment_evals) + + # ------------------------------------------------------------------ one training row + def _row(self, spec: Any) -> Any: + return self.student.load_row(spec) + + def single_train_step(self, batch: dict[str, Any], iteration: int) -> tuple[dict, dict, dict[str, LogScalar]]: + if self.normalizers is None: + raise RuntimeError("step-0 evaluation must run before training (it sets the loss normalizers)") + spec = batch["row"] + prepared = self._row(spec) + rollout = self._use_rollout(spec.id, iteration) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1)) + num_rungs = self.student.num_rungs + trained = self._trained_rungs(spec.id, iteration, num_rungs) + video, audio = prepared.video, prepared.audio + total = torch.zeros((), device=self.student.device) + metrics: dict[str, LogScalar] = { + "train/rollout": float(rollout), + "train/tokens": float(prepared.layout.sequence_length) + } + started = time.perf_counter() + for rung in range(num_rungs): + recon = None + with torch.no_grad(), self.teacher.rung_scope(prepared, rung, timesteps) as rt: + if rung in trained and self.recon_weight > 0: + # Student blocks train inside the teacher forward, on the teacher's sub-layer inputs. + norm = self.recon_norm[rung] if self.recon_norm else 1.0 + scale = self.recon_weight / (norm * 2 * len(self.student.transformer.transformer_blocks) * + len(trained) * accum) + with BlockReconstruction(self.teacher.transformer, self.student.transformer, scale, + self._recon_criterion(prepared)) as recon: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + else: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + if recon is not None: + recon_loss = recon.loss() + if not torch.isfinite(recon_loss): + raise FloatingPointError(f"non-finite reconstruction loss at step {iteration}, rung {rung}") + norm = self.recon_norm[rung] if self.recon_norm else 1.0 + total += self.recon_weight * recon_loss / (norm * len(trained)) + metrics[f"train/recon_vs_step0/rung{rung}"] = recon_loss / norm + for key, value in recon.metrics("train/recon").items(): + metrics[f"{key}/rung{rung}"] = value + if rung not in trained or self.output_weight == 0: + if rollout: # the student's own trajectory still needs its (no-grad) forward here + with torch.no_grad(), self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, + *(student_out if rollout else teacher_out)) + continue + with self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + pair = x0_pair(prepared, schedulers, rung, video, audio, { + "student": student_out, + "teacher": teacher_out + }) + loss = self._rung_loss(pair, rung, metrics) + if not torch.isfinite(loss): + raise FloatingPointError(f"non-finite QAD loss at step {iteration}, row {spec.id}, rung {rung}") + (self.output_weight * loss / (len(trained) * accum)).backward() + total += self.output_weight * loss.detach() / len(trained) + step_out = tuple(v.detach() for v in student_out) if rollout else teacher_out + video, audio = step_packed(schedulers, prepared, rung, video, audio, *step_out) + metrics["train/row_seconds"] = time.perf_counter() - started + if get_world_group().rank == 0: + logger.info("QAD step %d row %s (%s, %d tokens, rollout=%s): loss %.5f in %.1fs", iteration, spec.id, + spec.case, prepared.layout.sequence_length, rollout, float(total), metrics["train/row_seconds"]) + return {"total_loss": total}, {}, metrics + + def _recon_criterion(self, prepared: Any) -> Any: + from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import (GroupedRelError, local_row_groups, + packed_row_groups) + key = prepared.spec.id + if key not in self._recon_criteria: + self._recon_criteria = {} # one row at a time + groups = packed_row_groups(prepared.layout, int(prepared.layout.sequence_length)) + self._recon_criteria[key] = GroupedRelError(local_row_groups(groups, self.student.device), + self.recon_group_weights) + return self._recon_criteria[key] + + @torch.no_grad() + def _calibrate_recon(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]: + """Step-0 reconstruction error per forward (measure-only), the Stage-2 normalizer.""" + num_rungs = self.student.num_rungs + sums = torch.zeros(num_rungs, dtype=torch.float64, device=self.student.device) + count = torch.zeros((), dtype=torch.float64, device=self.student.device) + per_rank = -(-self.recon_norm_rows // self.dp_size) + for spec, counted in rows[:per_rank]: + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(num_rungs): + with self.teacher.rung_scope(prepared, rung, + timesteps) as rt, BlockReconstruction(self.teacher.transformer, + self.student.transformer, + 0.0, + self._recon_criterion(prepared), + measure_only=True) as recon: + out = self.teacher.forward_rung(prepared, video, audio, rt) + if counted: + sums[rung] += recon.loss().double() + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + count += float(counted) + if not self.is_sp_leader: + sums.zero_() + count.zero_() + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(sums) + torch.distributed.all_reduce(count) + self.recon_norm = [max(float(v), 1e-12) for v in (sums / count.clamp_min(1.0))] + return {f"init/recon_step0/rung{rung}": value for rung, value in enumerate(self.recon_norm)} + + def _noise_floor(self, rows: list[tuple[Any, bool]], step0: dict[str, float]) -> dict[str, LogScalar]: + """Re-evaluate step 0 at seed+1; per-metric tolerance = max(2 sigma, minimum) with sigma relative.""" + import dataclasses + shifted = [(dataclasses.replace(spec, seed=spec.seed + 1), counted) for spec, counted in rows] + other = self._evaluate(shifted, rollout=True, save_tag="seed1") + out: dict[str, LogScalar] = {} + for key in set(self.eligibility_keys) | {"endpoint/video", "score"}: + if key in step0 and key in other: + mean = 0.5 * (step0[key] + other[key]) + sigma = abs(step0[key] - other[key]) / (math.sqrt(2.0) * max(mean, 1e-12)) + self.tolerances[key] = max(2.0 * sigma, self.min_eligibility_tolerance) + out[f"init/noise_sigma/{key}"] = sigma + out[f"init/seed1/{key}"] = other[key] + return out + + def _tolerance(self, key: str) -> float: + return self.tolerances.get(key, self.fidelity_tolerance) + + def _sampling_key(self, row_id: str) -> str: + """Keep forward/backward choices in lockstep across the FSDP shard group. + + FSDP spans data-parallel groups, whose rows differ. A row-dependent choice can make one group gather + teacher block weights while another gathers student block weights, corrupting both forwards. Single-group + runs retain their existing row-dependent sampling. + """ + return row_id if self.dp_size == 1 else f"fsdp-shared:{self.sampling_seed}" + + def _use_rollout(self, row_id: str, iteration: int) -> bool: + return (iteration >= self.rollout_start_step + and _unit_hash(self._sampling_key(row_id), iteration) < self.rollout_fraction) + + def _trained_rungs(self, row_id: str, iteration: int, num_rungs: int) -> set[int]: + """All forwards, or a deterministic subset shared by every FSDP rank.""" + if self.rungs_per_row <= 0 or self.rungs_per_row >= num_rungs: + return set(range(num_rungs)) + order = sorted(range(num_rungs), + key=lambda rung: _unit_hash(self._sampling_key(row_id), iteration, "rung", rung)) + return set(order[:self.rungs_per_row]) + + def _rung_loss(self, pair: dict[str, dict[str, torch.Tensor]], rung: int, metrics: dict[str, + LogScalar]) -> torch.Tensor: + assert self.normalizers is not None + loss = torch.zeros((), device=self.student.device) + for modality in MODALITIES: + student, teacher = pair[modality]["student"], pair[modality]["teacher"] + mse = (student - teacher).pow(2).mean() + normalized = mse / self.normalizers[modality][rung] + loss = loss + self.weights[modality] * normalized + metrics[f"train/rel_mse_vs_ptq/{modality}/rung{rung}"] = normalized.detach() + return loss / sum(self.weights.values()) + + # ------------------------------------------------------------------ evaluation and guardrails + def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]: + # Smoke runs without Stage A's table: calibrate the FFN scales before anything uses them, whether step 0 + # is measured, restored from the cache, or skipped by a resume. A no-op once every scale is set. + calibration = self._provisional_calibration() + first = self.normalizers is None + if not first and (self.eval_every <= 0 or iteration % self.eval_every): + return dict(calibration) + rows = self._rank_rows() + metrics: dict[str, LogScalar] = dict(calibration) + started = time.perf_counter() + self._eval_iteration = iteration + cached = self._load_step0() if first else None + if cached is not None: + summary = self._restore_step0(cached) + metrics.update(cached) + else: + if first and iteration == 0 and not self.init_student_dcp: + # Fresh start only: a resumed run (e.g. Stage 2 from a Stage-1 checkpoint) is no longer the teacher. + metrics.update(self._initial_checks(rows)) + if first and self._needs_recon_norm(): + metrics.update(self._calibrate_recon(rows)) + full = first or (iteration % self.full_eval_every == 0) + summary = self._evaluate(rows, rollout=full) + metrics.update({f"eval/{k}": v for k, v in summary.items()}) + if first and self.noise_floor_eval: + metrics.update(self._noise_floor(rows, summary)) + if first and self.step0_cache and get_world_group().rank == 0: + save_step0(self.step0_cache, self._step0_inputs(), {k: float(v) for k, v in metrics.items()}) + if first: + self.normalizers = { + m: [max(summary[f"x0_mse/{m}/rung{r}"], 1e-20) for r in range(self.student.num_rungs)] + for m in MODALITIES + } + self.step0 = dict(summary) + self.best_score = summary.get(self.select_metric, summary["score"]) + else: + metrics.update(self._guardrails(iteration, summary)) + metrics["eval/seconds"] = time.perf_counter() - started + metrics["eval/score_vs_ptq"] = summary["score"] / max(self.step0.get("score", summary["score"]), 1e-30) + if get_world_group().rank == 0: + logger.info("QAD eval @%d: score %.5f (PTQ %.5f) %s", iteration, summary["score"], + self.step0.get("score", float("nan")), + json.dumps({ + k: round(float(v), 5) + for k, v in metrics.items() + })) + return metrics + + def _needs_recon_norm(self) -> bool: + return self.recon_normalize == "step0" and self.recon_weight > 0 + + def _step0_inputs(self) -> dict[str, Any]: + return step0_fingerprint_inputs(self._method_config_snapshot, + self.student.data_config, + self.student.qad_config, + self.student.amax_sha256, + len(self.student.eval_rows), + student_init_from=self.student._init_from, + teacher_init_from=self.teacher._init_from) + + def _load_step0(self) -> dict[str, float] | None: + if not self.step0_cache: + return None + cached = load_step0(self.step0_cache, self._step0_inputs()) + if cached is not None and get_world_group().rank == 0: + logger.info("QAD step 0 restored from %s (no step-0 evaluation)", self.step0_cache) + return cached + + def _restore_step0(self, cached: dict[str, float]) -> dict[str, float]: + summary, recon_norm, tolerances = restore_step0(cached, self.student.num_rungs, self._needs_recon_norm(), + self.noise_floor_eval, self.min_eligibility_tolerance) + if recon_norm is not None: + self.recon_norm = recon_norm + self.tolerances.update(tolerances) + return summary + + def _rank_rows(self) -> list[tuple[Any, bool]]: + """This data-parallel group's eval rows, padded so every group runs the same number of forwards. + + FSDP all-gathers span every rank, so groups must stay in lockstep; padded + repeats run but are not counted (``(row, counted)`` pairs). + """ + rows = self.student.eval_rows[:self.max_eval_rows] if self.max_eval_rows else self.student.eval_rows + if not rows: + return [] + per_rank = -(-len(rows) // self.dp_size) + return [(rows[(i * self.dp_size + self.dp_rank) % len(rows)], i * self.dp_size + self.dp_rank < len(rows)) + for i in range(per_rank)] + + def _eligible(self, summary: dict[str, float]) -> bool: + """Full-eval audio/keyframe metrics within step0 * (1 + tolerance); a teacher-forced-only eval cannot tell.""" + if "endpoint/video" not in summary: + return False + return all(summary[key] <= self.step0[key] * (1.0 + self._tolerance(key)) for key in self.eligibility_keys + if key in self.step0 and key in summary) + + def _guardrails(self, iteration: int, summary: dict[str, float]) -> dict[str, LogScalar]: + worse_score = summary["score"] > self.step0["score"] + self.bad_score_evals = self.bad_score_evals + 1 if worse_score else 0 + # Teacher alignment is measured by full evaluations only; a teacher-forced-only one leaves its count as is. + measured = [key for key in ALIGNMENT_KEYS if key in summary and key in self.step0] + if measured: + misaligned = any(summary[key] > self.step0[key] * (1.0 + self._tolerance(key)) for key in measured) + self.bad_alignment_evals = self.bad_alignment_evals + 1 if misaligned else 0 + self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals) + selected = summary.get(self.select_metric) + eligible = self._eligible(summary) + if selected is not None and selected < self.best_score and eligible: + self.best_score = selected + self._checkpoint_request = "best" if self.save_best else None + if self.bad_evals >= self.patience: + self.stop_requested = True + logger.warning("QAD stop: %d consecutive evaluations worse than the PTQ start (step %d)", self.bad_evals, + iteration) + return { + "eval/eligible": float(eligible), + "eval/bad_evals": float(self.bad_evals), + "eval/bad_score_evals": float(self.bad_score_evals), + "eval/bad_alignment_evals": float(self.bad_alignment_evals), + "eval/best_score": self.best_score, + "eval/stop_requested": float(self.stop_requested) + } + + @torch.no_grad() + def _evaluate(self, + rows: list[tuple[Any, bool]], + *, + rollout: bool, + attention: bool = True, + linears: bool = True, + save_tag: str = "") -> dict[str, float]: + num_rungs = self.student.num_rungs + stats = RungStats(num_rungs) + counted_stats = stats + # Same rollout count on every data-parallel group (lockstep), covering the first max_rollout_rows rows. + rollout_rows = -(-self.max_rollout_rows // self.dp_size) if self.max_rollout_rows else len(rows) + with self.student.quantization(linears=linears, attention=attention): + for position, (spec, counted) in enumerate(rows): + stats = counted_stats if counted else RungStats(num_rungs) + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(num_rungs): + with self.teacher.rung_scope(prepared, rung, timesteps) as rt: + teacher_out = self.teacher.forward_rung(prepared, video, audio, rt) + with self.student.rung_scope(prepared, rung, timesteps) as rt: + student_out = self.student.forward_rung(prepared, video, audio, rt) + stats.add_rung( + rung, + x0_pair(prepared, schedulers, rung, video, audio, { + "student": student_out, + "teacher": teacher_out + })) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *teacher_out) + if rollout and position < rollout_rows: + self._rollout_metrics(stats, + prepared, + schedulers, + timesteps, + video, + audio, + save=counted, + save_tag=save_tag) + means = counted_stats.reduce(contribute=self.is_sp_leader) + summary = summarize(means, num_rungs, self.weights) + for name in ("first", "last"): + student, teacher = summary.get(f"keyframe/student/{name}"), summary.get(f"keyframe/teacher/{name}") + if student is not None and teacher is not None: + summary[f"keyframe/drop/{name}"] = student / max(teacher, 1e-12) - 1.0 + return summary + + def _rollout_metrics(self, + stats: RungStats, + prepared: Any, + schedulers: tuple[Any, Any], + timesteps: Any, + teacher_video: torch.Tensor, + teacher_audio: torch.Tensor, + save: bool = False, + save_tag: str = "") -> None: + """Student's own 8-forward sample vs the teacher's: endpoint rel-L2 and latent keyframe fidelity.""" + video, audio = prepared.video, prepared.audio + for rung in range(self.student.num_rungs): + with self.student.rung_scope(prepared, rung, timesteps) as rt: + out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + cut = target_slices(prepared) + if save and self.eval_save_dir and self.is_sp_leader: + step_dir = f"step{self._eval_iteration:05d}" + (f"_{save_tag}" if save_tag else "") + path = Path(self.eval_save_dir) / step_dir / f"{prepared.spec.id}.pt" + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "student_video": video[cut["video"]].float().cpu(), + "teacher_video": teacher_video[cut["video"]].float().cpu(), + # Audio for decoded stats (spectral centroid, >8 kHz share, dynamic range, log-mel distance). + "student_audio": audio[cut["audio"]].float().cpu(), + "teacher_audio": teacher_audio[cut["audio"]].float().cpu(), + "latent_shape": prepared.latent_shape, + "case": prepared.spec.case, + "seed": prepared.spec.seed + }, + path) + for modality, student, teacher in (("video", video, teacher_video), ("audio", audio, teacher_audio)): + s, t = student[cut[modality]].float(), teacher[cut[modality]].float() + stats.add_scalar(f"endpoint.{modality}", (s - t).norm() / t.norm().clamp_min(1e-12)) + teacher_frames = keyframe_latents(prepared, teacher_video) + for name, (frame, reference) in keyframe_latents(prepared, video).items(): + stats.add_scalar(f"keyframe.vs_teacher.{name}", rel_l2(frame, teacher_frames[name][0])) + if reference is not None: + stats.add_scalar(f"keyframe.student.{name}", rel_l2(frame, reference)) + stats.add_scalar(f"keyframe.teacher.{name}", rel_l2(teacher_frames[name][0], reference)) + + def _initial_checks(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]: + """Bit-identity with quantization off, and the PTQ error with bf16 attention (Stage A's protocol).""" + out: dict[str, LogScalar] = {} + diff = torch.zeros((), device=self.student.device, dtype=torch.float64) + if rows: + prepared = self._row(rows[0][0]) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + with torch.no_grad(), self.student.quantization(linears=False, attention=False): + for rung in (0, self.student.num_rungs - 1): + outs = [] + for model in (self.teacher, self.student): + with model.rung_scope(prepared, rung, timesteps) as rt: + outs.append(model.forward_rung(prepared, prepared.video, prepared.audio, rt)) + for teacher_t, student_t in zip(*outs, strict=True): + diff = torch.maximum(diff, (teacher_t.double() - student_t.double()).abs().max()) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(diff, op=torch.distributed.ReduceOp.MAX) + out["init/noquant_vs_teacher_max_abs"] = float(diff) + sha = getattr(self.student, "amax_sha256", "") + if sha: # first 32 bits of the pinned calibration table's sha256, for W&B + out["init/amax_sha256_u32"] = float(int(sha[:8], 16)) + if self.require_bitwise_init and float(diff) != 0.0: + raise RuntimeError(f"student with quantization disabled differs from the teacher (max |diff| " + f"{float(diff):.3e}); refusing to train from a student that is not the teacher") + weights_only = self._evaluate(rows, rollout=False, attention=False) + out.update({ + f"init/ptq_bf16_attn/{k}": v + for k, v in weights_only.items() if k.startswith(("x0_rel", "v_rel", "score")) + }) + out.update(self._compare_ptq_reference(weights_only)) + if getattr(self.student, "compression", None) is not None: + # Training-free token compression alone (bf16 linears): how much of the step-0 error it contributes. + compress_only = self._evaluate(rows, rollout=False, linears=False) + out.update({ + f"init/compress_only/{k}": v + for k, v in compress_only.items() if k.startswith(("x0_rel", "v_rel", "score")) + }) + return out + + @torch.no_grad() + def _provisional_calibration(self) -> dict[str, LogScalar]: + """Smoke runs without Stage A's table: max-calibrate the FFN inputs on dense teacher trajectories.""" + from fastvideo.layers.quantization.nvfp4_qad import calibrating, nvfp4_qad_export_scales + pending = [m for m in self.student.qad_linears.values() if m.act_scale == "static" and m.input_amax is None] + if not pending: + return {} + count = int(getattr(self.student, "provisional_rows", 0)) + rows = self.student.train_rows[:count] + per_rank = -(-len(rows) // self.dp_size) + mine = [rows[(i * self.dp_size + self.dp_rank) % len(rows)] for i in range(per_rank)] + with self.student.quantization(linears=False, attention=False), calibrating(self.student.qad_linears): + for spec in mine: + prepared = self._row(spec) + schedulers = self.student.schedulers() + timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) + video, audio = prepared.video, prepared.audio + for rung in range(self.student.num_rungs): + with self.student.rung_scope(prepared, rung, timesteps) as rt: + out = self.student.forward_rung(prepared, video, audio, rt) + video, audio = step_packed(schedulers, prepared, rung, video, audio, *out) + table = nvfp4_qad_export_scales(self.student.qad_linears) + if get_world_group().rank == 0: + path = Path(self.training_config.checkpoint.output_dir) / "provisional_amax.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(table, indent=1, sort_keys=True)) + logger.warning("PROVISIONAL FFN scales from %d rows written to %s; use Stage A's table for real runs", + len(rows), path) + return {"init/provisional_calibration_rows": float(len(rows))} + + def _compare_ptq_reference(self, weights_only: dict[str, float]) -> dict[str, LogScalar]: + """Compare step-0 weights-only x0 rel-L2 per rung to Stage A's (``{"video": [...], "audio": [...]}``).""" + if not self.ptq_reference_json: + return {} + reference = json.loads(Path(self.ptq_reference_json).read_text()) + deviations = [] + for modality in MODALITIES: + for rung, value in enumerate(reference.get(modality, [])): + ours = weights_only.get(f"{self.ptq_reference_metric}/{modality}/rung{rung}") + if ours is not None and value: + deviations.append(abs(ours / float(value) - 1.0)) + worst = max(deviations, default=float("nan")) + if deviations and worst > self.ptq_tolerance: + raise RuntimeError(f"step-0 PTQ error deviates {worst:.1%} from Stage A's ({self.ptq_reference_json}); " + "the QAD student does not start at the calibrated PTQ point") + return {"init/ptq_vs_stage_a_max_rel_dev": worst} + + +__all__ = ["OmniRefPDDQADMethod"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py new file mode 100644 index 000000000..b50bd6abe --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Decoded-audio statistics and the teacher-seed-floor gate for QAD checkpoints. + +Latent-space audio error does not show what the Trim DMD post-mortem found mattered: decoded audio drifted dull +(spectral centroid 952 Hz teacher -> 530 Hz late checkpoint) while losses looked fine. These statistics, compared +with how much the bf16 teacher itself varies between two seeds, make that drift a selection gate. +""" +from __future__ import annotations + +import math + +import torch + +STATS = ("spectral_centroid_hz", "share_above_8khz", "dynamic_range_db", "loudness_iqr_db") + + +def _spectrogram(wave: torch.Tensor, n_fft: int = 2048, hop: int = 512) -> torch.Tensor: + """Power spectrogram ``[freq, frames]`` of a mono or stereo (``[samples, 2]``) waveform.""" + mono = wave.float().mean(dim=-1) if wave.ndim == 2 else wave.float() + window = torch.hann_window(n_fft) + spec = torch.stft(mono, n_fft=n_fft, hop_length=hop, window=window, return_complex=True) + return spec.abs().pow(2) + + +def audio_stats(wave: torch.Tensor, sample_rate: int) -> dict[str, float]: + power = _spectrogram(wave) + freqs = torch.linspace(0, sample_rate / 2, power.shape[0]) + energy = power.sum(dim=0) + voiced = energy > energy.max() * 1e-6 + centroid = (power * freqs[:, None]).sum(0) / energy.clamp_min(1e-20) + frame_db = 10 * torch.log10(energy.clamp_min(1e-20)) + voiced_db = frame_db[voiced] if bool(voiced.any()) else frame_db + q = torch.quantile(voiced_db, torch.tensor([0.05, 0.25, 0.75, 0.95])) + return { + "spectral_centroid_hz": float(centroid[voiced].median()) if bool(voiced.any()) else 0.0, + "share_above_8khz": float(power[freqs > 8000].sum() / power.sum().clamp_min(1e-20)), + "dynamic_range_db": float(q[3] - q[0]), + "loudness_iqr_db": float(q[2] - q[1]), + } + + +def log_mel_distance(a: torch.Tensor, b: torch.Tensor, sample_rate: int, n_mels: int = 80) -> float: + """Mean absolute difference of log-mel spectrograms (triangular mel filterbank, no torchaudio needed).""" + pa, pb = _spectrogram(a), _spectrogram(b) + frames = min(pa.shape[1], pb.shape[1]) + n_freq = pa.shape[0] + + def mel(f: torch.Tensor) -> torch.Tensor: + return 2595 * torch.log10(1 + f / 700) + + edges = 700 * (10**(torch.linspace(0, float(mel(torch.tensor(sample_rate / 2))), n_mels + 2) / 2595) - 1) + freqs = torch.linspace(0, sample_rate / 2, n_freq) + lo, mid, hi = edges[:-2, None], edges[1:-1, None], edges[2:, None] + bank = torch.clamp(torch.minimum((freqs - lo) / (mid - lo), (hi - freqs) / (hi - mid)), min=0) + la = torch.log(bank @ pa[:, :frames] + 1e-8) + lb = torch.log(bank @ pb[:, :frames] + 1e-8) + return float((la - lb).abs().mean()) + + +def seed_floor_gate(student: list[dict[str, float]], + teacher: list[dict[str, float]], + teacher_other_seed: list[dict[str, float]], + slack: float = 1.0) -> dict[str, float | bool]: + """Per statistic: median |student - teacher| (same seed) must be <= slack * median |teacher(s+1) - teacher(s)|. + + Rows are aligned lists. Returns per-stat student gap, floor and a ``pass`` flag over all statistics. + """ + result: dict[str, float | bool] = {} + passed = True + for name in STATS: + gap = torch.tensor([abs(s[name] - t[name]) for s, t in zip(student, teacher, strict=True)]) + floor = torch.tensor([abs(o[name] - t[name]) for o, t in zip(teacher_other_seed, teacher, strict=True)]) + gap_m, floor_m = float(gap.median()), float(floor.median()) + result[f"{name}/student_gap"] = gap_m + result[f"{name}/teacher_floor"] = floor_m + ok = gap_m <= slack * floor_m or math.isclose(gap_m, 0.0, abs_tol=1e-9) + result[f"{name}/pass"] = ok + passed &= ok + result["pass"] = passed + return result + + +__all__ = ["STATS", "audio_stats", "log_mel_distance", "seed_floor_gate"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py new file mode 100644 index 000000000..7ed74bee6 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-forward (T1) error and reference-fidelity bookkeeping for PDD QAD. + +All quantities are on the target rows only (conditioning rows are inputs). +``x0`` is the scheduler's own denoised estimate, ``x_t + sigma_t * v`` with +``sigma_t = 1 - timestep``, so x0 error is ``sigma_t`` times velocity error. +""" +from __future__ import annotations + +from collections import defaultdict +from typing import Any + +import torch +import torch.distributed as dist + +MODALITIES = ("video", "audio") + + +def target_slices(prepared: Any) -> dict[str, slice]: + return { + "video": slice(prepared.layout.num_condition_video_rows, None), + "audio": slice(prepared.layout.num_condition_audio_rows, None) + } + + +def denoised(state: torch.Tensor, velocity: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + """x0 exactly as ``MiniMaxH3Scheduler.step`` forms it (fp32).""" + return state.float() + (1 - timestep.to(torch.float32)) * velocity.float() + + +def scheduler_step(scheduler: Any, rung: int, velocity: torch.Tensor, state: torch.Tensor) -> torch.Tensor: + """``scheduler.step`` at forward ``rung`` without relying on its internal counter.""" + scheduler._step_index = rung + return scheduler.step(velocity.float(), scheduler.timesteps[rung], state, return_dict=False)[0] + + +def step_packed(schedulers: tuple[Any, Any], prepared: Any, rung: int, video: torch.Tensor, audio: torch.Tensor, + video_v: torch.Tensor, audio_v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Advance the packed target rows by one PDD forward (conditioning rows untouched), as the sampler does.""" + cut = target_slices(prepared) + video, audio = video.clone(), audio.clone() + video[cut["video"]] = scheduler_step(schedulers[0], rung, video_v[cut["video"]], video[cut["video"]]) + audio[cut["audio"]] = scheduler_step(schedulers[1], rung, audio_v[cut["audio"]], audio[cut["audio"]]) + return video, audio + + +def x0_pair(prepared: Any, schedulers: tuple[Any, Any], rung: int, video: torch.Tensor, audio: torch.Tensor, + outputs: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> dict[str, dict[str, torch.Tensor]]: + """``{modality: {"student": x0, "teacher": x0, "student_v": v, "teacher_v": v}}`` on target rows.""" + cut = target_slices(prepared) + states = {"video": video, "audio": audio} + result: dict[str, dict[str, torch.Tensor]] = {} + for index, modality in enumerate(MODALITIES): + timestep = schedulers[index].timesteps[rung] + entry = {} + for role, (video_v, audio_v) in outputs.items(): + velocity = (video_v, audio_v)[index][cut[modality]] + entry[role] = denoised(states[modality][cut[modality]], velocity, timestep) + entry[f"{role}_v"] = velocity.float() + result[modality] = entry + return result + + +class RungStats: + """Sums of per-rung error terms, reduced over data-parallel groups (one SP rank contributes).""" + + KEYS = ("mse", "err_sq", "ref_sq", "verr_sq", "vref_sq", "dot", "s_sq") + EXTRA_KEYS = ("endpoint.video", "endpoint.audio", "keyframe.vs_teacher.first", "keyframe.vs_teacher.last", + "keyframe.student.first", "keyframe.teacher.first", "keyframe.student.last", "keyframe.teacher.last") + + def __init__(self, num_rungs: int) -> None: + self.num_rungs = num_rungs + self.sums: dict[str, torch.Tensor] = {} + self.counts: dict[str, float] = defaultdict(float) + # Ranks hold different rows (cases), so every rank reduces this same fixed key list. + self.all_keys = [f"{m}.{r}.{k}" for m in MODALITIES for r in range(num_rungs) for k in self.KEYS] + self.all_keys += list(self.EXTRA_KEYS) + + def _add(self, key: str, value: torch.Tensor) -> None: + self.sums[key] = self.sums.get(key, 0.0) + value.detach().double() + + def add_rung(self, rung: int, pair: dict[str, dict[str, torch.Tensor]]) -> None: + for modality, entry in pair.items(): + s, t, sv, tv = entry["student"], entry["teacher"], entry["student_v"], entry["teacher_v"] + prefix = f"{modality}.{rung}" + self._add(f"{prefix}.mse", (s - t).pow(2).mean()) + self._add(f"{prefix}.err_sq", (s - t).pow(2).sum()) + self._add(f"{prefix}.ref_sq", t.pow(2).sum()) + self._add(f"{prefix}.verr_sq", (sv - tv).pow(2).sum()) + self._add(f"{prefix}.vref_sq", tv.pow(2).sum()) + self._add(f"{prefix}.dot", (s * t).sum()) + self._add(f"{prefix}.s_sq", s.pow(2).sum()) + for key in self.KEYS: + self.counts[f"{prefix}.{key}"] += 1 + + def add_scalar(self, key: str, value: torch.Tensor | float) -> None: + if key not in self.EXTRA_KEYS: + raise KeyError(f"{key} is not a reduced RungStats key") + self._add(key, torch.as_tensor(value, dtype=torch.float64, device=self._device())) + self.counts[key] += 1 + + def _device(self) -> torch.device: + return torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + + def reduce(self, contribute: bool) -> dict[str, float]: + """All-reduce sums and counts (call on every rank in the same key order); returns means.""" + keys = self.all_keys + device = self._device() + values = torch.tensor([float(self.sums.get(k, 0.0)) for k in keys] + [self.counts.get(k, 0.0) for k in keys], + dtype=torch.float64, + device=device) + if not contribute: + values.zero_() + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(values) + sums, counts = values[:len(keys)], values[len(keys):] + return {key: float(sums[i]) / float(counts[i]) for i, key in enumerate(keys) if float(counts[i]) > 0} + + +def summarize(means: dict[str, float], num_rungs: int, weights: dict[str, float]) -> dict[str, float]: + """Per-rung rel-L2 / cosine / MSE plus the weighted mean x0 rel-L2 ``score`` (lower is better).""" + out: dict[str, float] = {} + score_num = score_den = 0.0 + for modality in MODALITIES: + rel_all = [] + for rung in range(num_rungs): + p = f"{modality}.{rung}" + if f"{p}.ref_sq" not in means: + continue + rel = (means[f"{p}.err_sq"] / max(means[f"{p}.ref_sq"], 1e-30))**0.5 + vrel = (means[f"{p}.verr_sq"] / max(means[f"{p}.vref_sq"], 1e-30))**0.5 + cos = means[f"{p}.dot"] / max((means[f"{p}.s_sq"] * means[f"{p}.ref_sq"])**0.5, 1e-30) + out[f"x0_rel_l2/{modality}/rung{rung}"] = rel + out[f"v_rel_l2/{modality}/rung{rung}"] = vrel + out[f"x0_cos/{modality}/rung{rung}"] = cos + out[f"x0_mse/{modality}/rung{rung}"] = means[f"{p}.mse"] + rel_all.append(rel) + if rel_all: + out[f"x0_rel_l2/{modality}/mean"] = sum(rel_all) / len(rel_all) + score_num += weights[modality] * out[f"x0_rel_l2/{modality}/mean"] + score_den += weights[modality] + out["score"] = score_num / max(score_den, 1e-30) + for key, value in means.items(): + if key.startswith(("endpoint.", "keyframe.")): + out[key.replace(".", "/")] = value + return out + + +def keyframe_latents(prepared: Any, final_video: torch.Tensor) -> dict[str, tuple[torch.Tensor, torch.Tensor | None]]: + """Generated keyframe latents ``[C, H, W]`` (first; last for first_last_frame) and their reference image. + + OmniRef encodes keyframe references at high resolution with the target's + aspect ratio, so the reference latent is area-resized to the target grid: a + rough proxy, logged only. The guarded keyframe metric compares the student's + keyframe with the teacher's (``keyframe/vs_teacher``). + """ + import torch.nn.functional as F + + from fastvideo.pipelines.basic.minimax_h3.packing import unpatchify_video_tokens + + channels, frames, height, width = prepared.latent_shape + patch = (1, 2, 2) + cut = target_slices(prepared)["video"] + video = unpatchify_video_tokens(final_video[cut][None].float(), frames, height, width, channels, patch)[0] + images: list[torch.Tensor | None] = [] + offset = 0 + for ref in prepared.references: + if ref.media_type == "audio": + continue + count = ref.num_latent_frames * (ref.latent_height // 2) * (ref.latent_width // 2) + image = None + same_aspect = abs(ref.latent_height * width - ref.latent_width * height) <= 0.02 * ref.latent_width * height + if ref.media_type == "image" and same_aspect: + rows = prepared.reference_video_rows[offset:offset + count].to(video.device) + latent = unpatchify_video_tokens(rows[None].float(), 1, ref.latent_height, ref.latent_width, channels, + patch)[:, :, 0] + if (ref.latent_height, ref.latent_width) != (height, width): + latent = F.interpolate(latent, size=(height, width), mode="area") + image = latent[0] + images.append(image) + offset += count + pairs = {"first_frame": [("first", 0, 0)], "first_last_frame": [("first", 0, 0), ("last", -1, 1)]} + return { + name: (video[:, frame], images[ref] if ref < len(images) else None) + for name, frame, ref in pairs.get(prepared.spec.case, []) + } + + +def rel_l2(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12) + + +__all__ = ["RungStats", "keyframe_latents", "rel_l2", "step_packed", "summarize", "target_slices", "x0_pair"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py new file mode 100644 index 000000000..5cda222cc --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Layer-wise reconstruction for the FastH3 OmniRef QAD student, streamed per block. + +Supervised units R, per transformer block: the attention sub-layer (q/k/v/gate/out NVFP4 projections, QK +RMSNorm, RoPE and the sparse FP4 VSA-128 attention, as deployed) and the FFN sub-layer (NVFP4 fc_in/fc_out). +Each student unit receives the frozen teacher's input to that unit (stop-gradient) and matches the teacher +unit's output, so quantization error is corrected locally and cannot accumulate into a drift away from the +teacher. + +Locality makes this cheap: inside one teacher forward, right after teacher block ``b`` runs, student block +``b`` runs on the captured inputs and is backpropagated at once. There is no end-to-end student forward or +backward, no activation recompute, and memory holds one block's activations. + +Attention is sparse (VSA-128, the teacher's trained selection), not dense as in the paper's warmup: our +teacher was trained sparse and the student is deployed sparse, so dense attention would supervise a +configuration that never runs. + +The loss per unit is the relative squared error ``||s - t||^2 / ||t||^2`` (late FFN outputs reach ~1e6, so an +unnormalized sum would only train those), averaged over R. With a ``GroupedRelError`` it is computed per row group +(target video, target audio, conditioning) and weighted: over the whole packed sequence audio is ~1.4% of the rows +and massive text tokens dominate ``||t||^2``, so an ungrouped ratio barely trains audio. +""" +from __future__ import annotations + +from typing import Any + +import torch + + +def ensure_fsdp_root_initialized(module: torch.nn.Module) -> None: + """Make ``module`` its FSDP2 tree's root before any of its sharded blocks runs on its own. + + FSDP2 picks the root lazily, at the first forward of any sharded module. If a block runs first (as in + reconstruction, which calls student blocks directly), that block claims the root and the model's next full + forward fails with "FSDP state has already been lazily initialized". Idempotent; a no-op without FSDP2. + """ + get_state = getattr(module, "_get_fsdp_state", None) + if get_state is not None: + get_state()._lazy_init() + + +def _unwrap(module: torch.nn.Module) -> torch.nn.Module: + return getattr(module, "_checkpoint_wrapped_module", module) + + +GROUPS = ("video", "audio", "condition") # target video rows, target audio rows, text + reference rows + + +def packed_row_groups(layout: Any, sequence_length: int) -> torch.Tensor: + """Group id per packed row: 0 target video, 1 target audio, 2 conditioning (text, references).""" + groups = torch.full((sequence_length, ), 2, dtype=torch.long) + groups[layout.video_indices[layout.num_condition_video_rows:]] = 0 + groups[layout.audio_indices[layout.num_condition_audio_rows:]] = 1 + return groups + + +def local_row_groups(groups: torch.Tensor, device: torch.device) -> torch.Tensor: + """This sequence-parallel rank's rows of ``groups`` as the transformer shards them; padding is -1.""" + from fastvideo.distributed import get_sp_world_size, model_parallel_is_initialized + from fastvideo.distributed.communication_op import sequence_model_parallel_shard + + shifted = (groups + 1).to(device) + if model_parallel_is_initialized() and get_sp_world_size() > 1: + shifted, _ = sequence_model_parallel_shard(shifted, dim=0) # pads with 0 -> -1 after the shift + return shifted - 1 + + +class GroupedRelError: + """Weighted mean over row groups of ``||s - t||^2 / ||t||^2`` with sequence-parallel-global sums. + + Each rank backpropagates ``sp_size * w_g * num_local / den_global``. Summed over ranks, ``w_g * num_local / + den_global`` is the global ratio, so a group living on one rank (e.g. target audio) is weighted exactly as + if unsharded. The ``sp_size`` factor matches the x0 output loss: FSDP averages gradients over every rank, + SP ranks included, and the output loss is identical on every SP rank, where the transformer's output + all-gather backward (a SUM reduce-scatter) gives each rank ``sp_size`` times its share of the gradient. + Reconstruction has no gather, so without the factor the averaged gradient would be ``1 / sp_size`` of the + global ratio's. Groups absent from the sequence are dropped from the weight normalization. + """ + + def __init__(self, groups: torch.Tensor, weights: dict[str, float]) -> None: + self.groups = groups + self.weights = torch.tensor([float(weights.get(name, 0.0)) for name in GROUPS], device=groups.device) + + def __call__(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(local loss to backprop, global per-group relative error [3], detached)``. + + The local loss is this rank's ``sp_size``-scaled share (see the class docstring); the detached global + per-group errors are what to log. + """ + from fastvideo.distributed import get_sp_group, get_sp_world_size, model_parallel_is_initialized + + teacher = teacher.float().reshape(-1, teacher.shape[-1]) + err = (student.float().reshape(-1, teacher.shape[-1]) - teacher).pow(2).sum(-1) + ref = teacher.pow(2).sum(-1) + one_hot = (self.groups[:, None] == torch.arange(len(GROUPS), device=err.device)[None, :]).float() + num = err @ one_hot # [3], differentiable + totals = torch.stack((num.detach(), ref @ one_hot)) # [2, 3] + sp_size = get_sp_world_size() if model_parallel_is_initialized() else 1 + if sp_size > 1: + totals = get_sp_group().all_reduce(totals) + present = totals[1] > 0 + den = totals[1].clamp_min(1e-30) + weights = torch.where(present, self.weights, torch.zeros_like(self.weights)) + norm = weights.sum().clamp_min(1e-30) + loss = sp_size * (weights * num / den).sum() / norm + return loss, torch.where(present, totals[0] / den, torch.full_like(den, float("nan"))) + + +class _Replace: + """Forward pre-hook replacing a sub-layer's first input with the teacher's, plus an output capture.""" + + def __init__(self) -> None: + self.value: torch.Tensor | None = None + self.output: torch.Tensor | None = None + + def pre(self, module: torch.nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...]: + return (self.value, *args[1:]) + + def post(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self.output = output + + +class _Capture: + + def __init__(self) -> None: + self.input: torch.Tensor | None = None + self.output: torch.Tensor | None = None + + def __call__(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self.input, self.output = args[0], output + + +class BlockReconstruction: + """Stream reconstruction losses for every block during one teacher forward. + + ``scale`` multiplies each block's loss before its backward (e.g. ``1 / (units x rungs x accum)``). + Use as a context manager around the teacher's ``forward_rung``, with grad disabled for the teacher. + """ + + def __init__(self, + teacher: torch.nn.Module, + student: torch.nn.Module, + scale: float, + criterion: GroupedRelError | None = None, + measure_only: bool = False) -> None: + ensure_fsdp_root_initialized(teacher) + ensure_fsdp_root_initialized(student) + self.teacher_blocks = list(teacher.transformer_blocks) + self.student_blocks = list(student.transformer_blocks) + if len(self.teacher_blocks) != len(self.student_blocks): + raise ValueError("teacher and student block counts differ") + self.scale = scale + self.criterion = criterion + self.measure_only = measure_only + self.attn_rel: list[torch.Tensor] = [] + self.ff_rel: list[torch.Tensor] = [] + self.group_rel: list[torch.Tensor] = [] # [3] per unit, global, detached + self._handles: list[Any] = [] + + def _weighted(self, rel: torch.Tensor) -> torch.Tensor: + """Group-weighted mean of per-group relative errors ``[..., 3]`` (NaN = absent group, dropped).""" + assert self.criterion is not None + weights = self.criterion.weights.expand_as(rel) * (~rel.isnan()) + return (weights * rel.nan_to_num()).sum(-1) / weights.sum(-1).clamp_min(1e-30) + + def _unit_loss(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(loss to backprop, sequence-parallel-global relative error to log, detached)``.""" + if self.criterion is None: + teacher_f = teacher.float() + loss = (student.float() - teacher_f).pow(2).mean() / teacher_f.pow(2).mean().clamp_min(1e-20) + return loss, loss.detach() + loss, per_group = self.criterion(student, teacher) + self.group_rel.append(per_group) + return loss, self._weighted(per_group) + + def __enter__(self) -> BlockReconstruction: + for index, block in enumerate(self.teacher_blocks): + inner = _unwrap(block) + attn, ff = _Capture(), _Capture() + self._handles.append(inner.attn.register_forward_hook(attn)) + self._handles.append(inner.ff.register_forward_hook(ff)) + self._handles.append(block.register_forward_hook(self._after_block(index, attn, ff))) + return self + + def __exit__(self, *exc: Any) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + def _after_block(self, index: int, attn: _Capture, ff: _Capture): + + def hook(module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None: + self._student_block(index, args, attn, ff) + attn.input = attn.output = ff.input = ff.output = None + + return hook + + def _student_block(self, index: int, block_args: tuple[Any, ...], attn: _Capture, ff: _Capture) -> None: + block = self.student_blocks[index] + inner = _unwrap(block) + replace_attn, replace_ff = _Replace(), _Replace() + replace_attn.value, replace_ff.value = attn.input, ff.input + handles = [ + inner.attn.register_forward_pre_hook(replace_attn.pre), + inner.attn.register_forward_hook(replace_attn.post), + inner.ff.register_forward_pre_hook(replace_ff.pre), + inner.ff.register_forward_hook(replace_ff.post) + ] + try: + with torch.set_grad_enabled(not self.measure_only): + # The block's own hidden input only feeds norm1 and the residual, whose results are unused; the + # teacher attention input has the right shape. The call goes through the block, and the block + # input and output join the graph (zero weight), so FSDP's pre-/post-backward hooks re-gather the + # block's weights for backward and reduce-scatter its gradients. + hidden = attn.input.detach().requires_grad_(not self.measure_only) + block_out = block(hidden, *block_args[1:]) + attn_loss, attn_rel = self._unit_loss(replace_attn.output, attn.output) + ff_loss, ff_rel = self._unit_loss(replace_ff.output, ff.output) + if not self.measure_only: + ((attn_loss + ff_loss) * self.scale + 0.0 * block_out.float().sum()).backward() + finally: + for handle in handles: + handle.remove() + self.attn_rel.append(attn_rel) + self.ff_rel.append(ff_rel) + + def loss(self) -> torch.Tensor: + """Mean (group-weighted) relative error over all units of this forward, sequence-parallel-global, detached.""" + if self.criterion is not None and self.group_rel: + return self._weighted(torch.stack(self.group_rel)).mean() # [units, 3], NaN for absent groups + return (torch.stack(self.attn_rel).sum() + torch.stack(self.ff_rel).sum()) / (2 * len(self.attn_rel)) + + def metrics(self, prefix: str) -> dict[str, torch.Tensor]: + attn, ff = torch.stack(self.attn_rel), torch.stack(self.ff_rel) + third = max(1, len(attn) // 3) + groups = {} + if self.group_rel: + per_group = torch.stack(self.group_rel).nanmean(dim=0) + groups = {f"{prefix}/{name}_rel": per_group[i] for i, name in enumerate(GROUPS)} + return { + **groups, f"{prefix}/attn_rel": attn.mean(), + f"{prefix}/ff_rel": ff.mean(), + f"{prefix}/attn_rel_first_third": attn[:third].mean(), + f"{prefix}/attn_rel_last_third": attn[-third:].mean(), + f"{prefix}/ff_rel_last_third": ff[-third:].mean() + } + + +__all__ = ["BlockReconstruction", "GROUPS", "GroupedRelError", "local_row_groups", "packed_row_groups"] diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py new file mode 100644 index 000000000..7587c0b33 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reuse a QAD run's step-0 evaluation (PTQ baseline, seed+1 noise floor, reconstruction normalizers). + +A Stage-2 run from a fixed checkpoint spends most of its first ~1.75 h on step 0: the full held-out evaluation, +the same evaluation at seed+1 for the noise floor, and the reconstruction calibration. All of it is a deterministic +function of the teacher and student checkpoints (``init_from``, plus ``init_student_dcp``), the calibration table +(or the provisional-calibration rows), the quantization plan, the evaluation set and the score weights, so a rerun +with the same inputs can load it. The cache is keyed by a fingerprint of exactly those inputs; any change (another +init checkpoint, eval rows, tolerances, modality weights, ...) is a mismatch and the run refuses to start. + +No torch import, so the fingerprint can be computed offline from a run's ``cfg.yaml``. +""" +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any + +STEP0_FORMAT = 2 +METHOD_KEYS = ("init_student_dcp", "max_eval_rows", "max_rollout_rows", "noise_floor_eval", "min_eligibility_tolerance", + "eligibility_keys", "fidelity_tolerance", "objective", "recon_normalize", "recon_norm_rows", + "recon_group_weights", "rungs_per_row") +DATA_KEYS = ("manifests", "cases", "max_frames", "eval_manifest", "eval_resolutions", "eval_per_group", + "eval_max_frames", "eval_seed") +# 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") +QAD_KEYS = ("linears", "fp4_attention") +# Logged per evaluation but not part of the step-0 result. +VOLATILE = ("eval/seconds", "eval/score_vs_ptq") + + +def _plain(value: Any) -> Any: + """OmegaConf / mapping / sequence -> JSON-native containers.""" + if hasattr(value, "items"): + return {str(k): _plain(v) for k, v in value.items()} + if isinstance(value, list | tuple) or (hasattr(value, "__iter__") and not isinstance(value, str | bytes)): + return [_plain(v) for v in value] + return value + + +def step0_fingerprint_inputs(method_config: Any, data_config: Any, qad_config: Any, amax_sha256: str, + num_eval_rows: int, *, student_init_from: str, teacher_init_from: str) -> dict[str, Any]: + """Everything step 0 depends on. The init paths stand for their weights and ``fastvideo_inference.json``.""" + method, data, qad = _plain(method_config), _plain(data_config), _plain(qad_config) + linears = qad.get("linears") or {} + provisional = int(linears.get("provisional_calibration_rows") or 0) > 0 and not linears.get("amax_json") + return { + "format": STEP0_FORMAT, + "models": { + "student_init_from": str(student_init_from), + "teacher_init_from": str(teacher_init_from) + }, + "method": { + key: method.get(key) + for key in METHOD_KEYS + }, + # ``summarize`` weights the score with these; best_score and the guards compare against it. + "score_weights": { + modality: float(method.get(f"{modality}_weight", 1.0)) + for modality in ("video", "audio") + }, + "data": { + key: data.get(key) + for key in DATA_KEYS + }, + "provisional_plan": { + key: data.get(key) + for key in PROVISIONAL_DATA_KEYS + } if provisional else None, + "qad": { + key: qad.get(key) + for key in QAD_KEYS + }, + "amax_sha256": amax_sha256, + "num_eval_rows": int(num_eval_rows), + } + + +def fingerprint(inputs: dict[str, Any]) -> str: + return hashlib.sha256(json.dumps(inputs, sort_keys=True).encode()).hexdigest() + + +def save_step0(path: str | Path, inputs: dict[str, Any], metrics: dict[str, float]) -> None: + payload = { + "format": STEP0_FORMAT, + "fingerprint": fingerprint(inputs), + "inputs": inputs, + "metrics": { + key: float(value) + for key, value in metrics.items() if key not in VOLATILE + }, + } + Path(path).parent.mkdir(parents=True, exist_ok=True) + tmp = Path(f"{path}.tmp") + tmp.write_text(json.dumps(payload, indent=1, sort_keys=True)) + tmp.replace(path) + + +def load_step0(path: str | Path, inputs: dict[str, Any]) -> dict[str, float] | None: + """The cached step-0 metrics, None if there is no cache; raises on a fingerprint mismatch.""" + path = Path(path) + if not path.exists(): + return None + payload = json.loads(path.read_text()) + expected = fingerprint(inputs) + if payload.get("format") != STEP0_FORMAT or payload.get("fingerprint") != expected: + cached = payload.get("inputs", {}) + diff = sorted(k for k in set(cached) | set(inputs) if cached.get(k) != inputs.get(k)) + raise ValueError(f"step-0 cache {path} was computed for other inputs (differs in {diff}); " + f"delete it or point step0_cache elsewhere") + return {key: float(value) for key, value in payload["metrics"].items()} + + +def restore_step0(metrics: dict[str, float], num_rungs: int, recon: bool, noise_floor: bool, + min_tolerance: float) -> tuple[dict[str, float], list[float] | None, dict[str, float]]: + """(eval summary, reconstruction normalizers, eligibility tolerances) from cached step-0 metrics.""" + summary = {key[len("eval/"):]: value for key, value in metrics.items() if key.startswith("eval/")} + if "score" not in summary: + raise ValueError("step-0 cache has no eval/score") + recon_norm = None + if recon: + recon_norm = [max(metrics[f"init/recon_step0/rung{rung}"], 1e-12) for rung in range(num_rungs)] + tolerances: dict[str, float] = {} + if noise_floor: + prefix = "init/noise_sigma/" + tolerances = { + key[len(prefix):]: max(2.0 * value, min_tolerance) + for key, value in metrics.items() if key.startswith(prefix) and math.isfinite(value) + } + if not tolerances: + raise ValueError("noise_floor_eval is on but the step-0 cache has no noise-floor sigmas") + return summary, recon_norm, tolerances diff --git a/fastvideo/train/models/minimax_h3/omniref_data.py b/fastvideo/train/models/minimax_h3/omniref_data.py new file mode 100644 index 000000000..7cae77c19 --- /dev/null +++ b/fastvideo/train/models/minimax_h3/omniref_data.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniRef (FastH3 Ref2VA) rows for per-forward distillation of the PDD student. + +A row is one OmniRef training parquet row (precomputed Qwen3-VL presentation, +clean reference latent rows and their geometry); nothing is re-encoded. The +initial packed state of a (row, seed) pair is built exactly as +``generate_omniref_latents.py`` and ``MiniMaxH3LatentPreparationStage`` build +it (reference rows noise-augmented to 0.999, target video/audio noise from the +same CPU generator, PDD start scaling), so trajectories, teacher latents and +Stage A's teacher-forced error are all reproducible from the row id and seed. + +Row selection reuses the generator's deterministic plan (per (case, +resolution): seeded shuffle, then interleave), so the first rows of each group +are the rows the OmniRef latent generator and Stage A calibration used. +Held-out sources are excluded at both resolutions. +""" +from __future__ import annotations + +import hashlib +import json +import random +import re +from collections import defaultdict +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +OMNIREF_CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot", + "full_scene_audio_image_reference") +RESOLUTIONS = ("480p", "768p") +_BUCKET_FRAMES = re.compile(r"bucket=\d+x\d+-(\d+)f") +ROW_COLUMNS = [ + "id", "vae_latent_shape", "audio_latent_shape", "text_embedding_bytes", "text_embedding_shape", + "text_token_tags_bytes", "text_token_tags_shape", "reference_video_latent_rows_bytes", + "reference_video_latent_rows_shape", "reference_audio_latent_rows_bytes", "reference_audio_latent_rows_shape", + "reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights", + "reference_latent_widths", "reference_num_audio_latents", "width", "height", "num_frames" +] + + +@dataclass(frozen=True) +class RowSpec: + id: str + source: str + case: str + resolution: str + parquet: str + seed: int + + +def clip_seed(base_seed: int, clip_id: str) -> int: + digest = hashlib.sha256(f"{base_seed}:{clip_id}".encode()).digest() + return int.from_bytes(digest[:4], "little") & 0x7FFFFFFF + + +def load_manifest_groups(paths: list[str], cases: tuple[str, ...], + max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]: + """Manifest entries grouped by (case, resolution); several manifests (v4 + v5 additions) merge. + + An entry repeated across manifests is kept once per (id, resolution): the same source at 480p and 768p is + two rows. + """ + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + seen: set[tuple[str, str]] = set() + for path in paths: + with open(path) as handle: + for line in handle: + entry = json.loads(line) + match = _BUCKET_FRAMES.search(entry["parquet"]) + key = (entry["id"], entry["resolution"]) + if (match is None or entry["case"] not in cases or entry["resolution"] not in RESOLUTIONS or key in seen + or int(match.group(1)) > max_frames): + continue + seen.add(key) + groups[(entry["case"], entry["resolution"])].append(entry) + return groups + + +def limit_frames(groups: dict[tuple[str, str], list[dict[str, Any]]], + max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]: + """``groups`` without entries whose bucket is longer than ``max_frames`` (0 = no limit); a new mapping.""" + if not max_frames: + return dict(groups) + + def frames(entry: dict[str, Any]) -> int: + match = _BUCKET_FRAMES.search(entry["parquet"]) + return int(match.group(1)) if match else max_frames + 1 + + return {key: [entry for entry in entries if frames(entry) <= max_frames] for key, entries in groups.items()} + + +def select_rows(groups: dict[tuple[str, str], list[dict[str, Any]]], + seed: int, + per_group: dict[str, int], + exclude_sources: frozenset[str] = frozenset()) -> list[RowSpec]: + """The generator's plan: per group a seeded shuffle, minus excluded sources, then interleaved.""" + picked = [] + for key in sorted(groups): + entries = sorted(groups[key], key=lambda entry: entry["id"]) + random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries) + entries = [entry for entry in entries if entry["id"].split(":", 1)[-1] not in exclude_sources] + picked.append(entries[:int(per_group.get(key[1], 0))]) + plan = [] + for index in range(max((len(entries) for entries in picked), default=0)): + for entries in picked: + if index < len(entries): + entry = entries[index] + source = entry["id"].split(":", 1)[-1] + clip_id = f"omniref-{entry['case']}-{source}-{entry['resolution']}" + plan.append( + RowSpec(clip_id, source, entry["case"], entry["resolution"], entry["parquet"], + clip_seed(seed, clip_id))) + return plan + + +def load_eval_rows(manifest_json: str, + cases: tuple[str, ...], + resolutions: tuple[str, ...], + per_group: int | dict[str, int], + max_frames: int = 0) -> list[RowSpec]: + """Held-out rows from a ``generate_omniref_latents.py --eval-dir`` manifest (ids, parquets and seeds).""" + items = json.loads(Path(manifest_json).read_text()) + counts: dict[tuple[str, str], int] = defaultdict(int) + rows = [] + for item in sorted(items, key=lambda item: item["id"]): + key = (item["case"], item["resolution"]) + limit = int(per_group.get(item["resolution"], 0)) if isinstance(per_group, dict) else int(per_group) + if item["case"] not in cases or item["resolution"] not in resolutions or counts[key] >= limit: + continue + if max_frames and int(item.get("num_frames", 0)) > max_frames: + continue + counts[key] += 1 + rows.append( + RowSpec(item["id"], item["source"], item["case"], item["resolution"], item["parquet"], int(item["seed"]))) + return rows + + +def top_up_eval_rows(eval_rows: list[RowSpec], groups: dict[tuple[str, str], list[dict[str, Any]]], seed: int, + per_group: dict[str, int]) -> list[RowSpec]: + """Fill each (case, resolution) up to ``per_group`` with fresh held-out rows (call before the training plan). + + Extra rows come from the generator's plan rule at ``seed``, skipping sources already in ``eval_rows``; + the training plan must then exclude every eval source. Result is interleaved across groups. + """ + have: dict[tuple[str, str], int] = defaultdict(int) + for row in eval_rows: + have[(row.case, row.resolution)] += 1 + taken = frozenset(row.source for row in eval_rows) + extra = [] + for key in sorted(groups): + need = int(per_group.get(key[1], 0)) - have[key] + if need > 0: + extra += select_rows({key: groups[key]}, seed, {key[1]: need}, taken) + return interleave_by_group(eval_rows + extra) + + +def interleave_by_group(rows: list[RowSpec]) -> list[RowSpec]: + """Round-robin over (case, resolution) groups, so any prefix covers every case.""" + buckets: dict[tuple[str, str], list[RowSpec]] = defaultdict(list) + for row in rows: + buckets[(row.case, row.resolution)].append(row) + ordered = [buckets[key] for key in sorted(buckets)] + return [ + bucket[i] for i in range(max((len(b) for b in ordered), default=0)) for bucket in ordered if i < len(bucket) + ] + + +class RowStream: + """Infinite, deterministic per-data-parallel-rank row stream (every SP rank of a group sees the same row).""" + + def __init__(self, plan: list[RowSpec], *, dp_rank: int, dp_size: int, seed: int) -> None: + if not plan: + raise ValueError("the training row plan is empty") + self.plan = plan + self.dp_rank = dp_rank + self.dp_size = dp_size + self.seed = seed + self.position = 0 + + def __iter__(self) -> Iterator[dict[str, Any]]: + while True: + epoch, offset = divmod(self.position * self.dp_size + self.dp_rank, len(self.plan)) + order = list(range(len(self.plan))) + random.Random(f"{self.seed}:epoch{epoch}").shuffle(order) + self.position += 1 + yield {"row": self.plan[order[offset]], "epoch": epoch} + + def state_dict(self) -> dict[str, int]: + return {"position": self.position} + + def load_state_dict(self, state: dict[str, int]) -> None: + self.position = int(state["position"]) + + +# --------------------------------------------------------------------------- one row -> initial packed state +def _tensor(row: dict[str, Any], name: str) -> torch.Tensor: + shape = [int(value) for value in row[f"{name}_shape"]] + array = np.frombuffer(row[f"{name}_bytes"], dtype=np.float32) + if array.size != int(np.prod(shape, dtype=np.int64)): + raise ValueError(f"{name}: {array.size} values do not fill shape {shape}") + return torch.from_numpy(array.reshape(shape).copy()) + + +def read_row(parquet: str) -> dict[str, Any]: + import pyarrow.parquet as pq + + rows = pq.read_table(parquet, columns=ROW_COLUMNS).to_pylist() + if len(rows) != 1: + raise ValueError(f"{parquet}: expected one row, got {len(rows)}") + return rows[0] + + +def prepared_references(row: dict[str, Any]) -> list[Any]: + from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference + + fields = ("reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights", + "reference_latent_widths", "reference_num_audio_latents") + return [ + MiniMaxH3PreparedReference(media_type=media, + has_audio=bool(audio), + num_latent_frames=int(frames), + latent_height=int(height), + latent_width=int(width), + num_audio_latents=int(count)) + for media, audio, frames, height, width, count in zip(*(row[name] for name in fields), strict=True) + ] + + +@dataclass +class PreparedRow: + """Everything one PDD trajectory of a row needs; ``video``/``audio`` are the packed initial rows (fp32).""" + spec: RowSpec + layout: Any + prompt_embeds: torch.Tensor + video: torch.Tensor + audio: torch.Tensor + latent_shape: tuple[int, int, int, int] # (channels, frames, height, width) of the target video + references: list[Any] + reference_video_rows: torch.Tensor # clean, CPU + + +def prepare_row(spec: RowSpec, row: dict[str, Any], *, patch_size: tuple[int, int, int], latent_channels: int, + audio_channels: int, pdd_start_sigmas: tuple[float, float], device: torch.device) -> PreparedRow: + """Initial packed state, bit-identical to ``OmniRefLatentGenerator._batch`` + ``scale_pdd_initial_noise``.""" + from diffusers.utils.torch_utils import randn_tensor + + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_NOISE_AUG, + audio_latent_num_frames, build_ref2va_packed_sequence, + keyframe_condition_noise, patchify_video_latents, + video_latent_num_frames) + + num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"]) + frames, lat_h, lat_w = video_latent_num_frames(num_frames), height // 16, width // 16 + expected = tuple(int(value) for value in row["vae_latent_shape"][1:]) + if (frames, lat_h, lat_w) != expected or int(row["vae_latent_shape"][0]) != latent_channels: + raise ValueError(f"{spec.id}: target geometry {(frames, lat_h, lat_w)} disagrees with {expected}") + num_audio = audio_latent_num_frames(num_frames) + references = prepared_references(row) + tags = _tensor(row, "text_token_tags").round().to(torch.long) + embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None] + layout = build_ref2va_packed_sequence(tags, references, frames, lat_h, lat_w, num_audio, patch_size) + + generator = torch.Generator("cpu").manual_seed(spec.seed) + shapes = tuple( + (ref.num_latent_frames, ref.latent_height, ref.latent_width) for ref in references if ref.media_type != "audio") + noise = keyframe_condition_noise(shapes, patch_size, latent_channels, generator=generator, device=device) + clean = _tensor(row, "reference_video_latent_rows") + condition_video = MiniMaxH3Scheduler().scale_noise(clean.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise) + video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w), + generator=generator, + device=device, + dtype=torch.float32) + audio_rows = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels), + generator=generator, + device=device, + dtype=torch.float32) + condition_audio = _tensor(row, "reference_audio_latent_rows").to(device) + video = torch.cat((condition_video, patchify_video_latents(video_noise, patch_size))) + audio = torch.cat((condition_audio, audio_rows)) + if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"{spec.id}: packed rows do not match the Ref2VA layout") + video[layout.num_condition_video_rows:] *= pdd_start_sigmas[0] + audio[layout.num_condition_audio_rows:] *= pdd_start_sigmas[1] + return PreparedRow(spec, layout, embeds, video, audio, (latent_channels, frames, lat_h, lat_w), references, clean) + + +__all__ = [ + "OMNIREF_CASES", "PreparedRow", "RowSpec", "RowStream", "interleave_by_group", "load_eval_rows", + "load_manifest_groups", "top_up_eval_rows", "prepare_row", "read_row", "select_rows" +] diff --git a/fastvideo/train/models/minimax_h3/omniref_pdd.py b/fastvideo/train/models/minimax_h3/omniref_pdd.py new file mode 100644 index 000000000..ca5bb436b --- /dev/null +++ b/fastvideo/train/models/minimax_h3/omniref_pdd.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastH3 OmniRef PDD transformer as a per-forward (rung) model for distillation. + +One instance wraps ``transformer_ref`` of a FastH3 OmniRef PDD export, loaded +through the training FSDP loader with VSA-H3 attention at the trained tile +size. It exposes the eight PDD forwards of the deployed sampler one at a +time: ``rung_scope`` installs exactly what ``MiniMaxH3DenoisingStage`` sets +for forward ``k`` (VSA-H3 metadata with the reference keep rate, the fused +PDD output heads, per-row timesteps), ``forward_rung`` runs the transformer, +and ``step`` applies the scheduler's Euler step. Teacher and student load +the same way, so a student whose quantization is disabled is the teacher. + +The student may carry NVFP4 QAD linears (``qad.linears``) and the sparse-FP4 +attention emulator (``qad.fp4_attention``); see ``fastvideo.layers. +quantization.nvfp4_qad`` and ``fastvideo.attention.backends.fp4_vsa_qat``. +""" +from __future__ import annotations + +import contextlib +import json +import re +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.activation_checkpoint import apply_activation_checkpointing, resolve_checkpointing_type +from fastvideo.train.utils.moduleloader import load_module_from_path + +if TYPE_CHECKING: + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + from fastvideo.train.utils.training_config import TrainingConfig + +logger = init_logger(__name__) + +# Block linears and norms; AdaLN, embeddings, the token refiner and the output heads stay frozen. +DEFAULT_TRAINABLE = (r"transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out|to_gate_compress|norm_q|norm_k)" + r"|norm1|norm2|ff\.(fc_in|fc_out))\.") +INFERENCE_FILE = "fastvideo_inference.json" + + +@dataclass(frozen=True) +class PDDContract: + """The checkpoint's deployed sampling contract (``fastvideo_inference.json``).""" + pdd_step_indices: tuple[int, ...] + video_shift: float + audio_shift: float + vsa_sparsity: float + vsa_tile_size: int + vsa_ref_keep_rate: float | None + + @classmethod + def read(cls, model_path: str) -> PDDContract: + data = json.loads((Path(model_path) / INFERENCE_FILE).read_text()) + if data.get("model_type") != "ref2va" or "pdd_step_indices" not in data: + raise ValueError(f"{model_path} is not a FastH3 Ref2VA PDD export") + return cls(tuple(int(i) for i in data["pdd_step_indices"]), float(data["video_scheduler_shift"]), + float(data["audio_scheduler_shift"]), float(data["vsa_sparsity"]), int(data["vsa_tile_size"]), + data.get("vsa_ref_keep_rate")) + + +class MiniMaxH3OmniRefPDDModel(ModelBase): + """Teacher or QAD student for per-forward regression on the OmniRef PDD ladder.""" + + _transformer_cls_name = "MiniMaxH3Transformer3DModel" + _transformer_module_type = "transformer_ref" + + def __init__(self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + enable_gradient_checkpointing_type: str | None = None, + trainable_patterns: str = DEFAULT_TRAINABLE, + qad: dict[str, Any] | None = None, + data: dict[str, Any] | None = None, + frozen_bf16: bool = False) -> None: + super().__init__(trainable=trainable, attention_backend="VIDEO_SPARSE_ATTN_H3") + # Keep a frozen role's masters in bf16 (the checkpoint's dtype) instead of an fp32 upcast: compute is bf16 either + # way, so outputs are bit-identical, and the role needs half the memory. + self._frozen_bf16 = bool(frozen_bf16) + if training_config.pipeline_config is None: + raise ValueError("MiniMaxH3OmniRefPDDModel requires a resolved MiniMax-H3 pipeline config") + if str(training_config.dit_precision) != "fp32" and trainable: + raise ValueError("QAD needs fp32 master weights (training.dit_precision: fp32); BF16 masters " + "round away updates at QAD learning rates") + self.contract = self._read_contract(init_from) + self._check_tile() + # Inference keeps proj_in/out, audio_proj_in/out and the time embedder in FP32; FSDP mixed precision would + # compute them in BF16 and the harness would no longer be the deployed operating point (E1: -7.6% v_rel at + # forward 0). They are frozen here, so they stay FP32 and replicated outside FSDP. + training_config.pipeline_config.dit_config.uniform_parameter_dtype = False # type: ignore[attr-defined] + training_config.pipeline_config.dit_config.replicate_fp32_modules = True # type: ignore[attr-defined] + self._init_from = str(init_from) + self.training_config = training_config + self.qad_config = dict(qad or {}) + self.data_config = dict(data or {}) + with self._load_precision(trainable): + self.transformer = load_module_from_path(model_path=self._init_from, + module_type=self._transformer_module_type, + training_config=training_config, + override_transformer_cls_name=self._transformer_cls_name, + attention_backend=self.attention_backend) + checkpointing = resolve_checkpointing_type(enable_gradient_checkpointing_type, training_config) + if trainable and checkpointing: + self.transformer = apply_activation_checkpointing(self.transformer, checkpointing_type=checkpointing) + self.num_trainable = self._apply_trainable(trainable_patterns if trainable else None) + self._check_fp32_modules_frozen() + self.patch_size = tuple(int(v) for v in self.transformer.patch_size) + self._plan: Any = None + self._vsa_builder: Any = None + self.qad_linears: dict[str, Any] = {} + self.fp4_numerics: Any = None + self.provisional_rows = 0 + self.amax_sha256 = "" + self.eval_rows: list[Any] = [] + self.train_rows: list[Any] = [] + if self.qad_config.get("enabled", False): + self._install_qad() + self.dataloader: Any = None + self.start_step = 0 + + # ------------------------------------------------------------------ setup + def _load_precision(self, trainable: bool) -> contextlib.AbstractContextManager[None]: + """Context for the transformer load: a frozen role may keep bf16 masters (``frozen_bf16``).""" + if trainable or not self._frozen_bf16: + return contextlib.nullcontext() + return self._bf16_load() + + @contextlib.contextmanager + def _bf16_load(self) -> Iterator[None]: + config = self.training_config + pipeline_config = config.pipeline_config + saved = (config.dit_precision, pipeline_config.dit_precision) + config.dit_precision = pipeline_config.dit_precision = "bf16" + try: + yield + finally: + config.dit_precision, pipeline_config.dit_precision = saved + + def _read_contract(self, model_path: str) -> Any: + return PDDContract.read(model_path) + + def _check_tile(self) -> None: + if self.contract.vsa_tile_size != 128: + raise ValueError(f"the FP4 VSA emulator models tile 128; this export uses {self.contract.vsa_tile_size}") + + def load_row(self, spec: Any) -> PreparedRow: + from fastvideo.train.models.minimax_h3.omniref_data import read_row + return self.prepare(spec, read_row(spec.parquet)) + + def _apply_trainable(self, pattern: str | None) -> int: + regex = re.compile(pattern) if pattern else None + count = 0 + for name, parameter in self.transformer.named_parameters(): + clean = name.replace("_checkpoint_wrapped_module.", "") + wanted = regex is not None and regex.search(clean) is not None + parameter.requires_grad_(wanted) + count += parameter.numel() if wanted else 0 + self.transformer.train(regex is not None) + if regex is not None and count == 0: + raise ValueError(f"trainable_patterns {pattern!r} matched no parameters") + return count + + def _check_fp32_modules_frozen(self) -> None: + """FP32 boundary modules are replicated outside FSDP: training them would desynchronize the ranks.""" + selector = getattr(self.transformer, "_get_parameter_dtype", None) + for name, parameter in self.transformer.named_parameters(): + clean = name.replace("_checkpoint_wrapped_module.", "") + if callable(selector) and selector(clean, torch.bfloat16) == torch.float32: + if parameter.requires_grad: + raise ValueError(f"{clean} is an FP32 replicated module and must stay frozen") + if parameter.dtype != torch.float32: + raise ValueError(f"{clean} should be FP32 (as inference), got {parameter.dtype}") + + def _install_qad(self) -> None: + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table + + linears = dict(self.qad_config.get("linears") or {}) + if linears.get("enabled", True): + amax_path = linears.get("amax_json") + table = load_amax_table(amax_path) if amax_path else None + if amax_path: + self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"), + self.training_config.checkpoint.output_dir) + plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)), + quantize_gate=bool(linears.get("gate", True)), + quantize_ffn=bool(linears.get("ffn", True)), + skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ()))) + # Smoke tests only: without Stage A's table, max-calibrate the FFN scales on a few rows at step 0. + self.provisional_rows = int(linears.get("provisional_calibration_rows", 0)) + if table is None and self.provisional_rows <= 0: + raise ValueError("qad.linears.amax_json (Stage A calibration) is required") + self.qad_linears = install_nvfp4_qad(self.transformer, table, plan, allow_uncalibrated=table is None) + attention = dict(self.qad_config.get("fp4_attention") or {}) + if attention.get("enabled", True): + floor = attention.get("first_block_max_floor") + self.fp4_numerics = FP4AttentionNumerics(quantize=True, + two_level_p=bool(attention.get("two_level_p", False)), + smooth_k=bool(attention.get("smooth_k", False)), + first_block_max_floor=None if floor is None else float(floor)) + elif self._trainable: + # Tile 128 has no other grad-capable kernel: train through the emulator with exact softmax. + self.fp4_numerics = FP4AttentionNumerics(quantize=False) + install_fp4_vsa_attention(self.transformer, self.fp4_numerics) + logger.info("OmniRef QAD student: %d NVFP4 QAD linears, attention %s, %.2fB trainable parameters", + len(self.qad_linears), self.fp4_numerics, self.num_trainable / 1e9) + + @contextlib.contextmanager + def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]: + """Temporarily enable/disable the student's fake quantization (all off = the teacher's numerics).""" + from fastvideo.attention.backends.fp4_vsa_qat import install_fp4_vsa_attention + from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled + + set_nvfp4_qad_enabled(self.qad_linears, linears) + if not attention: + install_fp4_vsa_attention(self.transformer, None) + try: + yield + finally: + set_nvfp4_qad_enabled(self.qad_linears, True) + install_fp4_vsa_attention(self.transformer, self.fp4_numerics) + + def init_preprocessors(self, training_config: TrainingConfig) -> None: + from fastvideo.train.models.minimax_h3.omniref_data import (OMNIREF_CASES, RowStream, interleave_by_group, + limit_frames, load_eval_rows, load_manifest_groups, + select_rows, top_up_eval_rows) + cfg = self.data_config + cases = tuple(cfg.get("cases", OMNIREF_CASES)) + resolutions = tuple(cfg.get("eval_resolutions", ("480p", ))) + eval_per_group = cfg.get("eval_per_group", 2) + eval_max_frames = int(cfg.get("eval_max_frames", 0)) + eval_rows = load_eval_rows(cfg["eval_manifest"], cases, resolutions, eval_per_group, eval_max_frames) + groups = load_manifest_groups(list(cfg["manifests"]), cases, int(cfg.get("max_frames", 243))) + missing = sorted(set(cases) - {case for case, _ in groups}) + if missing: + logger.warning("OmniRef cases with no rows in the given manifests: %s", missing) + if isinstance(eval_per_group, dict): + # More held-out rows than the precomputed eval manifest holds: draw them before the training plan, + # under the same eval length cap as the manifest rows. + eval_groups = limit_frames({ + key: entries + for key, entries in groups.items() if key[1] in resolutions + }, eval_max_frames) + eval_rows = top_up_eval_rows(eval_rows, eval_groups, int(cfg.get("eval_seed", 20261008)), + dict(eval_per_group)) + else: + eval_rows = interleave_by_group(eval_rows) + plan = select_rows(groups, int(cfg.get("seed", 20261007)), dict(cfg.get("per_group", {"480p": 400})), + frozenset(row.source for row in eval_rows)) + if int(cfg.get("max_train_rows", 0)) > 0: # smoke tests: overfit a few rows + plan = plan[:int(cfg["max_train_rows"])] + world = get_world_group() + sp_size = int(training_config.distributed.sp_size or 1) + self.eval_rows = eval_rows + self.train_rows = plan + self.dataloader = RowStream(plan, + dp_rank=world.rank // sp_size, + dp_size=world.world_size // sp_size, + seed=int(training_config.data.seed or 0)) + logger.info("OmniRef QAD rows: %d train (%s), %d held-out eval", len(plan), dict(sorted(_count(plan).items())), + len(eval_rows)) + + # ------------------------------------------------------------------ PDD ladder + @property + def plan(self) -> Any: + if self._plan is None: + from fastvideo.layers.pdd import PDDModalitySchedule, build_pdd_sampling_plan + schedules = { + "video": PDDModalitySchedule(shift=self.contract.video_shift), + "audio": PDDModalitySchedule(shift=self.contract.audio_shift) + } + self._plan = build_pdd_sampling_plan(self.contract.pdd_step_indices, schedules, device=self.device) + return self._plan + + @property + def num_rungs(self) -> int: + return int(self.plan.num_steps) + + def schedulers(self) -> tuple[Any, Any]: + """Fresh video/audio schedulers on the PDD node sigmas (as the denoising stage sets them).""" + from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + video = MiniMaxH3Scheduler(shift=self.contract.video_shift) + audio = MiniMaxH3Scheduler(shift=self.contract.audio_shift) + video.set_timesteps(sigmas=self.plan.node_sigmas["video"].to(torch.float32), device=self.device) + audio.set_timesteps(sigmas=self.plan.node_sigmas["audio"].to(torch.float32), device=self.device) + return video, audio + + def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow: + from fastvideo.train.models.minimax_h3.omniref_data import prepare_row + start = (float(self.plan.node_sigmas["video"][0]), float(self.plan.node_sigmas["audio"][0])) + return prepare_row(spec, + row, + patch_size=self.patch_size, + latent_channels=24, + audio_channels=32, + pdd_start_sigmas=start, + device=self.device) + + def _metadata(self, prepared: PreparedRow, rung: int) -> Any: + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (_h3_vsa_ref2va_segments, + _h3_vsa_single_region_segments) + if self._vsa_builder is None: + self._vsa_builder = MiniMaxH3VSAMetadataBuilder() + keep = self.contract.vsa_ref_keep_rate + segments = (_h3_vsa_single_region_segments(prepared.layout, self.patch_size) + if keep is None else _h3_vsa_ref2va_segments(prepared.layout, self.patch_size)) + return self._vsa_builder.build(current_timestep=rung, + patch_size=self.patch_size, + VSA_sparsity=self.contract.vsa_sparsity, + packed_segments=segments, + device=self.device, + exempt=True, + dense_layers=(), + tile_size=self.contract.vsa_tile_size, + ref_keep_rate=keep) + + @contextlib.contextmanager + def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]: + """Forward context of PDD forward ``rung``; keep backward inside it (activation recompute reads it).""" + from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps + video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item()) + unique, inverse = build_row_timesteps(prepared.layout, + video_timestep=video_t, + audio_timestep=audio_t, + condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG), + condition_audio_timestep=1.0) + start, end = self.plan.block(rung) + fusion = self.transformer.fuse_pdd_block(start, end, self.plan.integration_weights, torch.float32) + with fusion, set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared, rung)): + yield unique.to(self.device), inverse.to(self.device) + + def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor, + row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Packed velocities ``([Lv, Cv], [La, Ca])`` for one forward; call inside ``rung_scope``.""" + layout = prepared.layout + device = self.device + # No autocast: like inference, the transformer casts its inputs to each weight's dtype (BF16 under + # FSDP mixed precision), so Q/K/V reach attention in BF16 exactly as on the deployed path. Autocast + # would run RMSNorm/RoPE in fp32 and change what the FP4 quantizers see. + video_v, audio_v = self.transformer(hidden_states=video[None], + audio_hidden_states=audio[None], + encoder_hidden_states=prepared.prompt_embeds, + timestep=row_timesteps[0], + timestep_indices=row_timesteps[1], + token_tags=layout.token_tags.to(device), + position_ids=layout.position_ids.to(device), + video_indices=layout.video_indices.to(device), + audio_indices=layout.audio_indices.to(device), + text_indices=layout.text_indices.to(device)) + return video_v[0], audio_v[0] + + # ------------------------------------------------------------------ ModelBase contract (unused paths) + def prepare_batch(self, raw_batch: dict[str, Any], **kwargs: Any) -> Any: # noqa: D102 + raise NotImplementedError("use prepare() + rung_scope() + forward_rung()") + + def add_noise(self, clean_latents: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("PDD rungs start from the generator's noise; see prepare()") + + def predict_noise(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError("use forward_rung()") + + def backward(self, loss: torch.Tensor, ctx: Any, *, grad_accum_rounds: int) -> None: + raise NotImplementedError("the PDD QAD method backpropagates inside rung_scope()") + + def sp_group(self) -> Any: + return get_sp_group() + + +def _pin_amax(path: str, expected_sha256: str | None, output_dir: str) -> str: + """Check the calibration table's sha256 (when pinned) and keep a copy next to the run's checkpoints. + + The static FFN scales are fixed buffers that QAD learns around, so a run must never pick up a table that was + rewritten in place. + """ + import hashlib + import shutil + + digest = hashlib.sha256(Path(path).read_bytes()).hexdigest() + if expected_sha256 and digest != expected_sha256: + raise ValueError(f"{path} has sha256 {digest}, but the config pins {expected_sha256}") + if output_dir and get_world_group().rank == 0: + Path(output_dir).mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, Path(output_dir) / f"amax-{digest[:12]}.json") + logger.info("NVFP4 calibration table %s sha256 %s", path, digest) + return digest + + +def _count(rows: list[Any]) -> dict[str, int]: + counts: dict[str, int] = {} + for row in rows: + key = f"{row.case}/{row.resolution}" + counts[key] = counts.get(key, 0) + 1 + return counts + + +__all__ = ["DEFAULT_TRAINABLE", "MiniMaxH3OmniRefPDDModel", "PDDContract"] diff --git a/fastvideo/train/models/minimax_h3/v2_compression.py b/fastvideo/train/models/minimax_h3/v2_compression.py new file mode 100644 index 000000000..b4ab509ef --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_compression.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Training-side alias of the shared FastH3 token compression (the implementation serves inference too).""" +from fastvideo.models.dits.minimax_h3_token_compression import (MODES, CompressionSpec, TokenCompression, + block_flops_share, compressed_grid, grid_ops, + token_reduction) + +__all__ = [ + "MODES", "CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops", + "token_reduction" +] diff --git a/fastvideo/train/models/minimax_h3/v2_data.py b/fastvideo/train/models/minimax_h3/v2_data.py new file mode 100644 index 000000000..ad645fdcb --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_data.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Data-free text-to-video rows for FastH3 V2 QAD. + +A row is a prompt with a precomputed Qwen3-VL text embedding (``precompute_v2_text.py``) plus a seed. The initial +packed state is pure noise (target video + audio), drawn from one CPU generator so (row, seed) reproduces the same +trajectory everywhere; V2 is a plain 8-step DMD ladder, so there is no PDD start scaling and no conditioning rows. +The result mimics ``omniref_data.PreparedRow`` so the QAD method and metrics run unchanged (``case == "t2v"`` +yields no keyframe metrics). +""" +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +V2_CASE = "t2v" + + +@dataclass(frozen=True) +class V2RowSpec: + id: str + source: str + prompt: str + embedding: str # path of the .pt written by precompute_v2_text.py + seed: int + width: int = 832 + height: int = 480 + num_frames: int = 124 + case: str = V2_CASE + + @property + def resolution(self) -> str: + return f"{self.height}p" + + +def load_row_specs(manifest: str, width: int, height: int, num_frames: int) -> list[V2RowSpec]: + """Rows of a ``build_v2_prompts.py`` manifest (JSONL: id, source, prompt, embedding, seed).""" + specs = [] + for line in Path(manifest).read_text().splitlines(): + if not line.strip(): + continue + item = json.loads(line) + specs.append( + V2RowSpec(item["id"], item.get("source", ""), item["prompt"], item["embedding"], int(item["seed"]), + int(item.get("width", width)), int(item.get("height", height)), + int(item.get("num_frames", num_frames)))) + if not specs: + raise ValueError(f"{manifest} has no rows") + return specs + + +def read_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]: + data = torch.load(path, map_location="cpu", weights_only=True) + embeds, tags = data["embeds"], data["tags"] + if embeds.ndim != 2 or tags.shape != (embeds.shape[0], ): + raise ValueError(f"{path}: bad embedding/tag shapes {tuple(embeds.shape)} / {tuple(tags.shape)}") + return embeds, tags.to(torch.long) + + +def prepare_row(spec: V2RowSpec, *, patch_size: tuple[int, int, int], latent_channels: int, audio_channels: int, + device: torch.device) -> Any: + """Initial packed state of ``spec`` (fp32 noise rows), laid out exactly as ``MiniMaxH3LatentPreparationStage``.""" + from diffusers.utils.torch_utils import randn_tensor + + from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, audio_latent_num_frames, + build_packed_sequence, patchify_video_latents, + video_latent_num_frames) + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + + frames, lat_h, lat_w = video_latent_num_frames(spec.num_frames), spec.height // 16, spec.width // 16 + embeds, tags = read_embedding(spec.embedding) + num_audio = audio_latent_num_frames(spec.num_frames) + layout = build_packed_sequence(tags, frames, lat_h, lat_w, num_audio, patch_size) + generator = torch.Generator("cpu").manual_seed(spec.seed) + video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w), + generator=generator, + device=device, + dtype=torch.float32) + audio = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels), + generator=generator, + device=device, + dtype=torch.float32) + video = patchify_video_latents(video_noise, patch_size) + if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"{spec.id}: packed rows do not match the T2V layout") + return PreparedRow(spec, layout, + embeds.to(device=device, dtype=torch.bfloat16)[None], video, audio, + (latent_channels, frames, lat_h, lat_w), [], torch.empty(0, video.shape[-1])) + + +__all__ = ["V2RowSpec", "V2_CASE", "load_row_specs", "prepare_row", "read_embedding"] diff --git a/fastvideo/train/models/minimax_h3/v2_t2v.py b/fastvideo/train/models/minimax_h3/v2_t2v.py new file mode 100644 index 000000000..53049715e --- /dev/null +++ b/fastvideo/train/models/minimax_h3/v2_t2v.py @@ -0,0 +1,271 @@ +# 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 in bf16 (native Triton fwd+bwd, no FP4 +attention emulation), 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`` also switches +off the token compression (the student's only non-linear deviation), so ``quantization(False, False)`` is the 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))) + if self.compression is not None and self._compression_on: + grid = self._token_grid(prepared) + sub = compressed_grid(self.compression.spec.mode, grid) + stack.enter_context(self.compression.scope(grid, self._metadata(prepared, rung, sub))) + yield unique.to(self.device), inverse.to(self.device) + + def set_compression(self, spec: CompressionSpec | None) -> None: + """Replace the token-compression spec (probes and sweeps); ``None`` removes it.""" + if self.compression is not None: + self.compression.remove() + self.compression = None if spec is None else TokenCompression(self.transformer, spec) + + def compression_summary(self, prepared: PreparedRow) -> dict[str, float]: + """Token and FLOP accounting of the configured compression for one row.""" + spec = self.compression.spec if self.compression else None + grid = self._token_grid(prepared) + other = prepared.layout.sequence_length - grid[0] * grid[1] * grid[2] + out = block_flops_share(spec, grid, len(self.transformer.transformer_blocks), other) + out["video_tokens"] = float(grid[0] * grid[1] * grid[2]) + out["sequence_tokens"] = float(prepared.layout.sequence_length) + return out + + # ------------------------------------------------------------------ PDD-only paths of the parent + def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor, + row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + return super().forward_rung(prepared, video, audio, row_timesteps) + + +__all__ = ["MiniMaxH3V2T2VModel", "V2Contract"] diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index 3ac85179d..b96bacafe 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -281,14 +281,28 @@ class Trainer: method, iteration=step, ) + last_step = step + + # Opt-in method hooks: a tagged checkpoint (e.g. the best evaluation + # so far) and an early stop (e.g. guardrail patience exhausted). + 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) + if bool(getattr(method, "stop_requested", False)): + if self.global_rank == 0: + self.tracker.log({"stopped_early_at": float(step)}, step) + break + else: + last_step = max_steps self.callbacks.on_train_end( method, - iteration=max_steps, + iteration=last_step, ) if checkpoint_manager is not None: - checkpoint_manager.save_final(max_steps) + checkpoint_manager.save_final(last_step) self.performance_monitor.close() self.tracker.finish() diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 7a3d6c105..1d520fa4c 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -39,6 +39,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() @@ -273,6 +284,33 @@ class CheckpointManager: self._cleanup_old_checkpoints() + def save_tagged(self, tag: str, step: int) -> None: + """Save the trainable roles' model weights only to ``/`` (overwritten). + + Tagged directories are not ``checkpoint-`` and are never pruned; + they hold no optimizer state (resume from regular checkpoints). + """ + # A leading alphanumeric excludes "." and ".." (the run directory and its parent, which rmtree would delete). + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", tag) or _CHECKPOINT_DIR_RE.match(tag): + raise ValueError(f"invalid checkpoint tag {tag!r}") + tagged_dir = Path(self.output_dir) / tag + if tagged_dir.resolve().parent != Path(self.output_dir).resolve(): + raise ValueError(f"checkpoint tag {tag!r} does not name a directory directly inside {self.output_dir}") + dcp_dir = tagged_dir / "dcp" + if _rank() == 0: + if tagged_dir.exists(): + shutil.rmtree(tagged_dir) + os.makedirs(dcp_dir, exist_ok=True) + self._write_metadata(tagged_dir, step) + 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,