docs/nvfp4.md

# Calibrated NVFP4 checkpoint and quantization study

## Quantization recipes

Each recipe quantizes only the 50 blocks' GEMMs (q/k/v fused, as ComfyUI stores them); AdaLN, the text refiner,
embedders and output heads keep their shipped precision. Scores are the mean relative error of the predicted
video/audio velocity against BF16, one denoising step at a time, on 8 held-out prompts (960×544, 124 frames).

![Recipe comparison](../assets/recipes.png)

<!-- BEGIN:recipes -->
| Recipe | Video error | Audio error | ≈ Size |
|---|---|---|---|
| Unquantized, factored AdaLN (noise floor) | 1.0% | 0.6% | 40 GB |
| FP8 W8A8 (all GEMMs) | 7.3% | 6.4% | 21 GB |
| NVFP4 FFN + FP8 attention | 13.3% | 10.8% | 16 GB |
| NVFP4, blocks 0-1/48-49 BF16 | 15.5% | 14.1% | 15 GB |
| NVFP4, local-Hessian scales | 15.8% | 14.9% | 13 GB |
| NVFP4 weights only (W4A16, RTN) | 16.0% | 13.6% | 13 GB |
| NVFP4, MSE scales | 16.9% | 15.6% | 13 GB |
| NVFP4, max scales (naive cast) | 18.4% | 16.7% | 13 GB |
| NVFP4, learned scales (overfit) \* | 21.2% | 16.2% | 13 GB |

\* Separate run; its own unquantized floor was 1.7% video and its local-Hessian start point 15.6% video, so compare it to that start point.
<!-- END:recipes -->

![Quality vs size](../assets/tradeoff.png)

![Error across denoising steps](../assets/steps.png)

<!-- BEGIN:steps -->
| Recipe | step 0 | step 4 | step 7 | step 11 | step 14 | step 18 |
|---|---|---|---|---|---|---|
| FP8 W8A8 (all GEMMs) | 6.1% | 13.6% | 10.5% | 6.0% | 4.0% | 3.5% |
| NVFP4 FFN + FP8 attention | 12.3% | 23.3% | 17.8% | 10.1% | 7.6% | 8.5% |
| NVFP4, local-Hessian scales | 15.4% | 28.4% | 20.6% | 11.9% | 8.8% | 9.6% |
| NVFP4, max scales (naive cast) | 20.4% | 32.3% | 23.3% | 13.9% | 10.0% | 10.7% |
<!-- END:steps -->

### Block-wise scale learning: a negative result

ModelOpt's `lsq` makes each NVFP4 block scale trainable (weights frozen). Tuning the 50 blocks one at a time from
the local-Hessian start point, each block fit to the BF16 stream's output, lowered the error on the calibration
prompts by 5–15% in the deep blocks, but **raised the held-out one-step error from 15.6% to 21.2% video** (audio
15.0% → 16.2%). The quantized stream drifts far from BF16 by the deep blocks (~0.25 relative MSE at block 40), so
asking each block's scales to undo that drift overfit the calibration prompts. A retry fitting each block only to
its own error, with keep/revert decided on held-out prompts, ran out of memory and was not repeated; even before
overfitting, the gains on early and middle blocks were 0.5–5%, so the ceiling looked low for this model.

![Block-wise scale learning](../assets/lsq_blocks.png)

## Published NVFP4 checkpoints, compared

Each file's stored 4-bit weights are decoded exactly into the same model, activations are quantized the way ComfyUI
does it (stored `input_scale`, else computed per call), and all are scored like the recipes above on the same
held-out prompts in one run.

<!-- BEGIN:ckpts -->
| Checkpoint | Size | Video error | Audio error |
|---|---|---|---|
| BF16, factored AdaLN (noise floor) | 40.5 GB | 0.6% | 0.6% |
| **This work**, NVFP4 local-Hessian | 12.7 GB | 15.4% | 15.3% |
| lilcheaty `pruned_nvfp4` (same layout) | 12.5 GB | 18.2% | 16.8% |
| rockerBOO `pruned_nvfp4` (FP8 qkv, BF16 out_proj + 5 edge blocks) | 20.1 GB | 12.7% | 9.1% |
<!-- END:ckpts -->

At the same size and speed, calibration lowers the error of an uncalibrated NVFP4 cast by ~15% (video). Keeping
attention in FP8 and edge blocks in BF16 buys more quality, especially audio, for ~7 GB and ~15% speed.

## Local results (ComfyUI, RTX 5090)

<!-- BEGIN:local -->
Stock ComfyUI (custom nodes disabled), NVIDIA GeForce RTX 5090, 960×544, 124 frames, PyTorch attention unless a config LoRA says otherwise. Median of 2 timed generations after a warm-up.

| Setup | File | Steps | s/step | Time per clip |
|---|---|---|---|---|
| ComfyUI repack, pruned INT8 ConvRot | ? | 20 | 3.87 | 87 s |
| this work: NVFP4, local-Hessian scales | 12.7 GB | 20 | 3.53 | 79 s |
| rockerBOO: NVFP4 MLP + INT8 ConvRot attention, edge blocks BF16 (20 GB) | 20.1 GB | 20 | 4.14 | 92 s |
| NVFP4 + 4-step LoRA merged (ComfyUI re-quantizes the layers) | 12.7 GB | 4 | – | 23 s |
| NVFP4 + 4-step LoRA in bypass mode (weights untouched) | 12.7 GB | 4 | – | 26 s |
| INT8 ConvRot + fast-attention config LoRA | 21.0 GB | 20 | 2.19 | 53 s |
| this work NVFP4 + fast-attention config LoRA | 12.7 GB | 20 | 1.76 | 43 s |
| rockerBOO NVFP4 + INT8 ConvRot + fast-attention config LoRA | 20.1 GB | 20 | 2.46 | 58 s |
| INT8 ConvRot + 4-step LoRA (merged) | 21.0 GB | 4 | – | 25 s |
| INT8 ConvRot + 4-step LoRA + Comfy attention LoRA | 21.0 GB | 4 | – | 21 s |
| INT8 ConvRot + 4-step LoRA + our attention LoRA (ref_speed_aw3) | 21.0 GB | 4 | – | 20 s |
| INT8 ConvRot + 4-step LoRA + our attention LoRA (x3) | 21.0 GB | 4 | – | 19 s |
| INT8 ConvRot + 4-step LoRA + our attention LoRA (x3.5_aw1) | 21.0 GB | 4 | – | 18 s |
<!-- END:local -->

## How these were made

1. Factor every block's AdaLN through a PCA basis of `silu(temb(t))` (rank 12, fp32).
2. Record the transformer's exact inputs/outputs at every denoising step for 8 evaluation and 16 calibration
   prompts.
3. Calibrate each ModelOpt recipe on the calibration traces; score one-step error on the evaluation traces.
4. Export to ComfyUI's format, packing weights from the calibrated ModelOpt quantizers (no re-quantization) and
   checking that the stored tensors decode to exactly what was scored.