Skip to content

Wan 2.2 T2V A14B diffusers PTQ: two-expert NVFP4-SVDQuant HF checkpoint export#1957

Draft
jingyu-ml wants to merge 4 commits into
mainfrom
jingyux/wan22-svdquant-export
Draft

Wan 2.2 T2V A14B diffusers PTQ: two-expert NVFP4-SVDQuant HF checkpoint export#1957
jingyu-ml wants to merge 4 commits into
mainfrom
jingyux/wan22-svdquant-export

Conversation

@jingyu-ml

@jingyu-ml jingyu-ml commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

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-14b now quantizes both experts by default — transformer (high-noise) first, then transformer_2 (low-noise), so the low-noise expert calibrates against the already-quantized high-noise expert, matching deployment.
  • The default Wan recipe (filter_func_wan_video: exclude the first 3 / last 3 of the 40 blocks, plus everything outside blocks) is now applied before calibration via the block_range mechanism. 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.
  • VAE backbones (--backbone vae) skip the transformer block-range recipe and keep their dedicated Conv3D recipe.
  • No svdquant_skip_layers for Wan: every quantized linear (self-attn, cross-attn incl. to_k/to_v, FFN) carries the full SVDQuant recipe (pre_quant_scale + rank-r svdquant_lora_a/b).
  • ONNX export paths untouched.

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_ckpt

Testing

  • New wan22_14b_nvfp4_svdquant case in tests/examples/diffusers/test_export_diffusers_hf_ckpt.py asserting per expert (transformer and transformer_2): NVFP4_SVD quant config + lora_rank; only blocks {3, 4} of the 8-block tiny fixture quantized and nothing outside blocks; svdquant_lora_a/b + pre_quant_scale present on all 10 block linears with rank-consistent shapes; NVFP4 weight_scale_2; no live-quantizer keys leaked.
  • Tiny Wan fixture updated: num_layers 2→8 (the block-range recipe needs ≥ 8 blocks), heads 2→4 (hidden 48, divisible by the NVFP4 block size 16; head_dim stays 12 so the RoPE axis split stays even).
  • All 11 affected example tests pass on GB200 (sm100), 653 s total: 3 Wan HF-export cases (int8 / fp8 / new svdquant), 7 Wan quantize+restore cases incl. --backbone vae fp8/fp4 (verifying the VAE gate), and the qwen_nvfp4_svdquant regression.
  • Real-model smoke: Wan2.2-T2V-A14B rank-32/64/128 SVDQuant exports with a structural verify gate (both experts 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"

  • Is this change backward compatible?: ❌ — --model wan2.2-t2v-14b previously quantized only transformer by default; it now quantizes both experts (a single-expert checkpoint of the two-expert pipeline is half-quantized). Pass --backbone transformer to restore the old behavior. The tiny Wan test fixture also grew from 2 to 8 blocks / hidden 24 to 48.
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: ✅
  • Did you update Changelog?: ✅
  • Did you get Claude approval on this PR?: ❌ (draft)

Additional 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, default 1.0 = migrate outliers fully to the weights; the diffusers example reuses --alpha). The AWQ-Lite search is gone from svdquant: two forward-loop passes instead of three, no per-alpha quantized GEMM sweep. (The transient --svdquant-alpha opt-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):

AWQ-Lite search SmoothQuant alpha=1.0
End-to-end quantize+export 1:32:48 1:12:07
Structural verify gate PASSED PASSED

Gate checks per expert: NVFP4_SVD + lora_rank=32; only blocks 3-36 quantized; 340/340 linears with svdquant_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. new test_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

…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>
@copy-pr-bot

copy-pr-bot Bot commented Jul 9, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9eeb3c5a-c70c-446f-b51f-3b5219659336

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch jingyux/wan22-svdquant-export

Comment @coderabbitai help to get the list of available commands.

@codecov

codecov Bot commented Jul 9, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.39535% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.76%. Comparing base (e911c3b) to head (1e44cbf).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
modelopt/torch/quantization/model_calib.py 80.95% 8 Missing ⚠️
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     
Flag Coverage Δ
unit 55.24% <81.39%> (+0.04%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jingyu-ml and others added 2 commits July 9, 2026 15:22
… 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>
@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1957/

Built to branch gh-pages at 2026-07-09 23:34 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

…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>
@jingyu-ml

Copy link
Copy Markdown
Contributor Author

Offline A/B, rank 32, Wan2.2-T2V-A14B (searched AWQ vs SmoothQuant α=1.0):

metric (per expert avg) AWQ-Lite search SmoothQuant α=1.0
weight recon rel-err mean 0.097 0.110
weight recon rel-err p95 0.104 0.145
weight recon rel-err max 0.165 0.177
pre_quant_scale max/min spread (median) 9.4–11.8 34.7–41.5
end-to-end quantize+export 1:32:48 1:12:07

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 pre_quant_scale spread shows α=1.0 actually flattening activations, which is the point of the migration for W4A4. Worst weight-side regressions concentrate in ffn.net.0.proj and early-block attn1 q/v (largest activation outliers, beyond what rank 32 absorbs). Decisive comparison needs a forward-pass/E2E quality eval; per-layer JSON kept alongside the ckpts.

🤖 Generated with Claude Code

@jingyu-ml

jingyu-ml commented Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

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):

ckpt lora_rank size quantize+export gate
r32 32 33 G 1:12:07 PASSED
r64 64 34 G 1:12:02 PASSED
r128 128 35 G 1:12:03 PASSED

(calib 1 × 20 steps @ 720×1280×81f smoke calibration — pipe-clean scope; quality-grade calibration is a knob rerun.)

🤖 Generated with Claude Code

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant