Wan 2.2 T2V A14B diffusers PTQ: two-expert NVFP4-SVDQuant HF checkpoint export#1957
Wan 2.2 T2V A14B diffusers PTQ: two-expert NVFP4-SVDQuant HF checkpoint export#1957jingyu-ml wants to merge 4 commits into
Conversation
…nt export - Quantize both experts by default for --model wan2.2-t2v-14b (transformer first, so transformer_2 calibrates against the quantized high-noise expert), producing a deployable two-expert checkpoint. - Apply the default Wan recipe (first-3/last-3 of 40 blocks, nothing outside blocks) BEFORE calibration via the block-range mechanism so SVDQuant leaves excluded weights bit-identical; VAE backbones keep their dedicated recipe. - Full SVDQuant on every quantized linear (self/cross-attn, FFN); no svdquant_skip_layers for Wan. - Tiny Wan fixture: 8 blocks (block-range needs >= 8) and hidden 48 (NVFP4 block-size divisible; head_dim stays 12 for even RoPE splits). - New wan22_14b_nvfp4_svdquant export test mirroring the Qwen SVDQuant structural assertions, per expert. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1957 +/- ##
==========================================
- Coverage 77.76% 77.76% -0.01%
==========================================
Files 519 519
Lines 57955 58081 +126
==========================================
+ Hits 45067 45164 +97
- Misses 12888 12917 +29
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
… search) - SVDQuantConfig.alpha (None = keep the AWQ-Lite search): when set, one SmoothQuant-style per-channel act-amax stats pass replaces AWQ-Lite's cache + 11-candidate search passes. The search optimizes pre-SVD, weight-only output MSE with activations unquantized -- a mismatched objective once the SVD low-rank branch absorbs the migrated outliers. alpha=1.0 migrates outliers fully to the weights (flat activations), per the SVDQuant paper. - Extract smoothquant's per-module smoothing into _smoothquant_postprocess and reuse it for the fixed-alpha pass (any quantizer format, not only int8; smoothquant() behavior unchanged). - diffusers example: --svdquant-alpha flag. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
…gration SVDQuant calibration is now a single per-channel act-amax stats pass at a fixed migration strength (SVDQuantConfig.alpha, default 1.0 = migrate outliers fully to the weights, which the SVD low-rank branch absorbs, per the SVDQuant paper). The AWQ-Lite alpha search is removed from this path: its objective -- pre-SVD, weight-only output MSE with activations unquantized -- does not match the shipped decomposition, and it cost a full extra forward-loop pass with one quantized GEMM per alpha candidate per linear. SVDQuant now runs two forward-loop passes instead of three. Checkpoint format is unchanged; scale values change (paper-aligned). The diffusers example reuses --alpha for the strength (the transient --svdquant-alpha flag from the previous commit is removed). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
Offline A/B, rank 32, Wan2.2-T2V-A14B (searched AWQ vs SmoothQuant α=1.0):
Reading: weight-reconstruction error is structurally biased toward the search (it is the search's objective and ignores activation quantization entirely — the search runs with input quantizers disabled). The 3–4× larger 🤖 Generated with Claude Code |
|
Rank ladder complete — all three Wan2.2-T2V-A14B SVDQuant exports on the final (SmoothQuant-only) path, each passing the full structural verify gate (both experts NVFP4_SVD, blocks 3–36 only, lora/pre_quant_scale on all 340 linears per expert, excluded region bit-identical in export precision, complete pipeline dirs):
(calib 1 × 20 steps @ 720×1280×81f smoke calibration — pipe-clean scope; quality-grade calibration is a knob rerun.) 🤖 Generated with Claude Code |
1 participant
What does this PR do?
Type of change: New example feature
Add Wan 2.2 T2V A14B NVFP4-SVDQuant HuggingFace checkpoint export to the diffusers quantization example, mirroring the Qwen-Image SVDQuant mechanics (#1706):
--model wan2.2-t2v-14bnow quantizes both experts by default —transformer(high-noise) first, thentransformer_2(low-noise), so the low-noise expert calibrates against the already-quantized high-noise expert, matching deployment.filter_func_wan_video: exclude the first 3 / last 3 of the 40blocks, plus everything outsideblocks) is now applied before calibration via theblock_rangemechanism. This is required for SVDQuant: the algorithm folds AWQ smoothing scales and subtracts the low-rank residual from every enabled linear during calibration, so post-hoc disabling (the previous filter-func-only path) would leave "excluded" weights silently mutated. With this PR they stay bit-identical to the source.--backbone vae) skip the transformer block-range recipe and keep their dedicated Conv3D recipe.svdquant_skip_layersfor Wan: every quantized linear (self-attn, cross-attn incl.to_k/to_v, FFN) carries the full SVDQuant recipe (pre_quant_scale+ rank-rsvdquant_lora_a/b).Usage
python examples/diffusers/quantization/quantize.py \ --model wan2.2-t2v-14b \ --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \ --format fp4 --quant-algo svdquant --lowrank 32 \ --batch-size 1 --calib-size 16 --n-steps 20 \ --hf-ckpt-dir ./hf_ckptTesting
wan22_14b_nvfp4_svdquantcase intests/examples/diffusers/test_export_diffusers_hf_ckpt.pyasserting per expert (transformerandtransformer_2):NVFP4_SVDquant config +lora_rank; only blocks {3, 4} of the 8-block tiny fixture quantized and nothing outsideblocks;svdquant_lora_a/b+pre_quant_scalepresent on all 10 block linears with rank-consistent shapes; NVFP4weight_scale_2; no live-quantizer keys leaked.num_layers2→8 (the block-range recipe needs ≥ 8 blocks), heads 2→4 (hidden 48, divisible by the NVFP4 block size 16;head_dimstays 12 so the RoPE axis split stays even).--backbone vaefp8/fp4 (verifying the VAE gate), and theqwen_nvfp4_svdquantregression.NVFP4_SVD; blocks 3–36 quantized; excluded blocks and non-block modules bit-identical to the source BF16; full pipeline dirs). Rank-32 in progress; results will be posted on this PR.Before your PR is "Ready for review"
--model wan2.2-t2v-14bpreviously quantized onlytransformerby default; it now quantizes both experts (a single-expert checkpoint of the two-expert pipeline is half-quantized). Pass--backbone transformerto restore the old behavior. The tiny Wan test fixture also grew from 2 to 8 blocks / hidden 24 to 48.CONTRIBUTING.md: N/AAdditional Information
Draft until the real-model rank-32/64/128 smoke exports and verify gates complete.
🤖 Generated with Claude Code
Update: AWQ-Lite search removed from the SVDQuant path
SVDQuant calibration is now SmoothQuant-style fixed migration only (
SVDQuantConfig.alpha, default1.0= migrate outliers fully to the weights; the diffusers example reuses--alpha). The AWQ-Lite search is gone fromsvdquant: two forward-loop passes instead of three, no per-alpha quantized GEMM sweep. (The transient--svdquant-alphaopt-in from the second commit was superseded in the fourth commit.)Why: the search optimized pre-SVD, weight-only output MSE with activations unquantized — a mismatched objective twice over for W4A4 SVDQuant, where the SVD low-rank branch absorbs the migrated outliers and activation flatness is the point of the migration (SVDQuant paper design).
Validation on Wan2.2-T2V-A14B (rank 32, calib 1 x 20 steps @ 720x1280x81f, GB200):
Gate checks per expert:
NVFP4_SVD+lora_rank=32; only blocks 3-36 quantized; 340/340 linears withsvdquant_lora_a/b+pre_quant_scale; excluded region bit-identical in export precision (the Wan-AI repo ships FP32 shards, so excluded tensors are pure BF16 downcasts - 0 value changes). Unit: 16/16 calib/mode tests incl. newtest_svdquant_smoothquant_calibration(asserts 2 forward loops + the SmoothQuant scale formula). Example: Wan + Qwen SVDQuant export tests pass on the new path.🤖 Generated with Claude Code