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