From 7066408fb31e7171abab94a0e1f26aea8cfc7af2 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:04:26 -0800 Subject: [PATCH] [feat] Pika H3 Flex NVFP4 QAD training: calibration, QAD, FP4 attention and export (#2051) Co-authored-by: Aryan Kumar --- docs-site/navigation.json | 4 + docs/contributing/env_vars.md | 3 + docs/training/pika_h3_flex_qad.md | 384 ++++++++++ .../inference/basic/basic_pika_h3_flex.py | 41 +- .../pika_h3_flex/amax-fl2va-balanced.json | 102 +++ .../pika_h3_flex/amax-ref2va-balanced.json | 102 +++ .../pika_h3_flex/fl2va_fp4attn_qad.yaml | 132 ++++ .../pika_h3_flex/fl2va_nvfp4_qad.yaml | 128 ++++ .../pika_h3_flex/ref2va_fp4attn_qad.yaml | 139 ++++ .../pika_h3_flex/ref2va_nvfp4_qad.yaml | 138 ++++ fastvideo/attention/backends/flash_attn.py | 12 +- fastvideo/attention/backends/fp4_dense_qat.py | 148 ++++ .../attention/utils/nvfp4_qk_precondition.py | 79 +++ fastvideo/envs.py | 19 + fastvideo/models/dits/minimax_h3.py | 13 + .../pipelines/basic/pika_h3_flex/adapter.py | 29 + .../pipelines/basic/pika_h3_flex/pipeline.py | 12 +- .../test_nvfp4_qk_transform_deploy.py | 129 ++++ .../tests/loader/test_pika_h3_flex_adapter.py | 32 + .../train/models/test_pika_h3_flex_qad.py | 576 +++++++++++++++ fastvideo/train/entrypoint/train.py | 4 + .../methods/knowledge_distillation/pdd_qad.py | 10 +- .../knowledge_distillation/pdd_qad_metrics.py | 9 +- fastvideo/train/models/minimax_h3/flex.py | 355 ++++++++++ .../train/models/minimax_h3/omniref_pdd.py | 5 +- fastvideo/train/models/minimax_h3/v2_t2v.py | 2 +- .../pika_h3_flex_qad/assemble_flex_release.py | 142 ++++ .../build_flex_ref2va_rows.py | 667 ++++++++++++++++++ .../pika_h3_flex_qad/calibrate_flex_nvfp4.py | 396 +++++++++++ .../check_fp4_attention_parity.py | 69 ++ .../check_nvfp4_export_parity.py | 131 ++++ .../pika_h3_flex_qad/export_student_bf16.py | 103 +++ .../pika_h3_flex_qad/merge_pika_adapter.py | 163 +++++ 33 files changed, 4264 insertions(+), 14 deletions(-) create mode 100644 docs/training/pika_h3_flex_qad.md create mode 100644 examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json create mode 100644 examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json create mode 100644 examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_fp4attn_qad.yaml create mode 100644 examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml create mode 100644 examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_fp4attn_qad.yaml create mode 100644 examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml create mode 100644 fastvideo/attention/backends/fp4_dense_qat.py create mode 100644 fastvideo/attention/utils/nvfp4_qk_precondition.py create mode 100644 fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py create mode 100644 fastvideo/tests/train/models/test_pika_h3_flex_qad.py create mode 100644 fastvideo/train/models/minimax_h3/flex.py create mode 100644 scripts/distill/pika_h3_flex_qad/assemble_flex_release.py create mode 100644 scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py create mode 100644 scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py create mode 100644 scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py create mode 100644 scripts/distill/pika_h3_flex_qad/check_nvfp4_export_parity.py create mode 100644 scripts/distill/pika_h3_flex_qad/export_student_bf16.py create mode 100644 scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py diff --git a/docs-site/navigation.json b/docs-site/navigation.json index 0db6ece0e..489ed306f 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -314,6 +314,10 @@ "label": "FastH3 NVFP4 QAD", "href": "/training/minimax_h3_fp4_qad" }, + { + "label": "Pika H3 Flex NVFP4 QAD", + "href": "/training/pika_h3_flex_qad" + }, { "label": "HunyuanVideo 1.5 T2V", "href": "/training/hunyuan15_overfit" diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 8b5ecdba0..2de43a4a3 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -179,6 +179,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_CACHE_ROOT` | path | computed | path | Root directory for FastVideo cache files. Defaults to ~/.cache/fastvideo, or $XDG_CACHE_HOME/fastvideo when XDG_CACHE_HOME is set. | | `FASTVIDEO_REASON1_WEIGHTS_PATH` | str | unset | path | Local path or Hugging Face id of Reason1 weights to load instead of the checkpoint's own. | | `FASTVIDEO_HOST_IP` | str | `""` | distributed | IP address of this node when the node has several network interfaces. Set it on each node for multi-node inference. | +| `FASTVIDEO_PG_TIMEOUT_MIN` | int | `0` | distributed | Timeout in minutes for the training process groups' collectives (0 keeps PyTorch's 10-minute default). Raise it for large models whose rank-0 load and scatter, checkpoint saves or long evaluations outlast the default. | | `FASTVIDEO_LOOPBACK_IP` | str | `""` | distributed | Loopback IP address to use instead of the detected one. | | `FASTVIDEO_RAY_PER_WORKER_GPUS` | float | `1.0` | distributed | GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share the GPUs with FastVideo. | | `FASTVIDEO_NCCL_SO_PATH` | str | unset | distributed | Path to the NCCL library file. Needed because the nccl>=2.19 that PyTorch ships has a bug (https://github.com/NVIDIA/nccl/issues/1234). | @@ -200,6 +201,8 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_VSA_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel; grad-tracking forwards need its autograd operator, else Triton. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. | | `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | | `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | +| `FASTVIDEO_NVFP4_FA4_QK_TRANSFORM` | one of none, smooth_k, hadamard, smooth_k+hadamard | `none` | attention | Exact Q/K preconditioning applied right before NVFP4 quantization on the FlashAttention-4 NVFP4 path: smooth_k subtracts each head's token-mean of K, hadamard rotates Q and K by the head-dim Hadamard matrix. Must match the qk_transform the model was QAD-trained with. An explicit nvfp4_qk_transform attention implementation argument takes precedence. | +| `FASTVIDEO_NVFP4_FA4_BF16_BLOCKS` | str | `""` | attention | Comma-separated MiniMax-H3 transformer_blocks indices whose attention stays in BF16 while FASTVIDEO_NVFP4_FA4 quantizes the rest, for example 45,49. | | `FASTVIDEO_FA4_PV_MODE` | one of bf16, fp8 | `bf16` | attention | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence. | | `FASTVIDEO_DISABLE_ATTENTION_COMPILE` | bool | `1` | attention | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile. | | `FASTVIDEO_MLX_WINDOW` | int | `0` | attention | MLX FastWan windowed attention size in tokens. 0 uses full attention. | diff --git a/docs/training/pika_h3_flex_qad.md b/docs/training/pika_h3_flex_qad.md new file mode 100644 index 000000000..7e6a8bbf3 --- /dev/null +++ b/docs/training/pika_h3_flex_qad.md @@ -0,0 +1,384 @@ +# Pika H3 Flex NVFP4 QAD + +This page describes how the NVFP4 version of +[Pika H3 Flex](https://huggingface.co/Pika-Labs/Pika-H3-Flex) is trained: a +calibration pass followed by quantization-aware distillation (QAD) against +the bf16 Flex model, for both Flex adapters. It applies the shared FastH3 QAD +recipe ([FastH3 NVFP4 QAD](minimax_h3_fp4_qad.md)) to Flex's six-forward +re-noise sampler. + +| Adapter | Trunk | Tasks | Model plugin | +|---|---|---|---| +| Ref2VA | `transformer_ref` | image / video / audio references | `flex.MiniMaxH3FlexRef2VAModel` | +| FL2VA | `transformer` | T2VA, I2VA (first frame), L2VA (last frame), FL2VA (both) | `flex.MiniMaxH3FlexFL2VAModel` | + +Training runs in two stages per adapter: + +1. **NVFP4 linears** (`*_nvfp4_qad.yaml`): W4A4 linears with BF16 dense + attention. +2. **Dense FP4 attention** (`*_fp4attn_qad.yaml`): continues from stage 1's + best checkpoint with FlashAttention-4 NVFP4 attention numerics. + +The stage-1 result is the released checkpoint below; stage 2 is optional and +trades a small quality cost for faster long clips (see +[Dense FP4 attention](#stage-2-dense-fp4-attention)). + +## Released checkpoint + +| Role | Checkpoint | +|---|---| +| Teacher (bf16 adapters) | `Pika-Labs/Pika-H3-Flex` on `MiniMaxAI/MiniMax-H3` @ `42ed227` | +| QAD result (packed NVFP4, both trunks) | `FastVideo/Pika-H3-Flex-NVFP4-QAD` | + +The release contains packed NVFP4 `transformer_ref` (Ref2VA) and `transformer` +(FL2VA) trunks with the adapters already merged, an NVFP4 Qwen3-VL text +encoder, an NVFP4 VAE decoder and the Pika schedulers. It targets Blackwell +GPUs and loads without a conversion step (validated on GB200): + +```bash +FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 \ +python examples/inference/basic/basic_pika_h3_flex.py \ + --model-path FastVideo/Pika-H3-Flex-NVFP4-QAD --nvfp4-dit h3_dit \ + --nvfp4-decoder /vae_decoder_nvfp4/decoder-nvfp4.safetensors \ + --task ref2va --reference image:character.png --reference video:motion.mp4 \ + --prompt "$(cat prompt.txt)" --resolution 480p --duration 5 +``` + +## Target numerics + +- **Linears:** NVFP4 W4A4 for `attn.to_q`, `attn.to_k`, `attn.to_v`, + `attn.to_out`, `ff.fc_in` and `ff.fc_out` in all 50 blocks, active from + step 0. FFN inputs use static activation scales (`448 * 6 / amax`) from the + calibrated table; attention linears use the unit activation scale. +- **Attention:** BF16 dense attention in stage 1. Flex has no VSA gate, so + there is no sparse-attention stage and no gate to freeze. Stage 2 switches + to the dense FP4 emulator (below). +- **Trained:** only the six quantized linear families. Norms, AdaLN, the token + refiner, embeddings and the fp32 boundary modules stay frozen. +- **Teacher:** the merged bf16 Flex trunk, frozen. The student starts from the + same weights with fp32 masters. + +### Merged adapters + +Flex ships each adapter as an fp32 PEFT LoRA (rank and alpha 128) over the +MiniMax-H3 trunk. QAD needs the teacher as plain weights, so +`merge_pika_adapter.py` merges `W + (alpha / rank) * B @ A` for every adapted +module in fp32 and rounds once to bf16. The merged directory also carries +`pika_flex_qad.json`, the sampling contract the model plugin reads (six +forwards at knots 1000/750/500/250/100/40, video shift 12, audio shift 3, +stochastic re-noise between forwards). + +Exported trunks that already contain the adapter are marked by +`pika_h3_flex_merged.json` at the model root; the Pika pipeline then skips +loading the adapter again. + +## Training objective + +Each optimizer step takes one row. The frozen teacher walks the six-forward +ladder; three of the six forwards are trained per row. At each trained forward +the student predicts from the teacher's state: + +- **Output:** relative L2 of the student's x0 prediction against the + teacher's, for video and audio. +- **Reconstruction:** relative L2 of every block's output against the + teacher's block output (`objective: output+recon`, `recon_lambda: 1.0`). + +Both terms are normalized by their step-0 values. From step 20, 25% of rows +roll out from the student's own states instead of the teacher's. The learning +rate is 5e-6 (cosine to 0.3x after a 10-step warmup); gradients are clipped at +1.0. The re-noise draws are seeded per row and per forward, so teacher and +student see identical noise. + +## Evaluation and guards + +Every `eval_every` steps the run scores held-out rows (teacher-forced x0 error +per forward, averaged into `eval/score`); every `full_eval_every` steps it also +samples each row with the student's own six forwards and compares the result +with the teacher's (`endpoint/video`, `endpoint/audio`) and, for keyframe +tasks, the first and last frames. + +- **Eligibility:** a checkpoint can become `best/` only if its full-eval + audio, endpoint and keyframe metrics stay within `step0 * (1 + tolerance)`. + With `noise_floor_eval: true` the tolerance per metric is twice its seed + noise at step 0 (re-evaluated at seed + 1), at least 2%. +- **Patience:** the run stops after `patience` consecutive evaluations that are + worse than step 0 on the score, or (full evaluations only) misaligned on the + endpoint and keyframe metrics. + +Keyframe metrics matter most for FL2VA: the last-frame match is the metric most +sensitive to the learning rate (see [Expected behavior](#expected-behavior)). + +## Data preparation + +QAD regresses on the teacher, so rows need conditioning (text embeddings and +reference latents) and target geometry, not curated targets. The starting +point is a set of encoded Ref2VA rows in the packed schema +(`fastvideo/dataset/dataloader/schema.py`, `pyarrow_schema_ref2va`) with a +`MANIFEST.jsonl` listing `case`, `id`, `parquet` and `resolution`, plus a +held-out manifest, as in the FastH3 OmniRef recipe (see +[Data preparation](minimax_h3_fp4_qad.md#data-preparation) there). Raw clips can +be encoded into this format with +`fastvideo/pipelines/preprocess/preprocess_minimax_h3_packed_ref2va.py`. Row +ids must contain the source video's id as `asset-<32 hex digits>`; clips of the +same video share it. Below, the source rows are `data/omniref/MANIFEST.jsonl` +and their held-out manifest is `data/omniref/eval-manifest.json`. + +1. **Merge each adapter** into its trunk: + + ```bash + python scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py \ + --base --adapter Pika-Labs/Pika-H3-Flex \ + --task ref2va --output data/models/pika-flex-ref2va-merged + python scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py \ + --base --adapter Pika-Labs/Pika-H3-Flex \ + --task fl2va --output data/models/pika-flex-fl2va-merged + ``` + +2. **Build Ref2VA rows** with `build_flex_ref2va_rows.py` (`plan` on CPU, + `encode` sharded over GPUs, then `finalize`). Each reference is decoded + from its stored latents and re-encoded at Flex's reference sizing, so the + source media is not needed. The plan re-tiers the rows the QAD config will + train on and adds input mixes the source set does not cover: a standalone + audio reference, mixed image and video references, 6-9 image references, + and two video references. + + `plan` reads the QAD plan from `--plan-config`. The Flex rows do not exist + yet at this point, so plan with a copy of `ref2va_nvfp4_qad.yaml` + (`plan-ref2va.yaml` below) whose `manifests` lists only + `data/omniref/MANIFEST.jsonl` and whose `eval_manifest` is + `data/omniref/eval-manifest.json`. `--eval-manifest` is the source held-out + manifest; its clips are kept out of the new rows. + + ```bash + python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py plan --mode ref2va \ + --source-manifest data/omniref/MANIFEST.jsonl --plan-config plan-ref2va.yaml \ + --eval-manifest data/omniref/eval-manifest.json --output-dir data/flex-ref2va + CUDA_VISIBLE_DEVICES=$i python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py encode \ + --model-path data/models/pika-flex-ref2va-merged --output-dir data/flex-ref2va \ + --shard $i --num-shards 8 + python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py finalize \ + --output-dir data/flex-ref2va --eval-manifest data/omniref/eval-manifest.json + ``` + + `finalize` merges the shard manifests into `data/flex-ref2va/MANIFEST.jsonl` + and writes `data/flex-ref2va/eval-manifest.json`: the source held-out rows, + with each row that was re-tiered pointing at its Flex-sized parquet. The + Ref2VA configs list the Flex manifest ahead of the source manifest, so + Flex-sized rows replace source rows with the same id and resolution. + +3. **Build FL2VA rows** from the same source rows. `--mode fl2va` builds T2VA, + I2VA, L2VA and FL2VA rows whose keyframes are the clip's exact first/last + frames, with the caption's reference paragraphs rewritten as keyframe + definitions. `--fl2va-counts` and `--resolution-split` size the set + (defaults: 1,100 rows, 80% 480p). FL2VA holds out its eval rows from its own + groups, so `finalize` takes no `--eval-manifest` and writes an empty one. + + ```bash + python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py plan --mode fl2va \ + --source-manifest data/omniref/MANIFEST.jsonl \ + --eval-manifest data/omniref/eval-manifest.json --output-dir data/flex-fl2va + CUDA_VISIBLE_DEVICES=$i python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py encode \ + --model-path data/models/pika-flex-fl2va-merged --output-dir data/flex-fl2va \ + --shard $i --num-shards 8 + python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py finalize --output-dir data/flex-fl2va + ``` + +Rows that fail to encode (for example a reference video longer than Flex's +15-second limit) are recorded per shard; `encode --retry-skipped` retries only +those. The configs train on rows up to 345 frames at 480p and 768p +(`per_group: {480p: 160, 768p: 40}` per case) and hold out 3 rows per case. + +## Calibration + +`calibrate_flex_nvfp4.py` picks the static FFN activation scales on rows walked +by the bf16 teacher over its six forwards: + +1. `collect` (torchrun): forward pre-hooks record each row's per-layer input + amax, split into target video, target audio and conditioning tokens. + `--max-rows` takes the first rows of the plan; leave rows uncollected for + `mse`, which stops if every planned row was collected. +2. `candidates` (CPU): per layer, the max over rows plus clipped alternatives + (per-row p99/p90, and 0.7/0.5/0.35/0.25 of the max). +3. `mse` (torchrun): on disjoint rows plus the rows that set the maxima, each + candidate's NVFP4 GEMM output is compared with the bf16 output. +4. `choose` (CPU): per layer, the candidate with the lowest mean of video and + audio relative error (`amax-mse-balanced.json`), next to the all-token + criterion (`amax-mse.json`) and the plain max. + +```bash +torchrun --nproc_per_node 4 scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py collect \ + --config examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml \ + --output-dir calib/ref2va --shard $NODE --num-shards $NODES --max-rows $COLLECT_ROWS \ + --training.distributed.num_gpus 4 --training.distributed.hsdp_shard_dim 4 +python scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py candidates --output-dir calib/ref2va +torchrun --nproc_per_node 4 scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py mse \ + --config --output-dir calib/ref2va --shard $NODE --num-shards $NODES \ + --training.distributed.num_gpus 4 --training.distributed.hsdp_shard_dim 4 +python scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py choose --output-dir calib/ref2va +``` + +Compare the candidate tables by the QAD harness's step-0 score (run the QAD +config with each table and read `eval/score` at step 0) and pin the best one. +For both adapters the video/audio-balanced table scored best: + +| Adapter | Calibration rows | Step-0 score: max / all-token MSE / balanced | +|---|---|---| +| Ref2VA | 1,133 | 0.3311 / 0.3299 / **0.3260** | +| FL2VA | 546 | 0.2730 / 0.2741 / **0.2724** | + +The selected tables are vendored and pinned by sha256: + +| Table | sha256 | +|---|---| +| `examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json` | `ad60b45f…da90b373` | +| `examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json` | `b690d117…553e0daa` | + +## Running + +Point the placeholder paths in the config at the merged model, the row +manifests and an output directory, then launch with `torchrun`. The stage-1 +configs assume 16 GPUs and the stage-2 configs 12 (sequence parallel 4, +HSDP over all GPUs); the released FL2VA stage-1 run used 24. Rows reach about +58k tokens; on GB200 the runs use about 90-100 GB per GPU. + +```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/pika_h3_flex/ref2va_nvfp4_qad.yaml +``` + +Loading and the step-0 evaluation of a multi-node run can exceed the default +10-minute process-group timeout; set `FASTVIDEO_PG_TIMEOUT_MIN=120`. Any config +key can be overridden on the command line (`--training.distributed.num_gpus 24 +--training.distributed.hsdp_shard_dim 24` for 6 nodes, for example). Stage 2 +launches the same way with `ref2va_fp4attn_qad.yaml` or `fl2va_fp4attn_qad.yaml` +once stage 1 has written `best/`. + +## Stage 2: dense FP4 attention + +Deployment can additionally run attention through FlashAttention-4's NVFP4 +path: Q and K are quantized to NVFP4 per 16 channels with a unit global scale, +P and V stay BF16. Stage 2 adapts the stage-1 student to these numerics with +`fp4_dense_qat.py`, which routes every dense attention layer through the FP4 +attention training kernel with full block lists and straight-through +gradients. `scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py` +checks that the emulator matches the FA4 kernel (on a Blackwell GPU): the +emulator-to-kernel distance is under 1% of the FP4 quantization error. + +Plain FP4 attention costs 13-17% on the step-0 score. Two exact Q/K transforms +and two BF16 blocks recover most of it before any training: + +- **smooth-K:** subtract each head's token-mean of K before quantizing + (softmax is invariant to a per-query shift). +- **Hadamard:** rotate Q and K by the same orthonormal 128-point Hadamard + matrix (`Q H (K H)^T = Q K^T`), spreading outlier channels across the + 16-channel scale blocks. +- **BF16 blocks 45 and 49:** the two attention layers with the largest FP4 + error keep BF16 attention. + +| Per-layer FP4 attention error vs BF16 (real activations) | FL2VA | Ref2VA | +|---|---|---| +| Plain FP4 | 0.115 | 0.113 | +| smooth-K + Hadamard | 0.095 | 0.095 | +| + blocks 45 and 49 in BF16 | 0.083 | 0.080 | + +The stage-2 configs set `qk_transform: smooth_k+hadamard` and +`bf16_blocks: [45, 49]`, start from stage 1's `best/` +(`method.init_student_dcp`) and train at 5e-6 (Ref2VA) or 3e-6 (FL2VA). Deploy +a stage-2 checkpoint with the same settings: + +```bash +python examples/inference/basic/basic_pika_h3_flex.py ... \ + --nvfp4-attention smooth_k+hadamard --nvfp4-attention-bf16-blocks 45,49 +``` + +(equivalently `FASTVIDEO_NVFP4_FA4=1`, +`FASTVIDEO_NVFP4_FA4_QK_TRANSFORM=smooth_k+hadamard`, +`FASTVIDEO_NVFP4_FA4_BF16_BLOCKS=45,49`). + +Measured on one GB200 (FL2VA, NVFP4 DiT and decoder, warm), FP4 attention +helps long clips and is neutral on short ones: + +| Clip | BF16 attention | FP4 attention | FP4 + transforms | +|---|---|---|---| +| 480p, 5 s | 5.87 s | 5.66 s | 6.08 s | +| 768p, 15 s | 99.7 s | 84.7 s | 87.4 s | + +The transforms currently run as separate PyTorch ops before the FP4 quantizer; +fusing them into the quantizer removes most of their cost. + +## Export + +Checkpoints hold fp32 masters in DCP format. Export a checkpoint to the packed +NVFP4 format the inference loader reads, then verify it: + +```bash +# 1. bf16 trunk from the student's masters (CPU) +python scripts/distill/pika_h3_flex_qad/export_student_bf16.py \ + --dcp outputs/pika_flex_ref2va_nvfp4_qad/best/dcp \ + --base data/models/pika-flex-ref2va-merged --trunk transformer_ref \ + --out exports/ref2va/transformer_bf16 +# 2. pack NVFP4 linears with the pinned activation scales +python scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py \ + --src exports/ref2va/transformer_bf16 --dst exports/ref2va/transformer_ref \ + --quantize-ffn --quantize-attention \ + --act-amax examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json +# 3. byte-level parity with the QAD numerics (1 GPU) +python scripts/distill/pika_h3_flex_qad/check_nvfp4_export_parity.py \ + --export exports/ref2va/transformer_ref --bf16 exports/ref2va/transformer_bf16 \ + --amax examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json +``` + +The parity check confirms that packed weights and block scales are +byte-identical to what training quantized and that the deployed NVFP4 GEMM +matches the training GEMM bit for bit. Repeat for FL2VA with +`--trunk transformer` and its table, then assemble a model directory: + +```bash +python scripts/distill/pika_h3_flex_qad/assemble_flex_release.py --output release/Pika-H3-Flex-NVFP4 \ + --base --pika \ + --transformer-ref exports/ref2va/transformer_ref --ref2va-merged data/models/pika-flex-ref2va-merged \ + --transformer exports/fl2va/transformer --fl2va-merged data/models/pika-flex-fl2va-merged +``` + +`--text-encoder` and `--vae-decoder-nvfp4` add NVFP4 text-encoder and decoder +exports. + +## Code map + +| Path | Purpose | +|---|---| +| `fastvideo/train/models/minimax_h3/flex.py` | Flex model plugins: sampling contract, re-noise ladder, FL2VA keyframe rows, QAD install | +| `fastvideo/train/methods/knowledge_distillation/pdd_qad.py` | Shared QAD method: losses, evaluation, guards | +| `fastvideo/attention/backends/fp4_dense_qat.py` | Dense FP4 attention emulator and install hook | +| `fastvideo/attention/utils/nvfp4_qk_precondition.py` | Exact Q/K transforms shared by training and inference | +| `fastvideo/attention/backends/flash_attn.py` | Deployed FA4 NVFP4 attention (`nvfp4_fa4`, `nvfp4_qk_transform`) | +| `fastvideo/pipelines/basic/pika_h3_flex/adapter.py` | Adapter loading and the merged-adapter marker | +| `scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py` | Merge an adapter into its trunk | +| `scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py` | Ref2VA and FL2VA training rows | +| `scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py` | FFN activation-scale calibration | +| `scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py` | Emulator vs FA4 NVFP4 kernel parity | +| `scripts/distill/pika_h3_flex_qad/export_student_bf16.py` | DCP masters to a bf16 trunk | +| `scripts/distill/pika_h3_flex_qad/check_nvfp4_export_parity.py` | Packed export vs QAD numerics parity | +| `scripts/distill/pika_h3_flex_qad/assemble_flex_release.py` | Self-contained NVFP4 model directory | +| `examples/train/configs/knowledge_distillation/pika_h3_flex/` | Stage-1 and stage-2 configs for both adapters | +| `examples/train/calibration/pika_h3_flex/` | Pinned FFN activation `amax` tables | + +## Expected behavior + +Teacher-forced error falls steadily for the first 30-60 steps and flattens by +step 90-120. Video improves faster than audio. Results of the released runs +(held-out rows, lower is better, distances to the bf16 Flex teacher): + +| Run | Step 0 (PTQ) | Selected | Notes | +|---|---|---|---| +| Ref2VA stage 1, 16 GPUs | 0.326 | **0.270** at step 120 (-17.1%) | video x0 0.185, audio x0 0.356 | +| FL2VA stage 1, 24 GPUs | 0.272 | **0.234** at step 60 (-14.1%) | first-frame 0.072, last-frame 0.247 | + +- FL2VA's last-frame match can drift above its tolerance while the score keeps + improving; at 5e-6 this stopped stage 2 at step 60 with no eligible + checkpoint, so its stage-2 config uses 3e-6. +- A lower teacher-forced score with a failed guard is not a better model: + promote only eligible checkpoints, and review the full-eval samples + (`eval_save_dir`) before release. diff --git a/examples/inference/basic/basic_pika_h3_flex.py b/examples/inference/basic/basic_pika_h3_flex.py index c3761c75d..337001334 100644 --- a/examples/inference/basic/basic_pika_h3_flex.py +++ b/examples/inference/basic/basic_pika_h3_flex.py @@ -14,7 +14,14 @@ does not expose a step-count flag. ``--lossless-accel`` turns on the existing training-free FastVideo switches (bit-identical kernels, reference memo, parallel VAE tiles) before the generator starts. ``--nvfp4-decoder`` and ``--nvfp4-text-encoder`` opt into the existing Blackwell FP4 weight paths; -they change weights, not the six-forward sampler. Kernel/VSA work is separate. +``--nvfp4-dit h3_dit`` loads a packed NVFP4 DiT export (e.g. +FastVideo/Pika-H3-Flex-NVFP4-QAD). They change weights, not the six-forward sampler. Kernel/VSA work is separate. + +``--nvfp4-attention`` turns on FlashAttention-4 NVFP4 attention (Q/K in NVFP4, P/V in BF16; Blackwell and FA4 only) +with the named exact Q/K preconditioning, which must match the ``qk_transform`` the checkpoint was QAD-trained with +(``smooth_k+hadamard`` for Pika-H3-Flex-NVFP4-QAD). ``--nvfp4-attention-bf16-blocks 45,49`` keeps those +``transformer_blocks`` on BF16 attention, as in training. ``none`` quantizes Q/K without preconditioning. +Without ``--nvfp4-attention`` attention stays BF16. ``--resolution`` (480p/768p), ``--aspect`` (one of Pika's eleven buckets), and ``--duration`` (5-15 s) pick the bucket and the ``17n + 5`` frame count; @@ -47,7 +54,7 @@ import os from pathlib import Path from typing import Any -from fastvideo import VideoGenerator +from fastvideo import VideoGenerator, envs from fastvideo.api import ( CompileConfig, ComponentConfig, @@ -59,8 +66,10 @@ from fastvideo.api import ( OutputConfig, ParallelismConfig, PipelineSelection, + QuantizationConfig, SamplingConfig, ) +from fastvideo.attention.utils.nvfp4_qk_precondition import QK_TRANSFORMS, parse_block_indices from fastvideo.models.vision_utils import load_image from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference from fastvideo.pipelines.basic.pika_h3_flex.buckets import ( @@ -130,8 +139,23 @@ def parse_args() -> argparse.Namespace: help="Blackwell-only NVFP4 H3 decoder (quality-36L or fast-8L decoder-nvfp4.safetensors)") parser.add_argument("--nvfp4-text-encoder", help="Blackwell-only serialized NVFP4 Qwen3-VL text_encoder directory") + parser.add_argument("--nvfp4-dit", + choices=("h3_dit", "h3_dit_ffn"), + help="load a packed NVFP4 DiT export with this layer profile (h3_dit: attention + FFN linears)") + parser.add_argument("--nvfp4-attention", + choices=tuple(QK_TRANSFORMS), + help="FA4 NVFP4 Q/K attention with this exact Q/K preconditioning (default: BF16 attention)") + parser.add_argument("--nvfp4-attention-bf16-blocks", + default="", + help="comma-separated transformer_blocks indices that keep BF16 attention, e.g. 45,49") parser.add_argument("--output", default="outputs/pika_h3_flex") args = parser.parse_args() + if args.nvfp4_attention_bf16_blocks and not args.nvfp4_attention: + parser.error("--nvfp4-attention-bf16-blocks needs --nvfp4-attention") + try: + parse_block_indices(args.nvfp4_attention_bf16_blocks) + except ValueError as error: + parser.error(f"--nvfp4-attention-bf16-blocks: {error}") if args.task in ("i2va", "fl2va") and not args.image: parser.error(f"--task {args.task} needs --image") @@ -202,8 +226,19 @@ def _pipeline(args: argparse.Namespace) -> PipelineSelection: ) +def _apply_nvfp4_attention(args: argparse.Namespace) -> None: + """Workers inherit the environment, so set it before VideoGenerator starts.""" + if not args.nvfp4_attention: + return + envs.FASTVIDEO_NVFP4_FA4.set(True) + envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.set(args.nvfp4_attention) + envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.set(args.nvfp4_attention_bf16_blocks) + envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1") + + def main() -> None: args = parse_args() + _apply_nvfp4_attention(args) if args.lossless_accel: for key, value in LOSSLESS_ACCEL_ENV.items(): os.environ.setdefault(key, value) @@ -223,6 +258,8 @@ def main() -> None: parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), offload=offload, compile=CompileConfig(enabled=args.compile), + quantization=(QuantizationConfig(transformer_quant="NVFP4", layer_profile=args.nvfp4_dit) + if args.nvfp4_dit else None), ), pipeline=_pipeline(args), )) diff --git a/examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json b/examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json new file mode 100644 index 000000000..3e522b84d --- /dev/null +++ b/examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json @@ -0,0 +1,102 @@ +{ + "b0.ff.fc_in": 5.625, + "b0.ff.fc_out": 2008.0, + "b1.ff.fc_in": 4.5, + "b1.ff.fc_out": 2432.0, + "b2.ff.fc_in": 4.25, + "b2.ff.fc_out": 1936.0, + "b3.ff.fc_in": 4.84375, + "b3.ff.fc_out": 1432.0, + "b4.ff.fc_in": 6.5625, + "b4.ff.fc_out": 2912.0, + "b5.ff.fc_in": 17.549999999999955, + "b5.ff.fc_out": 4592.0, + "b6.ff.fc_in": 26.0, + "b6.ff.fc_out": 5008.0, + "b7.ff.fc_in": 19.862499999999997, + "b7.ff.fc_out": 4224.0, + "b8.ff.fc_in": 12.424999999999999, + "b8.ff.fc_out": 2208.0, + "b9.ff.fc_in": 12.53125, + "b9.ff.fc_out": 2576.0, + "b10.ff.fc_in": 10.1875, + "b10.ff.fc_out": 4832.0, + "b11.ff.fc_in": 17.625, + "b11.ff.fc_out": 6406.4, + "b12.ff.fc_in": 18.125, + "b12.ff.fc_out": 1464.0, + "b13.ff.fc_in": 6.9125, + "b13.ff.fc_out": 20736.0, + "b14.ff.fc_in": 7.46875, + "b14.ff.fc_out": 1444.8, + "b15.ff.fc_in": 9.1875, + "b15.ff.fc_out": 1688.7999999999993, + "b16.ff.fc_in": 10.381249999999994, + "b16.ff.fc_out": 3216.0, + "b17.ff.fc_in": 25.125, + "b17.ff.fc_out": 2480.0, + "b18.ff.fc_in": 10.971874999999997, + "b18.ff.fc_out": 12736.0, + "b19.ff.fc_in": 9.6875, + "b19.ff.fc_out": 3852.7999999999997, + "b20.ff.fc_in": 8.35625, + "b20.ff.fc_out": 2755.2, + "b21.ff.fc_in": 11.0, + "b21.ff.fc_out": 8630.39999999998, + "b22.ff.fc_in": 13.9375, + "b22.ff.fc_out": 7188.799999999996, + "b23.ff.fc_in": 12.95, + "b23.ff.fc_out": 6208.0, + "b24.ff.fc_in": 12.6875, + "b24.ff.fc_out": 7572.799999999996, + "b25.ff.fc_in": 9.31875, + "b25.ff.fc_out": 10304.0, + "b26.ff.fc_in": 14.25, + "b26.ff.fc_out": 5632.0, + "b27.ff.fc_in": 11.125, + "b27.ff.fc_out": 9705.599999999991, + "b28.ff.fc_in": 10.671874999999986, + "b28.ff.fc_out": 9542.4, + "b29.ff.fc_in": 9.21875, + "b29.ff.fc_out": 12550.399999999994, + "b30.ff.fc_in": 10.0, + "b30.ff.fc_out": 18176.0, + "b31.ff.fc_in": 9.846874999999997, + "b31.ff.fc_out": 13440.0, + "b32.ff.fc_in": 11.5, + "b32.ff.fc_out": 8704.0, + "b33.ff.fc_in": 10.125, + "b33.ff.fc_out": 12185.599999999999, + "b34.ff.fc_in": 9.42812499999998, + "b34.ff.fc_out": 8512.0, + "b35.ff.fc_in": 7.375, + "b35.ff.fc_out": 6720.0, + "b36.ff.fc_in": 17.875, + "b36.ff.fc_out": 30208.0, + "b37.ff.fc_in": 10.53125, + "b37.ff.fc_out": 13056.0, + "b38.ff.fc_in": 13.9125, + "b38.ff.fc_out": 29535.99999999997, + "b39.ff.fc_in": 7.7437499999999995, + "b39.ff.fc_out": 302182.3999999999, + "b40.ff.fc_in": 10.3125, + "b40.ff.fc_out": 11712.0, + "b41.ff.fc_in": 13.1875, + "b41.ff.fc_out": 10368.0, + "b42.ff.fc_in": 16.25, + "b42.ff.fc_out": 15641.599999999977, + "b43.ff.fc_in": 17.625, + "b43.ff.fc_out": 28544.0, + "b44.ff.fc_in": 36.25, + "b44.ff.fc_out": 113254.4, + "b45.ff.fc_in": 44.52499999999998, + "b45.ff.fc_out": 156262.4, + "b46.ff.fc_in": 57.25, + "b46.ff.fc_out": 14272.0, + "b47.ff.fc_in": 56.5, + "b47.ff.fc_out": 30400.0, + "b48.ff.fc_in": 219.0999999999999, + "b48.ff.fc_out": 75264.0, + "b49.ff.fc_in": 111.0, + "b49.ff.fc_out": 166912.0 +} \ No newline at end of file diff --git a/examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json b/examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json new file mode 100644 index 000000000..449f7c8b2 --- /dev/null +++ b/examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json @@ -0,0 +1,102 @@ +{ + "b0.ff.fc_in": 5.75, + "b0.ff.fc_out": 1792.0, + "b1.ff.fc_in": 8.9375, + "b1.ff.fc_out": 2704.0, + "b2.ff.fc_in": 4.5625, + "b2.ff.fc_out": 2096.0, + "b3.ff.fc_in": 6.875, + "b3.ff.fc_out": 2384.0, + "b4.ff.fc_in": 11.125, + "b4.ff.fc_out": 2720.0, + "b5.ff.fc_in": 18.549999999999997, + "b5.ff.fc_out": 8225.280000000006, + "b6.ff.fc_in": 33.0, + "b6.ff.fc_out": 4352.0, + "b7.ff.fc_in": 18.75, + "b7.ff.fc_out": 4608.0, + "b8.ff.fc_in": 13.375, + "b8.ff.fc_out": 2320.0, + "b9.ff.fc_in": 13.212499999999999, + "b9.ff.fc_out": 2890.880000000001, + "b10.ff.fc_in": 23.875, + "b10.ff.fc_out": 4320.0, + "b11.ff.fc_in": 13.737499999999999, + "b11.ff.fc_out": 6227.2, + "b12.ff.fc_in": 14.874999999999998, + "b12.ff.fc_out": 1632.0, + "b13.ff.fc_in": 15.049999999999999, + "b13.ff.fc_out": 22144.0, + "b14.ff.fc_in": 7.46875, + "b14.ff.fc_out": 1552.0, + "b15.ff.fc_in": 9.4375, + "b15.ff.fc_out": 1444.8, + "b16.ff.fc_in": 11.272500000000008, + "b16.ff.fc_out": 3216.0, + "b17.ff.fc_in": 30.5, + "b17.ff.fc_out": 3760.0, + "b18.ff.fc_in": 11.4375, + "b18.ff.fc_out": 14656.0, + "b19.ff.fc_in": 12.5, + "b19.ff.fc_out": 4144.0, + "b20.ff.fc_in": 15.5, + "b20.ff.fc_out": 4544.0, + "b21.ff.fc_in": 10.149999999999999, + "b21.ff.fc_out": 9920.0, + "b22.ff.fc_in": 14.524999999999999, + "b22.ff.fc_out": 6656.0, + "b23.ff.fc_in": 18.75, + "b23.ff.fc_out": 8288.0, + "b24.ff.fc_in": 12.987500000000004, + "b24.ff.fc_out": 11520.0, + "b25.ff.fc_in": 14.625, + "b25.ff.fc_out": 9963.520000000004, + "b26.ff.fc_in": 10.0625, + "b26.ff.fc_out": 9390.080000000016, + "b27.ff.fc_in": 11.5625, + "b27.ff.fc_out": 6784.0, + "b28.ff.fc_in": 10.587499999999999, + "b28.ff.fc_out": 9900.8, + "b29.ff.fc_in": 8.44375, + "b29.ff.fc_out": 11991.040000000008, + "b30.ff.fc_in": 10.25, + "b30.ff.fc_out": 15411.199999999999, + "b31.ff.fc_in": 11.355000000000004, + "b31.ff.fc_out": 8576.0, + "b32.ff.fc_in": 12.9375, + "b32.ff.fc_out": 12928.0, + "b33.ff.fc_in": 11.125, + "b33.ff.fc_out": 9907.200000000004, + "b34.ff.fc_in": 10.230000000000004, + "b34.ff.fc_out": 14656.0, + "b35.ff.fc_in": 13.75, + "b35.ff.fc_out": 15360.0, + "b36.ff.fc_in": 14.417500000000004, + "b36.ff.fc_out": 29568.0, + "b37.ff.fc_in": 11.4375, + "b37.ff.fc_out": 12759.040000000008, + "b38.ff.fc_in": 19.085000000000008, + "b38.ff.fc_out": 30848.0, + "b39.ff.fc_in": 8.44375, + "b39.ff.fc_out": 290816.0, + "b40.ff.fc_in": 12.272500000000008, + "b40.ff.fc_out": 20787.199999999997, + "b41.ff.fc_in": 10.8125, + "b41.ff.fc_out": 18887.68000000014, + "b42.ff.fc_in": 11.06875, + "b42.ff.fc_out": 11264.0, + "b43.ff.fc_in": 16.275, + "b43.ff.fc_out": 25728.0, + "b44.ff.fc_in": 36.75, + "b44.ff.fc_out": 54476.799999999996, + "b45.ff.fc_in": 45.51000000000005, + "b45.ff.fc_out": 118272.0, + "b46.ff.fc_in": 40.20000000000002, + "b46.ff.fc_out": 20070.399999999998, + "b47.ff.fc_in": 54.0, + "b47.ff.fc_out": 31232.0, + "b48.ff.fc_in": 113.5, + "b48.ff.fc_out": 40960.0, + "b49.ff.fc_in": 128.0, + "b49.ff.fc_out": 167731.19999999998 +} \ No newline at end of file diff --git a/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_fp4attn_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_fp4attn_qad.yaml new file mode 100644 index 000000000..7da70d18b --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_fp4attn_qad.yaml @@ -0,0 +1,132 @@ +# Pika H3 Flex FL2VA stage 2: adapt the stage-1 NVFP4 QAD student to dense FP4 attention. +# +# Same recipe, data and calibration as fl2va_nvfp4_qad.yaml, starting from that run's best +# checkpoint (method.init_student_dcp). Attention runs the FlashAttention-4 +# NVFP4 numerics (Q/K in NVFP4 per 16 channels, P and V in BF16) through the +# dense FP4 emulator with straight-through gradients. Q/K are preconditioned +# exactly before quantization (token-mean subtraction on K, then a shared +# 128-point Hadamard rotation of Q and K), and blocks 45 and 49, the two most +# FP4-sensitive attention layers, keep BF16 attention. Deploy with the same +# settings: FASTVIDEO_NVFP4_FA4_QK_TRANSFORM=smooth_k+hadamard and +# FASTVIDEO_NVFP4_FA4_BF16_BLOCKS=45,49 (or the example's --nvfp4-attention +# flags). +# +# Data and model paths are placeholders. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json + amax_sha256: b690d1176702b399975d4e0e838784edc89584fda520abb94911e53f553e0daa + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: true + qk_transform: smooth_k+hadamard + bf16_blocks: [45, 49] + data: + manifests: + - data/flex-fl2va/MANIFEST.jsonl + cases: + - t2va + - i2va + - l2va + - fl2va + per_group: + 480p: 160 + 768p: 40 + seed: 20261007 + max_frames: 345 + # FL2VA holds out rows from its own groups (eval_per_group top-up); finalize writes an empty + # manifest here, which the trainer still expects. + eval_manifest: data/flex-fl2va/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + # Stage-1 best checkpoint (its tagged student weights). + init_student_dcp: outputs/pika_flex_fl2va_nvfp4_qad/best + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + # Ignored when init_student_dcp is set (the student no longer starts from the teacher's weights). + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/pika_flex_fl2va_fp4attn_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_fl2va_fp4attn_qad/step0_cache.json + +training: + distributed: + num_gpus: 12 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 12 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 3.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/pika_flex_fl2va_fp4attn_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-fl2va-fp4attn-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/pika_h3_flex/fl2va_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml new file mode 100644 index 000000000..7afa9ec4e --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml @@ -0,0 +1,128 @@ +# Pika H3 Flex FL2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder). +# +# Teacher: Pika H3 Flex FL2VA in bf16, i.e. MiniMax-H3 `transformer` with the +# released `fl2va` adapter merged by +# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same +# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks, +# static FFN scales from the Flex FL2VA calibration table, BF16 dense attention +# (fl2va_fp4attn_qad.yaml adapts the result to dense FP4 attention). +# +# Rows are T2VA / I2VA / L2VA / FL2VA keyframe-anchor rows built from encoded +# Ref2VA rows by scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py +# --mode fl2va (exact first/last frames, re-encoded at the canvas size). +# +# Data and model paths are placeholders: point them at the merged model, the +# row manifests and an output dir. The calibration table is vendored and pinned. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-fl2va-balanced.json + amax_sha256: b690d1176702b399975d4e0e838784edc89584fda520abb94911e53f553e0daa + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: false + data: + manifests: + - data/flex-fl2va/MANIFEST.jsonl + cases: + - t2va + - i2va + - l2va + - fl2va + per_group: + 480p: 160 + 768p: 40 + seed: 20261007 + max_frames: 345 + # FL2VA holds out rows from its own groups (eval_per_group top-up); finalize writes an empty + # manifest here, which the trainer still expects. + eval_manifest: data/flex-fl2va/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + 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 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + 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/pika_flex_fl2va_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_fl2va_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/pika_flex_fl2va_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-fl2va-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/pika_h3_flex/ref2va_fp4attn_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_fp4attn_qad.yaml new file mode 100644 index 000000000..30ea1b170 --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_fp4attn_qad.yaml @@ -0,0 +1,139 @@ +# Pika H3 Flex Ref2VA stage 2: adapt the stage-1 NVFP4 QAD student to dense FP4 attention. +# +# Same recipe, data and calibration as ref2va_nvfp4_qad.yaml, starting from that run's best +# checkpoint (method.init_student_dcp). Attention runs the FlashAttention-4 +# NVFP4 numerics (Q/K in NVFP4 per 16 channels, P and V in BF16) through the +# dense FP4 emulator with straight-through gradients. Q/K are preconditioned +# exactly before quantization (token-mean subtraction on K, then a shared +# 128-point Hadamard rotation of Q and K), and blocks 45 and 49, the two most +# FP4-sensitive attention layers, keep BF16 attention. Deploy with the same +# settings: FASTVIDEO_NVFP4_FA4_QK_TRANSFORM=smooth_k+hadamard and +# FASTVIDEO_NVFP4_FA4_BF16_BLOCKS=45,49 (or the example's --nvfp4-attention +# flags). +# +# Data and model paths are placeholders. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json + amax_sha256: ad60b45fea00906c55e370a0dec82638eb413d6b68537d1ceb38fc52da90b373 + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: true + qk_transform: smooth_k+hadamard + bf16_blocks: [45, 49] + data: + manifests: + # Flex-sized rows (build_flex_ref2va_rows.py) first: they replace source rows with the same id + # and resolution. + - data/flex-ref2va/MANIFEST.jsonl + - data/omniref/MANIFEST.jsonl + cases: + - first_frame + - first_last_frame + - storyboard + - continue_scene + - continue_shot + - audio_reference + - mixed_image_video + - many_images + - multi_video + per_group: + 480p: 160 + 768p: 40 + seed: 20261007 + max_frames: 345 + # Held-out rows: the source held-out manifest with Flex-sized rows (written by finalize). + eval_manifest: data/flex-ref2va/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + # Stage-1 best checkpoint (its tagged student weights). + init_student_dcp: outputs/pika_flex_ref2va_nvfp4_qad/best + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + # Ignored when init_student_dcp is set (the student no longer starts from the teacher's weights). + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/pika_flex_ref2va_fp4attn_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_ref2va_fp4attn_qad/step0_cache.json + +training: + distributed: + num_gpus: 12 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 12 + 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/pika_flex_ref2va_fp4attn_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-ref2va-fp4attn-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/pika_h3_flex/ref2va_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml new file mode 100644 index 000000000..67d29fb64 --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml @@ -0,0 +1,138 @@ +# Pika H3 Flex Ref2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder). +# +# Teacher: Pika H3 Flex Ref2VA in bf16, i.e. MiniMax-H3 `transformer_ref` with the +# released `ref2va` adapter merged by +# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same +# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks +# active from step 0, static FFN activation scales from the Flex calibration +# table. Attention stays BF16 dense in this stage (Flex has no VSA gate to +# freeze); ref2va_fp4attn_qad.yaml adapts the result to dense FP4 attention. +# Norms, AdaLN, the token refiner and the FP32 boundary modules stay frozen. +# +# The objective, evaluation, guards and checkpoint contract are the shared +# recipe's (docs/training/minimax_h3_fp4_qad.md). Rows are OmniRef Ref2VA +# parquet rows; the teacher computes every target on the fly, so the OmniRef +# held-out eval manifest is reused as is. +# +# Data and model paths are placeholders: point them at the merged model, the +# row manifests and an output dir. The calibration table is vendored and pinned. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json + amax_sha256: ad60b45fea00906c55e370a0dec82638eb413d6b68537d1ceb38fc52da90b373 + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: false + data: + manifests: + # Flex-sized rows (build_flex_ref2va_rows.py) first: they replace source rows with the same id + # and resolution. + - data/flex-ref2va/MANIFEST.jsonl + - data/omniref/MANIFEST.jsonl + cases: + - first_frame + - first_last_frame + - storyboard + - continue_scene + - continue_shot + - audio_reference + - mixed_image_video + - many_images + - multi_video + per_group: + 480p: 160 + 768p: 40 + seed: 20261007 + max_frames: 345 + # Held-out rows: the source held-out manifest with Flex-sized rows (written by finalize). + eval_manifest: data/flex-ref2va/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + 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 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + 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/pika_flex_ref2va_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_ref2va_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/pika_flex_ref2va_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-ref2va-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/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index c130b0dfd..b2ad1de5c 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -28,6 +28,7 @@ from fastvideo.attention.backends.abstract import ( PackedSequenceMetadata, ) from fastvideo.attention.backends.attn_qat_infer import _resolve_fa4_pv_mode +from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -284,6 +285,12 @@ class FlashAttentionImpl(AttentionImpl): if nvfp4_fa4 is None: nvfp4_fa4 = envs.FASTVIDEO_NVFP4_FA4.get() self.nvfp4_fa4 = bool(nvfp4_fa4) + # Exact Q/K preconditioning ahead of NVFP4 quantization, shared with the QAD emulator + # (fastvideo.attention.utils.nvfp4_qk_precondition). Same arg-over-env precedence as nvfp4_fa4. + qk_transform = extra_impl_args.get("nvfp4_qk_transform") + if qk_transform is None: + qk_transform = envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() + self.nvfp4_smooth_k, self.nvfp4_hadamard = parse_qk_transform(qk_transform) # PV-mode knob for the FA4-FP4 path. Resolved + validated # unconditionally (shared helper) so a typo fails at construction, not # at the first forward on a Blackwell box, and so the ATTN_QAT_INFER @@ -301,7 +308,8 @@ class FlashAttentionImpl(AttentionImpl): assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}") assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). " "Install via instructions in docs/inference/optimizations.md") - logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s)", self.fa4_pv_mode) + logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s, smooth_k=%s, hadamard=%s)", + self.fa4_pv_mode, self.nvfp4_smooth_k, self.nvfp4_hadamard) def forward( self, @@ -536,6 +544,8 @@ class FlashAttentionImpl(AttentionImpl): orig_seqlen_q = query.shape[1] orig_seqlen_k = key.shape[1] + query, key = precondition_qk(query, key, smooth_k=self.nvfp4_smooth_k, hadamard=self.nvfp4_hadamard) + # Quantize Q/K to FP4 (internally pads to multiple of 128 for SF layout) q_fp4, q_sf = _nvfp4_quantize_for_fa4(query) k_fp4, k_sf = _nvfp4_quantize_for_fa4(key) diff --git a/fastvideo/attention/backends/fp4_dense_qat.py b/fastvideo/attention/backends/fp4_dense_qat.py new file mode 100644 index 000000000..2d242ab90 --- /dev/null +++ b/fastvideo/attention/backends/fp4_dense_qat.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dense FP4 attention emulator for QAD of models deployed with FlashAttention-4 NVFP4 Q/K (dense attention). + +The deployed dense FP4 path (``FLASH_ATTN`` with ``nvfp4_fa4``) quantizes Q and K to NVFP4 per 16 channels and keeps +P and V in BF16. The FP4 block-sparse training kernel (``fp4_vsa_attn_qat`` with ``quantize_pv=False``) computes the +same numerics with straight-through gradients; given every key block for every query block it is dense attention. + +``install_fp4_dense_attention(transformer, numerics)`` routes every dense attention impl of ``transformer`` +(``FlashAttentionImpl``, or ``SDPAImpl`` where the ``FLASH_ATTN`` selection fell back to SDPA because flash-attn is +not installed) through that kernel (``None`` restores the original forward). Inputs arrive as ``[B, L, H, D]`` after +the sequence-parallel all-to-all (full sequence, local heads); they are padded to 128-token tiles, attended over the +full block lists and cut back. +""" +from __future__ import annotations + +import re +import types +from typing import Any + +import torch + +from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, load_qat_vsa_modules +from fastvideo.attention.utils.nvfp4_qk_precondition import precondition_qk + +TILE = 128 +_ORIGINAL = "_fp4_dense_original_forward" +_BLOCK = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.") + + +def dense_block_lists(seq_len: int, batch: int, heads: int, reference: Any, + device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Kernel block lists for full (dense) attention over ``seq_len`` tokens in 128-token tiles.""" + n_tiles = -(-seq_len // TILE) + tile_valid = torch.full((n_tiles, ), TILE, dtype=torch.int32, device=device) + tile_valid[-1] = seq_len - (n_tiles - 1) * TILE + tile_mask = torch.ones((batch, heads, n_tiles, n_tiles), dtype=torch.bool, device=device) + return reference.tile128_mask_to_fp4_blocks(tile_mask, tile_valid) + + +class FP4DenseAttention: + """Dense attention ``forward`` replacement: ``(q, k, v, metadata) -> [B, L, H, D]`` through the FP4 kernel.""" + + def __init__(self, + numerics: FP4AttentionNumerics, + modules: tuple[Any, Any] | None = None, + qk_transform: tuple[bool, bool] = (False, False)) -> None: + self.numerics = numerics + # (smooth K, Hadamard-rotate Q/K) applied before quantization, exactly as the deployed FA4 NVFP4 path does. + self.smooth_k_mean, self.hadamard = qk_transform + self._reference, self._train = modules or load_qat_vsa_modules() + self._lists: dict[tuple[int, int, int, str], tuple[torch.Tensor, ...]] = {} + + def _block_lists(self, seq_len: int, batch: int, heads: int, device: torch.device) -> tuple[torch.Tensor, ...]: + key = (seq_len, batch, heads, str(device)) + if key not in self._lists: + self._lists = {key: dense_block_lists(seq_len, batch, heads, self._reference, device)} + return self._lists[key] + + def __call__(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, metadata: Any) -> torch.Tensor: + if getattr(metadata, "packed_sequence", None) is not None: + raise NotImplementedError("dense FP4 attention emulator does not support packed sequences") + batch, seq_len, heads, _ = query.shape + query, key = precondition_qk(query, key, smooth_k=self.smooth_k_mean, hadamard=self.hadamard) + pad = -seq_len % TILE + q, k, v = (torch.nn.functional.pad(t.to(torch.bfloat16).transpose(1, 2), (0, 0, 0, pad)) + for t in (query, key, value)) + q2k_idx, q2k_num, kv_valid = self._block_lists(seq_len, batch, heads, query.device) + out = self._train.fp4_vsa_attn_qat(q, + k, + v, + q2k_idx, + q2k_num, + kv_valid, + quantize=self.numerics.quantize, + quantize_pv=self.numerics.quantize_pv, + two_level_p=self.numerics.two_level_p, + smooth_k=self.numerics.smooth_k, + first_block_max_floor=self.numerics.first_block_max_floor) + return out[:, :, :seq_len].transpose(1, 2).to(query.dtype) + + +def _dense_impl_types() -> tuple[type, ...]: + from fastvideo.attention.backends.sdpa import SDPAImpl + + try: + from fastvideo.attention.backends.flash_attn import FlashAttentionImpl + except (ImportError, RuntimeError): # no flash-attn install: FLASH_ATTN selection resolves to SDPA + return (SDPAImpl, ) + return (FlashAttentionImpl, SDPAImpl) + + +def named_dense_attention_impls(transformer: torch.nn.Module) -> list[tuple[str, Any]]: + """``(module name, impl)`` for every dense (FlashAttention or SDPA) impl behind a distributed attention layer.""" + types_ = _dense_impl_types() + impls = [] + for name, module in transformer.named_modules(): + impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None) + if isinstance(impl, types_): + impls.append((name, impl)) + return impls + + +def dense_attention_impls(transformer: torch.nn.Module) -> list[Any]: + """Every dense (FlashAttention or SDPA) impl reached through the transformer's distributed attention layers.""" + return [impl for _, impl in named_dense_attention_impls(transformer)] + + +def block_index(module_name: str) -> int | None: + """``transformer_blocks.`` index of a module name (``None`` outside the block stack).""" + match = _BLOCK.search(module_name + ".") + return int(match.group(1)) if match else None + + +def install_fp4_dense_attention( + transformer: torch.nn.Module, + numerics: FP4AttentionNumerics | None, + modules: tuple[Any, Any] | None = None, + *, + qk_transform: tuple[bool, bool] = (False, False), + bf16_blocks: tuple[int, ...] = ()) -> int: + """Route dense attention layers through the FP4 emulator; returns how many now run FP4. + + ``numerics=None`` restores every original forward. ``bf16_blocks`` keeps the attention of those + ``transformer_blocks`` indices on its original (BF16) forward, matching a deployment that opts them out of NVFP4. + """ + override = FP4DenseAttention(numerics, modules, qk_transform) if numerics is not None else None + named = named_dense_attention_impls(transformer) + skip = frozenset(int(b) for b in bf16_blocks) + present = {index for name, _ in named if (index := block_index(name)) is not None} + missing = skip - present + if missing: + raise ValueError(f"bf16_blocks {sorted(missing)} match no transformer_blocks attention layer") + installed = 0 + for name, impl in named: + if _ORIGINAL not in impl.__dict__: + impl.__dict__[_ORIGINAL] = impl.forward + if override is None or block_index(name) in skip: + impl.forward = impl.__dict__[_ORIGINAL] + else: + impl.forward = types.MethodType(lambda self, q, k, v, meta, _o=override: _o(q, k, v, meta), impl) + installed += 1 + return installed if override is not None else len(named) + + +__all__ = [ + "FP4DenseAttention", "block_index", "dense_block_lists", "dense_attention_impls", "install_fp4_dense_attention", + "named_dense_attention_impls" +] diff --git a/fastvideo/attention/utils/nvfp4_qk_precondition.py b/fastvideo/attention/utils/nvfp4_qk_precondition.py new file mode 100644 index 000000000..e130b266e --- /dev/null +++ b/fastvideo/attention/utils/nvfp4_qk_precondition.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact Q/K preconditioning ahead of NVFP4 (per-16 block) Q/K quantization. + +Both transforms leave the attention output unchanged in exact arithmetic and only redistribute what the FP4 blocks +have to represent: + +* ``smooth_k``: subtract each head's token-mean of K. ``q . k_mean`` is the same for every key of a query row, and + softmax is shift-invariant per row. +* ``hadamard``: rotate Q and K by the same orthonormal head-dim Hadamard matrix ``H`` (``Q H (K H)^T = Q K^T``), + which spreads outlier channels across the 16-wide scale blocks. + +The QAD emulator and the deployed FlashAttention-4 NVFP4 path apply the same function, so training and inference +quantize identical tensors. +""" +from __future__ import annotations + +import math +from functools import lru_cache + +import torch + +QK_TRANSFORMS = { + "none": (False, False), + "smooth_k": (True, False), + "hadamard": (False, True), + "smooth_k+hadamard": (True, True), +} + + +def parse_qk_transform(name: str | None) -> tuple[bool, bool]: + """``"smooth_k+hadamard"`` -> ``(smooth_k, hadamard)``; ``None``/``"none"`` -> no transform.""" + key = (name or "none").strip().lower() + if key not in QK_TRANSFORMS: + raise ValueError(f"unknown NVFP4 Q/K transform {name!r}; expected one of {', '.join(QK_TRANSFORMS)}") + return QK_TRANSFORMS[key] + + +def parse_block_indices(raw: str | None) -> frozenset[int]: + """``"45, 49"`` -> ``frozenset({45, 49})``; empty/``None`` -> empty set.""" + items = [item.strip() for item in (raw or "").split(",") if item.strip()] + try: + indices = frozenset(int(item) for item in items) + except ValueError as error: + raise ValueError(f"expected a comma-separated list of block indices, got {raw!r}") from error + if any(index < 0 for index in indices): + raise ValueError(f"block indices must be non-negative, got {raw!r}") + return indices + + +@lru_cache(maxsize=8) +def _hadamard(n: int, device: torch.device) -> torch.Tensor: + if n <= 0 or n & (n - 1): + raise ValueError(f"Hadamard rotation needs a power-of-two head dim, got {n}") + h = torch.ones(1, 1) + while h.shape[0] < n: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) + return (h / math.sqrt(n)).to(device) + + +def hadamard_matrix(n: int, device: torch.device | str) -> torch.Tensor: + """Orthonormal Sylvester Hadamard matrix of size ``n`` (fp32).""" + return _hadamard(n, torch.device(device)) + + +def precondition_qk(query: torch.Tensor, key: torch.Tensor, *, smooth_k: bool, + hadamard: bool) -> tuple[torch.Tensor, torch.Tensor]: + """Transform ``[B, L, H, D]`` Q/K (computed in fp32, returned in the input dtype).""" + if not (smooth_k or hadamard): + return query, key + q, k = query.float(), key.float() + if smooth_k: + k = k - k.mean(dim=1, keepdim=True) + if hadamard: + rotation = hadamard_matrix(q.shape[-1], q.device) + q, k = q @ rotation, k @ rotation + return q.to(query.dtype), k.to(key.dtype) + + +__all__ = ["QK_TRANSFORMS", "hadamard_matrix", "parse_block_indices", "parse_qk_transform", "precondition_qk"] diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 45468ee68..ee7e968f2 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -276,6 +276,12 @@ FASTVIDEO_HOST_IP = EnvStr( category="distributed", doc="IP address of this node when the node has several network interfaces. Set it on each node for multi-node " "inference.") +FASTVIDEO_PG_TIMEOUT_MIN = EnvInt( + 0, + category="distributed", + doc="Timeout in minutes for the training process groups' collectives (0 keeps PyTorch's 10-minute default). " + "Raise it for large models whose rank-0 load and scatter, checkpoint saves or long evaluations outlast the " + "default.") FASTVIDEO_LOOPBACK_IP = EnvStr("", category="distributed", doc="Loopback IP address to use instead of the detected one.") @@ -378,6 +384,19 @@ FASTVIDEO_NVFP4_FA4 = EnvBool( category="attention", doc="FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument " "takes precedence.") +FASTVIDEO_NVFP4_FA4_QK_TRANSFORM = EnvChoice( + "none", + choices=("none", "smooth_k", "hadamard", "smooth_k+hadamard"), + category="attention", + doc="Exact Q/K preconditioning applied right before NVFP4 quantization on the FlashAttention-4 NVFP4 path: " + "smooth_k subtracts each head's token-mean of K, hadamard rotates Q and K by the head-dim Hadamard matrix. " + "Must match the qk_transform the model was QAD-trained with. An explicit nvfp4_qk_transform attention " + "implementation argument takes precedence.") +FASTVIDEO_NVFP4_FA4_BF16_BLOCKS = EnvStr( + "", + category="attention", + doc="Comma-separated MiniMax-H3 transformer_blocks indices whose attention stays in BF16 while " + "FASTVIDEO_NVFP4_FA4 quantizes the rest, for example 45,49.") FASTVIDEO_FA4_PV_MODE = EnvChoice( "bf16", choices=("bf16", "fp8"), diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 3bd6fea21..06ad52e68 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -19,6 +19,7 @@ from fastvideo.attention.backends.abstract import AttentionMetadata, layer_idx_f from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata from fastvideo.attention.layer import DistributedAttention_VSA from fastvideo.attention.selector import get_attn_backend +from fastvideo.attention.utils.nvfp4_qk_precondition import parse_block_indices from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather_with_unpad, @@ -226,6 +227,7 @@ class MiniMaxH3Attention(nn.Module): exact_rope: bool = False, mk_qkv: bool = False, vsa_fp4: dict[str, Any] | None = None, + bf16_attention: bool = False, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -274,6 +276,8 @@ class MiniMaxH3Attention(nn.Module): supported_attention_backends=supported_attention_backends) use_vsa = resolved_backend.get_name() == "VIDEO_SPARSE_ATTN_H3" attention_cls = DistributedAttention_VSA if use_vsa else DistributedAttention + # Blocks listed in FASTVIDEO_NVFP4_FA4_BF16_BLOCKS opt out of the process-wide NVFP4 FA4 Q/K quantization. + precision_args: dict[str, Any] = {"nvfp4_fa4": False} if bf16_attention else {} self.distributed_attention = attention_cls( num_heads=num_attention_heads, head_size=attention_head_dim, @@ -282,6 +286,7 @@ class MiniMaxH3Attention(nn.Module): prefix=prefix, fa4_packed_varlen=fa4_packed_varlen, cudnn_sdpa=cudnn_sdpa, + **precision_args, ) # Opt-in inference route: VSA-H3 selection on the block-sparse FP4 # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path. @@ -698,6 +703,7 @@ class MiniMaxH3TransformerBlock(nn.Module): exact_kernels: frozenset[str] = frozenset(), vsa_fp4: dict[str, Any] | None = None, mk_tier: bool = False, + bf16_attention: bool = False, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -715,6 +721,7 @@ class MiniMaxH3TransformerBlock(nn.Module): exact_rope="rope" in exact_kernels, vsa_fp4=vsa_fp4, mk_qkv=mk_tier, + bf16_attention=bf16_attention, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -977,6 +984,11 @@ class MiniMaxH3Transformer3DModel(BaseDiT): config.quant_config, prefix=f"{config.prefix}.token_refiner", ) + bf16_blocks = parse_block_indices(envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.get()) + unknown_blocks = sorted(index for index in bf16_blocks if index >= arch.num_layers) + if unknown_blocks: + raise ValueError(f"FASTVIDEO_NVFP4_FA4_BF16_BLOCKS {unknown_blocks} exceed the {arch.num_layers} " + "transformer_blocks") self.transformer_blocks = nn.ModuleList([ MiniMaxH3TransformerBlock( arch.hidden_size, @@ -998,6 +1010,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): exact_kernels=self.enabled_exact_kernels, vsa_fp4=config.vsa_fp4, mk_tier=self.use_mk_kernels, + bf16_attention=index in bf16_blocks, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/pipelines/basic/pika_h3_flex/adapter.py b/fastvideo/pipelines/basic/pika_h3_flex/adapter.py index 7177c0d23..a318b4899 100644 --- a/fastvideo/pipelines/basic/pika_h3_flex/adapter.py +++ b/fastvideo/pipelines/basic/pika_h3_flex/adapter.py @@ -234,6 +234,33 @@ def load_adapter(source: str | os.PathLike[str], *, task: str, revision: str | N ) +PIKA_MERGED_MARKER = "pika_h3_flex_merged.json" + + +def merged_adapter(model_dir: str | os.PathLike[str], task: str) -> dict[str, Any] | None: + """The marker entry when ``model_dir``'s trunk for ``task`` already contains the Pika adapter, else None. + + Exported checkpoints (e.g. NVFP4 QAD releases) merge the adapter into the trunk before quantizing it, so + applying the released adapter again would adapt the model twice. Such a checkpoint ships + ``pika_h3_flex_merged.json``: ``{"format_version": 1, "merged": {"": {"trunk": ..., "release_id": ..., + "adapter_sha256": ...}}}``. A marker that does not cover ``task`` is an error: that trunk is absent or unadapted. + """ + path = Path(model_dir) / PIKA_MERGED_MARKER + if not path.is_file(): + return None + marker = _load_manifest(path) + if _as_int(marker.get("format_version"), "format_version") != PIKA_ADAPTER_FORMAT_VERSION: + raise PikaAdapterError(f"Unsupported {PIKA_MERGED_MARKER} format {marker.get('format_version')!r}.") + merged = marker.get("merged") + if not isinstance(merged, Mapping) or task not in merged: + raise PikaAdapterError(f"{str(path)!r} declares merged Pika adapters for {sorted(merged or {})}, not " + f"{task!r}; this checkpoint does not carry the {task} model.") + entry = merged[task] + if not isinstance(entry, Mapping) or entry.get("trunk") != expected_trunk(task): + raise PikaAdapterError(f"{str(path)!r}: the merged {task} adapter must sit on trunk {expected_trunk(task)!r}.") + return dict(entry) + + def adapter_hf_patterns(task: str | None = None) -> list[str]: """The adapter files the pipeline must include in a partial Hub snapshot.""" tasks = tuple(TASK_TO_TRUNK) if task is None else (task, ) @@ -251,6 +278,7 @@ __all__ = [ "PIKA_ADAPTER_RANK", "PIKA_BASE_REPO", "PIKA_MANIFEST_FILENAME", + "PIKA_MERGED_MARKER", "PIKA_SAMPLING_CONTRACT", "PIKA_TARGET_MODULES", "PikaAdapter", @@ -259,6 +287,7 @@ __all__ = [ "adapter_hf_patterns", "expected_trunk", "load_adapter", + "merged_adapter", "task_for_pipeline", "validate_adapter_manifest", ] diff --git a/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py b/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py index 8bfcba982..5befd6da1 100644 --- a/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py +++ b/fastvideo/pipelines/basic/pika_h3_flex/pipeline.py @@ -14,7 +14,8 @@ from fastvideo.configs.pipelines.pika_h3_flex import ( from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline -from fastvideo.pipelines.basic.pika_h3_flex.adapter import adapter_hf_patterns, load_adapter, task_for_pipeline +from fastvideo.pipelines.basic.pika_h3_flex.adapter import (PIKA_MERGED_MARKER, adapter_hf_patterns, load_adapter, + merged_adapter, task_for_pipeline) from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage @@ -36,7 +37,7 @@ class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline): # The concrete class picks the adapter, so a T2VA download does not # fetch the 2.7 GB Ref2VA adapter (and vice versa). task = "ref2va" if issubclass(cls, PikaH3FlexRefPipeline) else "fl2va" - return [*patterns, *adapter_hf_patterns(task)] + return [*patterns, *adapter_hf_patterns(task), PIKA_MERGED_MARKER] def _resolve_pika_adapter(self) -> None: """Point ``lora_path`` at the released adapter before LoRA layers attach. @@ -47,6 +48,13 @@ class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline): """ task = task_for_pipeline(self._ref2va) explicit = self.fastvideo_args.lora_path + merged = merged_adapter(self.model_path, task) if explicit is None else None + if merged is not None: + # The trunk already carries the adapter (an exported, e.g. NVFP4 QAD, checkpoint): attach nothing. + logger.info("Pika H3 Flex: the %s trunk of %s already contains adapter %s; not applying it again", task, + self.model_path, merged.get("release_id")) + self._pika_adapter = None + return source = explicit if explicit is not None else self.model_path revision = None if explicit is not None else self.fastvideo_args.revision adapter = load_adapter(source, task=task, revision=revision) diff --git a/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py b/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py new file mode 100644 index 000000000..36982f804 --- /dev/null +++ b/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deployed NVFP4 FA4 Q/K preconditioning and per-block BF16 opt-out. + +The transform must be applied to Q/K right before NVFP4 quantization with the same function the QAD emulator uses, +and MiniMax-H3 blocks listed in ``FASTVIDEO_NVFP4_FA4_BF16_BLOCKS`` must build BF16 attention impls. +""" +from unittest.mock import MagicMock, patch + +import pytest + +torch = pytest.importorskip("torch") + +import fastvideo.envs as envs # noqa: E402 +from fastvideo.attention.utils.nvfp4_qk_precondition import ( # noqa: E402 + parse_block_indices, parse_qk_transform, precondition_qk, +) + +flash_attn_module = pytest.importorskip( + "fastvideo.attention.backends.flash_attn", + reason="no usable flash-attention package installed", + exc_type=ImportError, +) +FlashAttentionImpl = flash_attn_module.FlashAttentionImpl + + +def _impl(**extra) -> "FlashAttentionImpl": + return FlashAttentionImpl(num_heads=2, head_size=64, causal=False, softmax_scale=0.125, num_kv_heads=2, **extra) + + +@pytest.fixture +def fa4_runtime_available(): + with ( + patch.object(flash_attn_module, "_FA4_FP4_AVAILABLE", True), + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + ): + yield + + +def test_block_list_parsing() -> None: + assert parse_block_indices("") == frozenset() + assert parse_block_indices(None) == frozenset() + assert parse_block_indices(" 45, 49,45 ") == frozenset({45, 49}) + for bad in ("a,1", "-1", "4.5"): + with pytest.raises(ValueError): + parse_block_indices(bad) + + +def test_transform_env_default_and_validation(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override(None)) + assert envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() == "none" + assert _impl().nvfp4_smooth_k is False and _impl().nvfp4_hadamard is False + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("Smooth_K+Hadamard")) + impl = _impl() + assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, True) + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("bogus")) + with pytest.raises(envs.EnvVarError): + envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() + + +def test_impl_arg_wins_over_transform_env(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("hadamard")) + impl = _impl(nvfp4_qk_transform="smooth_k") + assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, False) + with pytest.raises(ValueError): + _impl(nvfp4_qk_transform="bogus") + + +@pytest.mark.parametrize("name", ["none", "smooth_k", "hadamard", "smooth_k+hadamard"]) +def test_transform_applied_before_quantization(name, env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(None)) + impl = _impl(nvfp4_fa4=True, nvfp4_qk_transform=name) + torch.manual_seed(0) + q = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0 + k = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0 + v = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + seen = [] + + def fake_quant(tensor): + seen.append(tensor.clone()) + return tensor, tensor + + fp4 = MagicMock(side_effect=lambda *a, **kw: a[2]) + with ( + patch.object(flash_attn_module, "_nvfp4_quantize_for_fa4", side_effect=fake_quant), + patch.object(flash_attn_module, "flash_attn_fp4_func", fp4), + ): + impl._forward_nvfp4(q, k, v) + + smooth_k, hadamard = parse_qk_transform(name) + expected_q, expected_k = precondition_qk(q, k, smooth_k=smooth_k, hadamard=hadamard) + assert torch.equal(seen[0], expected_q) and torch.equal(seen[1], expected_k) + assert (name == "none") == (torch.equal(seen[0], q) and torch.equal(seen[1], k)) + + +def test_smooth_k_removes_mean_and_preserves_attention() -> None: + torch.manual_seed(1) + q, k, v = (torch.randn(1, 32, 2, 64) for _ in range(3)) + k = k + 5.0 + qt, kt = precondition_qk(q, k, smooth_k=True, hadamard=True) + ref = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (q, k, v))) + out = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (qt, kt, v))) + torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4) + + +def _build_h3_attention(bf16_attention: bool): + h3 = pytest.importorskip("fastvideo.models.dits.minimax_h3") + backend = MagicMock() + backend.get_name.return_value = "FLASH_ATTN" + distributed = MagicMock() + with ( + patch.object(h3, "get_attn_backend", return_value=backend), + patch.object(h3, "get_compute_dtype", return_value=torch.bfloat16), + patch.object(h3, "ReplicatedLinear", side_effect=lambda *a, **kw: torch.nn.Identity()), + patch.object(h3, "DistributedAttention", distributed), + ): + h3.MiniMaxH3Attention(64, 2, 32, 1e-6, (), None, prefix="transformer_blocks.3.attn", + bf16_attention=bf16_attention) + return distributed.call_args.kwargs + + +def test_h3_attention_bf16_opt_out_wiring() -> None: + assert _build_h3_attention(True)["nvfp4_fa4"] is False + assert "nvfp4_fa4" not in _build_h3_attention(False) + + +def test_bf16_block_impl_ignores_env_opt_in(env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(True)) + assert _impl(nvfp4_fa4=_build_h3_attention(True)["nvfp4_fa4"]).nvfp4_fa4 is False + assert _impl(**{k: v for k, v in _build_h3_attention(False).items() if k == "nvfp4_fa4"}).nvfp4_fa4 is True # env wins diff --git a/fastvideo/tests/loader/test_pika_h3_flex_adapter.py b/fastvideo/tests/loader/test_pika_h3_flex_adapter.py index ec077e065..8e741dbbc 100644 --- a/fastvideo/tests/loader/test_pika_h3_flex_adapter.py +++ b/fastvideo/tests/loader/test_pika_h3_flex_adapter.py @@ -93,3 +93,35 @@ def test_load_adapter_requires_the_manifest(tmp_path: Path) -> None: (task_dir / "manifest.json").unlink() with pytest.raises(adapter.PikaAdapterError, match="manifest"): adapter.load_adapter(tmp_path, task="fl2va") + + +def _write_marker(root: Path, merged: dict) -> None: + (root / adapter.PIKA_MERGED_MARKER).write_text(json.dumps({"format_version": 1, "merged": merged}), + encoding="utf-8") + + +def test_merged_marker_skips_the_released_adapter(tmp_path): + assert adapter.merged_adapter(tmp_path, "ref2va") is None + _write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref", "release_id": "pika-flex-ref2va-v1"}}) + assert adapter.merged_adapter(tmp_path, "ref2va")["release_id"] == "pika-flex-ref2va-v1" + + +def test_merged_marker_must_cover_the_task_and_trunk(tmp_path): + _write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref"}}) + with pytest.raises(adapter.PikaAdapterError, match="does not carry the fl2va model"): + adapter.merged_adapter(tmp_path, "fl2va") + _write_marker(tmp_path, {"ref2va": {"trunk": "transformer"}}) + with pytest.raises(adapter.PikaAdapterError, match="must sit on trunk"): + adapter.merged_adapter(tmp_path, "ref2va") + + +def test_assemble_script_marker_name_matches_the_loader(): + import importlib.util + from pathlib import Path + + from fastvideo.pipelines.basic.pika_h3_flex.adapter import PIKA_MERGED_MARKER + script = Path(__file__).resolve().parents[3] / "scripts/distill/pika_h3_flex_qad/assemble_flex_release.py" + spec = importlib.util.spec_from_file_location("assemble_flex_release", script) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + assert module.PIKA_MERGED_MARKER == PIKA_MERGED_MARKER diff --git a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py new file mode 100644 index 000000000..9b3bf928a --- /dev/null +++ b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py @@ -0,0 +1,576 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for the Pika H3 Flex QAD profile: ladder, re-noise determinism, contract, adapter merge.""" +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import scheduler_step +from fastvideo.train.models.minimax_h3.flex import (CONTRACT_FILE, CONTRACT_SCHEMA, FlexContract, + MiniMaxH3FlexRef2VAModel, RowRenoiseScheduler, flex_sigmas, + renoise_seed) + +# Pika's published per-modality ladders (knots 1000, 750, 500, 250, 100, 40, 0 shifted by 12 / 3). +PIKA_VIDEO_SIGMAS = (1.0, 0.97297, 0.92308, 0.8, 0.57143, 0.33333, 0.0) +PIKA_AUDIO_SIGMAS = (1.0, 0.9, 0.75, 0.5, 0.25, 0.11111, 0.0) +STEPS = (1000, 750, 500, 250, 100, 40) + + +def _merge_module(): + path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py" + spec = importlib.util.spec_from_file_location("merge_pika_adapter", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _CPUFlexModel(MiniMaxH3FlexRef2VAModel): + device = torch.device("cpu") + + +def _model_stub() -> MiniMaxH3FlexRef2VAModel: + model = object.__new__(_CPUFlexModel) + model.contract = FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0) + return model + + +def _row(seed: int) -> SimpleNamespace: + return SimpleNamespace(spec=SimpleNamespace(seed=seed)) + + +def test_flex_sigmas_match_pika_ladders(): + torch.testing.assert_close(flex_sigmas(STEPS, 12.0), torch.tensor(PIKA_VIDEO_SIGMAS), atol=1e-5, rtol=0) + torch.testing.assert_close(flex_sigmas(STEPS, 3.0), torch.tensor(PIKA_AUDIO_SIGMAS), atol=1e-5, rtol=0) + + +def test_schedulers_expose_six_forwards_in_clean_time(): + video, audio = _model_stub().schedulers(_row(7)) + assert video.timesteps.numel() == audio.timesteps.numel() == 6 + torch.testing.assert_close(video.timesteps, 1 - torch.tensor(PIKA_VIDEO_SIGMAS[:-1]), atol=1e-5, rtol=0) + assert video.stochastic_sampling and audio.stochastic_sampling + + +def test_schedulers_require_the_row(): + with pytest.raises(ValueError, match="prepared row"): + _model_stub().schedulers() + + +def test_renoise_step_is_x0_plus_fixed_fresh_noise(): + video, _ = _model_stub().schedulers(_row(11)) + state, velocity = torch.randn(5, 4), torch.randn(5, 4) + rung = 2 + out = scheduler_step(video, rung, velocity, state) + timestep = video.timesteps[rung] + x0 = state + (1 - timestep) * velocity + generator = torch.Generator().manual_seed(renoise_seed(11, rung, 0)) + noise = torch.randn(state.shape, generator=generator, dtype=torch.float32) + sigma_next = video.sigmas[rung + 1] + torch.testing.assert_close(out, (1 - sigma_next) * x0 + sigma_next * noise) + + +def test_renoise_noise_depends_only_on_row_forward_and_modality(): + """Teacher-forced pass, rollout and re-evaluation of one row must draw identical noise at each forward.""" + state, velocity = torch.randn(6, 3), torch.randn(6, 3) + first = _model_stub().schedulers(_row(3)) + second = _model_stub().schedulers(_row(3)) + # Out-of-order use (a rollout after a full teacher-forced pass) still reproduces forward 1. + for rung in range(5): + scheduler_step(first[0], rung, velocity, state) + torch.testing.assert_close(scheduler_step(first[0], 1, velocity, state), + scheduler_step(second[0], 1, velocity, state)) + assert not torch.equal(scheduler_step(first[0], 1, velocity, state), scheduler_step(first[0], 2, velocity, + state)) + other_row = _model_stub().schedulers(_row(4)) + assert not torch.equal(scheduler_step(first[0], 1, velocity, state), + scheduler_step(other_row[0], 1, velocity, state)) + assert renoise_seed(3, 1, 0) != renoise_seed(3, 1, 1) + + +def test_last_forward_returns_x0_without_noise(): + video, audio = _model_stub().schedulers(_row(5)) + state, velocity = torch.randn(4, 2), torch.randn(4, 2) + for scheduler in (video, audio): + out = scheduler_step(scheduler, 5, velocity, state) + torch.testing.assert_close(out, state + (1 - scheduler.timesteps[5]) * velocity) + + +def test_row_renoise_scheduler_refuses_an_implicit_index(): + scheduler = RowRenoiseScheduler(12.0, 0, 0) + scheduler.set_timesteps(sigmas=flex_sigmas(STEPS, 12.0)) + with pytest.raises(ValueError, match="explicit step index"): + scheduler.step(torch.zeros(2, 2), scheduler.timesteps[0], torch.zeros(2, 2)) + + +def _write_contract(path: Path, **overrides) -> None: + data = { + "schema": CONTRACT_SCHEMA, + "task": "ref2va", + "trunk": "transformer_ref", + "denoising_steps": list(STEPS), + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "sampler": "renoise", + } + data.update(overrides) + (path / CONTRACT_FILE).write_text(json.dumps(data)) + + +def test_contract_reads_the_released_ladder(tmp_path): + _write_contract(tmp_path) + assert FlexContract.read(str(tmp_path)) == FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0) + + +@pytest.mark.parametrize("overrides", [{ + "denoising_steps": [1000, 500, 0] +}, { + "video_scheduler_shift": 10.0 +}, { + "trunk": "transformer" +}, { + "sampler": "euler" +}]) +def test_contract_rejects_anything_but_the_release(tmp_path, overrides): + _write_contract(tmp_path, **overrides) + with pytest.raises(ValueError): + FlexContract.read(str(tmp_path)) + + +def test_contract_must_match_the_profile_trunk(tmp_path): + _write_contract(tmp_path, task="fl2va", trunk="transformer") + with pytest.raises(ValueError, match="needs ref2va on transformer_ref"): + _model_stub()._read_contract(str(tmp_path)) + + +def _reference_module(): + from fastvideo.attention.backends.fp4_vsa_qat import _load_checkout + return _load_checkout("attn_qat_vsa_reference") + + +class _ExactKernel: + """Stand-in for ``attn_qat_vsa_train``: exact softmax over the listed (here: all) blocks and valid keys.""" + + def __init__(self, reference): + self.reference = reference + self.calls = [] + + def fp4_vsa_attn_qat(self, q, k, v, q2k_idx, q2k_num, kv_valid, **numerics): + self.calls.append((q.shape, numerics)) + n_blocks = q.shape[2] // 128 + assert bool((q2k_num == n_blocks).all()) + keys = self.reference.key_valid_mask(kv_valid, n_blocks, q.device).reshape(-1) + scores = (q.float() @ k.float().transpose(-1, -2)) / q.shape[-1]**0.5 + scores = scores.masked_fill(~keys[None, None, None, :], float("-inf")) + return (scores.softmax(-1) @ v.float()).to(torch.bfloat16) + + +def test_dense_block_lists_cover_every_key_tile(): + from fastvideo.attention.backends.fp4_dense_qat import dense_block_lists + reference = _reference_module() + q2k_idx, q2k_num, kv_valid = dense_block_lists(300, 1, 2, reference, torch.device("cpu")) + assert q2k_num.shape == (1, 2, 3) and bool((q2k_num == 3).all()) + assert sorted(q2k_idx[0, 0, 0].tolist()) == [0, 1, 2] + assert kv_valid.tolist() == [64, 64, 64, 64, 44, 0] + + +def test_fp4_dense_attention_pads_to_tiles_and_matches_dense_attention(): + from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + reference = _reference_module() + kernel = _ExactKernel(reference) + attention = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False), (reference, kernel)) + torch.manual_seed(0) + q, k, v = (torch.randn(1, 200, 2, 16, dtype=torch.bfloat16) for _ in range(3)) + out = attention(q, k, v, None) + expected = torch.nn.functional.scaled_dot_product_attention(*(t.float().transpose(1, 2) for t in (q, k, v))) + assert out.shape == q.shape and out.dtype == q.dtype + torch.testing.assert_close(out.float(), expected.transpose(1, 2), atol=2e-2, rtol=2e-2) + shape, numerics = kernel.calls[0] + assert shape == (1, 2, 256, 16) and numerics["quantize_pv"] is False and numerics["quantize"] is True + + +def test_install_fp4_dense_attention_overrides_and_restores_flash_attention(): + from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + try: + from fastvideo.attention.backends.flash_attn import FlashAttentionImpl + except (ImportError, RuntimeError) as error: # FlashAttention is a CUDA-host dependency + pytest.skip(f"FlashAttention backend not importable here: {error}") + + class Layer(torch.nn.Module): + + def __init__(self): + super().__init__() + self.distributed_attention = SimpleNamespace(attn_impl=object.__new__(FlashAttentionImpl)) + + model = torch.nn.Sequential(Layer(), Layer()) + impls = [layer.distributed_attention.attn_impl for layer in model] + originals = [impl.forward for impl in impls] + reference = _reference_module() + kernel = _ExactKernel(reference) + assert install_fp4_dense_attention(model, FP4AttentionNumerics(quantize_pv=False), (reference, kernel)) == 2 + q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16) + impls[0].forward(q, q, q, None) + assert len(kernel.calls) == 1 + install_fp4_dense_attention(model, None) + assert all(impl.forward == original for impl, original in zip(impls, originals, strict=True)) + + +def test_install_fp4_dense_attention_covers_sdpa_fallback(): + """Without flash-attn the FLASH_ATTN selection resolves to SDPA; the emulator must still be installed.""" + from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + from fastvideo.attention.backends.sdpa import SDPAImpl + + class Layer(torch.nn.Module): + + def __init__(self): + super().__init__() + self.distributed_attention = SimpleNamespace(attn_impl=SDPAImpl(1, 8, causal=False, softmax_scale=8**-0.5)) + + model = torch.nn.Sequential(Layer(), Layer()) + impls = [layer.distributed_attention.attn_impl for layer in model] + originals = [impl.forward for impl in impls] + reference = _reference_module() + kernel = _ExactKernel(reference) + assert install_fp4_dense_attention(model, FP4AttentionNumerics(quantize_pv=False), (reference, kernel)) == 2 + q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16) + assert impls[1].forward(q, q, q, None).shape == q.shape + assert len(kernel.calls) == 1 + with pytest.raises(NotImplementedError, match="packed"): + impls[0].forward(q, q, q, SimpleNamespace(packed_sequence=object())) + install_fp4_dense_attention(model, None) + assert all(impl.forward == original for impl, original in zip(impls, originals, strict=True)) + + +@pytest.mark.parametrize("name", ["smooth_k", "hadamard", "smooth_k+hadamard"]) +def test_precondition_qk_leaves_attention_unchanged(name): + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk + torch.manual_seed(0) + q, k, v = (torch.randn(2, 37, 3, 16, dtype=torch.float32) for _ in range(3)) + k = k + 4.0 # shared K bias that smoothing removes + smooth, hadamard = parse_qk_transform(name) + tq, tk = precondition_qk(q, k, smooth_k=smooth, hadamard=hadamard) + + def attend(a, b): + return torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (a, b, v))).transpose(1, 2) + + torch.testing.assert_close(attend(tq, tk), attend(q, k), atol=1e-5, rtol=1e-5) + if smooth: + assert float(tk.mean(dim=1).abs().max()) < 1e-4 + + +def test_parse_qk_transform_rejects_unknown_names(): + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform + assert parse_qk_transform(None) == (False, False) + with pytest.raises(ValueError, match="unknown NVFP4 Q/K transform"): + parse_qk_transform("rotate") + + +def test_fp4_dense_attention_with_qk_transform_matches_dense_attention(): + from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + reference = _reference_module() + kernel = _ExactKernel(reference) + attention = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False), (reference, kernel), + qk_transform=(True, True)) + torch.manual_seed(0) + q, k, v = (torch.randn(1, 200, 2, 16, dtype=torch.bfloat16) for _ in range(3)) + out = attention(q, k, v, None) + expected = torch.nn.functional.scaled_dot_product_attention(*(t.float().transpose(1, 2) for t in (q, k, v))) + torch.testing.assert_close(out.float(), expected.transpose(1, 2), atol=3e-2, rtol=3e-2) + + +def test_install_fp4_dense_attention_keeps_bf16_blocks_on_original_forward(): + from fastvideo.attention.backends.fp4_dense_qat import block_index, install_fp4_dense_attention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + from fastvideo.attention.backends.sdpa import SDPAImpl + + class Block(torch.nn.Module): + + def __init__(self): + super().__init__() + self.attn = torch.nn.Module() + self.attn.distributed_attention = SimpleNamespace( + attn_impl=SDPAImpl(1, 8, causal=False, softmax_scale=8**-0.5)) + + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList([Block(), Block(), Block()]) + impls = [block.attn.distributed_attention.attn_impl for block in model.transformer_blocks] + originals = [impl.forward for impl in impls] + reference = _reference_module() + kernel = _ExactKernel(reference) + numerics = FP4AttentionNumerics(quantize_pv=False) + assert block_index("transformer_blocks.12._checkpoint_wrapped_module.attn") == 12 + assert block_index("token_refiner.blocks.0.attn") is None + assert install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(2, )) == 2 + assert impls[2].forward == originals[2] and impls[0].forward != originals[0] + q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16) + impls[2].forward(q, q, q, None) + assert kernel.calls == [] + impls[1].forward(q, q, q, None) + assert len(kernel.calls) == 1 + with pytest.raises(ValueError, match="match no transformer_blocks"): + install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(7, )) + + +def _write_trunk(trunk: Path, weights: dict[str, torch.Tensor]) -> None: + trunk.mkdir(parents=True) + names = sorted(weights) + shards = {"model-00001-of-00002.safetensors": names[:2], "model-00002-of-00002.safetensors": names[2:]} + weight_map = {} + for shard, keys in shards.items(): + save_file({key: weights[key] for key in keys}, str(trunk / shard)) + weight_map.update({key: shard for key in keys}) + (trunk / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + (trunk / "config.json").write_text("{}") + + +def test_merge_trunk_adds_scaled_lora_in_fp32(tmp_path): + merge = _merge_module() + torch.manual_seed(0) + weights = { + "blocks.0.attn.to_q.weight": torch.randn(6, 4).to(torch.bfloat16), + "blocks.0.attn.to_q.bias": torch.randn(6).to(torch.bfloat16), + "blocks.0.ff.net.2.weight": torch.randn(4, 8).to(torch.bfloat16), + "proj_out.weight": torch.randn(3, 4), + } + factors = { + "blocks.0.attn.to_q": (torch.randn(2, 4), torch.randn(6, 2)), + "blocks.0.ff.net.2": (torch.randn(2, 8), torch.randn(4, 2)), + } + _write_trunk(tmp_path / "base", weights) + report = merge.merge_trunk(tmp_path / "base", tmp_path / "out", factors, scale=0.5) + assert report["modules"] == 2 + from safetensors.torch import load_file + merged = {} + for shard in (tmp_path / "out").glob("*.safetensors"): + merged.update(load_file(str(shard))) + for module, (lora_a, lora_b) in factors.items(): + key = f"{module}.weight" + expected = (weights[key].float() + 0.5 * lora_b @ lora_a).to(torch.bfloat16) + assert merged[key].dtype == torch.bfloat16 + assert torch.equal(merged[key], expected) + assert torch.equal(merged["blocks.0.attn.to_q.bias"], weights["blocks.0.attn.to_q.bias"]) + assert torch.equal(merged["proj_out.weight"], weights["proj_out.weight"]) + + +def test_merge_trunk_refuses_adapter_modules_without_weights(tmp_path): + merge = _merge_module() + _write_trunk(tmp_path / "base", {f"w{i}.weight": torch.zeros(2, 2) for i in range(3)}) + with pytest.raises(ValueError, match="no weight"): + merge.merge_trunk(tmp_path / "base", tmp_path / "out", {"missing": (torch.zeros(1, 2), torch.zeros(2, 1))}, + scale=1.0) + + +def test_adapter_factors_pair_peft_keys(tmp_path): + merge = _merge_module() + path = tmp_path / "adapter.safetensors" + save_file( + { + "a.to_q.lora_A.default.weight": torch.zeros(2, 4), + "a.to_q.lora_B.default.weight": torch.zeros(6, 2), + }, str(path)) + factors = merge.adapter_factors(str(path)) + assert list(factors) == ["a.to_q"] and factors["a.to_q"][0].shape == (2, 4) + save_file({"a.to_q.lora_A.default.weight": torch.zeros(2, 4)}, str(path)) + with pytest.raises(ValueError, match="both LoRA factors"): + merge.adapter_factors(str(path)) + + +def _calibrate_module(): + path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py" + spec = importlib.util.spec_from_file_location("calibrate_flex_nvfp4", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_calibration_keys_follow_the_qad_table(): + calibrate = _calibrate_module() + assert calibrate.amax_key("transformer.transformer_blocks.7.ff.fc_out") == "b7.ff.fc_out" + assert calibrate.amax_key("transformer_blocks.3._checkpoint_wrapped_module.ff.fc_in") == "b3.ff.fc_in" + assert calibrate.amax_key("transformer_blocks.3.attn.to_q") is None + + +def test_calibration_candidates_and_outliers(): + calibrate = _calibrate_module() + rows = [{"id": f"r{i}", "amax": {"b0.ff.fc_in": [float(i), 0.0, 1.0], "b0.ff.fc_out": [1.0, 2.0, 3.0]}} + for i in range(1, 101)] + rows[0]["amax"]["b0.ff.fc_out"] = [50.0, 0.0, 0.0] + table = calibrate.candidate_table(rows) + assert table["b0.ff.fc_in"]["max"] == 100.0 + assert table["b0.ff.fc_in"]["0.5x_max"] == 50.0 + assert table["b0.ff.fc_in"]["row_p90"] == pytest.approx(90.1) + assert set(calibrate.outlier_row_ids(rows, 2)) == {"r1", "r100"} + + + +def test_calibration_candidates_accept_a_single_row(tmp_path): + import argparse + calibrate = _calibrate_module() + row = {"id": "r1", "amax": {"b0.ff.fc_in": [2.0, 1.0, 3.0]}} + (tmp_path / "collect-shard0-dp0.jsonl").write_text(json.dumps(row) + "\n") + calibrate.run_candidates(argparse.Namespace(output_dir=str(tmp_path), outlier_rows=8)) + assert json.loads((tmp_path / "amax-max.json").read_text()) == {"b0.ff.fc_in": 3.0} + +def test_calibration_choice_all_tokens_or_balanced(): + calibrate = _calibrate_module() + candidates = {"b0.ff.fc_in": {"max": 10.0, "0.5x_max": 5.0}} + # "max" wins on all-token error (video dominates) but loses badly on audio; the balanced objective picks the clip. + mse = {"b0.ff.fc_in": {"ref": [100.0, 1.0, 10.0], "max": [1.0, 0.5, 0.1], "0.5x_max": [2.0, 0.05, 0.1]}} + table, report = calibrate.choose_table(mse, candidates, "balanced") + assert table == {"b0.ff.fc_in": 5.0} + assert report["b0.ff.fc_in"]["best"] == "0.5x_max" + table, _ = calibrate.choose_table(mse, candidates, "all") + assert table == {"b0.ff.fc_in": 10.0} + + +def _builder_module(): + path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py" + spec = importlib.util.spec_from_file_location("build_flex_ref2va_rows", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_PROMPT = "subject_definitions: is the exact first frame of the target video.\n\nsummary: A dog runs." +_ASSET_A = "asset-" + "a" * 32 + + +def _entry(case: str, clip: str, parquet: str, resolution: str = "480p") -> dict: + entry_id = f"{case}:{clip}" if case in ("first_frame", "first_last_frame", "storyboard") else clip + return {"case": case, "id": entry_id, "parquet": parquet, "resolution": resolution} + + +def test_builder_adds_subject_definitions_in_order(): + builder = _builder_module() + after = builder._subject(_PROMPT, "