mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-12 04:40:29 +08:00
jingyu-mlandClaude Opus 4.8
6b4ad85849
### What does this PR do? Type of change: New feature Adds **Qwen-Image** (`Qwen/Qwen-Image`, `QwenImageTransformer2DModel`) to the diffusers quantization example and exports HuggingFace checkpoints in three precisions — **FP8**, **NVFP4**, and **NVFP4 + SVDQuant** — through the unified HF export. - Registers `--model qwen-image` (lazy diffusers import; no `trust_remote_code`). - Transformer-block-range recipe: quantizes only the linears under `transformer_blocks`, keeping the **first 2 / last 2** blocks (and everything outside `transformer_blocks`) in original precision. Applied **before** calibration so SVDQuant never mutates the excluded blocks. Expressed with the top-level `enable` `QuantizerCfgEntry` field (disable-all → re-enable `transformer_blocks` → disable first/last-N). - SVDQuant export (AWQ-style): promotes quantizer-owned tensors to clean module-level safetensors keys at export time — `weight_quantizer.svdquant_lora_a/b → <module>.svdquant_lora_a/b` and `input_quantizer._pre_quant_scale → <module>.pre_quant_scale` — with a documented `NVFP4_SVD` `quantization_config` (`group_size`, `has_zero_point: false`, `pre_quant_scale: true`, `lora_rank`). **Core SVDQuant quantization code (`modelopt/torch/quantization`) is unchanged.** - Shared export-path change — **intentionally global** (applies to all diffusers exports — SDXL / Flux / Wan, not just Qwen; the full export suite was verified green on GB200): `hide_quantizers_from_state_dict` now strips quantizer state from *all* modules (not just quant-linears) so calibrated norm-layer input quantizers no longer leak `input_quantizer._amax`. (An earlier `max_shard_size` workaround was dropped after merging `main`: #1794 makes the ComfyUI layerwise-metadata post-processing a no-op unless explicitly opted in, so a default sharded export no longer hits the unsupported-sharded path.) ### Usage ```bash python examples/diffusers/quantization/quantize.py \ --model qwen-image --override-model-path <Qwen-Image> --model-dtype BFloat16 \ --format fp4 --quant-algo svdquant --lowrank 32 \ --calib-size 64 --n-steps 20 \ --hf-ckpt-dir <out> # FP8: --format fp8 --quant-algo max # NVFP4: --format fp4 --quant-algo max ``` ### Testing - Focused unit + example tests pass on GB200 (sm_100): block-range recipe, `NVFP4_SVD` config schema, SVDQuant forward/fold (LoRA stays on `weight_quantizer`), Qwen dummy-input / strict-QKV-fusion / promotion, pipeline loading, and the diffusers HF-export test for Qwen FP8 / NVFP4 / SVDQuant. - Full `tests/examples/diffusers/test_export_diffusers_hf_ckpt.py` is green (SDXL, Flux, Qwen, Wan2.2) — confirms the shared export changes do not regress other models. - End-to-end on the real `Qwen/Qwen-Image` (~20B): all three formats export valid HF checkpoints — only `transformer_blocks` 2..57 quantized, nothing outside, no quantizer/`_amax` leak, correct `weight_scale`(`_2`)/`input_scale`, promoted SVDQuant keys (rank-consistent shapes), and the expected `quantization_config`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ <!-- live-model LoRA storage unchanged; existing exports unaffected --> - 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?: ❌ <!-- new-example feature; add a CHANGELOG.rst entry if required --> - Did you get Claude approval on this PR?: ❌ <!-- run /claude review --> ### Additional Information All changes are confined to the diffusers example (`examples/diffusers/quantization`) plus the shared export path (`modelopt/torch/export`); the core quantization library is untouched. ### Follow-up (next step): fused-QKV SVDQuant for sglang / Nunchaku This export keeps attention `q/k/v` (and `add_q/k/v_proj`) as **separate** projections — the diffusers-native layout. That matches sglang's bf16 / FP8 / plain-NVFP4 paths (which also keep QKV separate) and ModelOpt/TRT-LLM consumers, so those load 1:1. sglang's **NVFP4-SVDQuant (Nunchaku)** path, however, builds a **fused** `to_qkv` with a *single* fused rank-r LoRA in Nunchaku-native format (`proj_down`/`proj_up`, `smooth_factor`, `wscales`/`wtscale`). Our per-projection tensors (`svdquant_lora_a/b` + `pre_quant_scale`; three independent rank-r decompositions) are not directly loadable there — and cannot be fused at load time, because the fp16 weight residual needed to derive a single fused rank-r is not preserved after export. **Planned next step:** an opt-in fused-QKV SVDQuant export mode that fuses q/k/v **before** SVDQuant calibration (yielding one rank-r over the fused weight) and emits a Nunchaku-compatible layout, enabling lower-latency fused-QKV inference in sglang. Tracked as a separate follow-up. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added Qwen-Image (`QWEN_IMAGE`) model quantization and Diffusers export support * Added NVFP4_SVD (SVDQuant) export configuration support * Added transformer block-range quantization recipes (exclude first/last blocks) * **Bug Fixes** * Improved missing-pipeline error messaging for Qwen-Image * Prevented quantizer-related tensor/buffer leakage by promoting and cleaning quantizer outputs during export * **Tests** * Added Qwen-Image HF checkpoint export tests and offline fixtures * Added unit coverage for SVDQuant promotion/clean state-dict keys <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>