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

CHANGELOG.rst

@@ -6,6 +6,7 @@ Changelog
 
 **Backward Breaking Changes**
 
+- SVDQuant calibration now uses a fixed SmoothQuant-style outlier migration (``SVDQuantConfig.alpha``, default ``1.0`` = migrate fully to the weights; diffusers example flag ``--alpha``) instead of the AWQ-Lite alpha search. The search optimized pre-SVD, weight-only output MSE — a mismatched objective once the SVD low-rank branch absorbs the migrated outliers — and cost an extra forward-loop pass with one quantized GEMM per alpha candidate per linear. SVDQuant now runs two forward-loop passes (stats + max calibration) instead of three; existing ``svdquant`` recipes will produce different (paper-aligned) scales. Checkpoint format is unchanged.
 - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained.
 - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy <https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/auto_deploy>`_ directly together with ModelOpt PTQ in ``examples/llm_ptq``.
 
@@ -19,6 +20,7 @@ Changelog
 
 **New Features**
 
+- Add Wan 2.2 T2V A14B NVFP4-SVDQuant HF checkpoint export in the diffusers quantization example: ``--model wan2.2-t2v-14b`` now quantizes both experts (``transformer`` and ``transformer_2``) by default, and the default Wan recipe (first-3/last-3 of the 40 ``blocks`` excluded, nothing outside ``blocks`` quantized) is applied **before** calibration via the block-range mechanism so SVDQuant leaves excluded weights bit-identical. VAE backbones (``--backbone vae``) keep their dedicated recipe. See ``examples/diffusers/README.md``.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.

examples/diffusers/README.md

@@ -116,6 +116,28 @@ python quantize.py \
     --hf-ckpt-dir ./hf_ckpt
 ```
 
+#### Wan 2.2 T2V A14B NVFP4 SVDQuant [Script](./quantization/quantize.py)
+
+Wan 2.2 T2V A14B is a two-expert pipeline (high-noise `transformer` + low-noise
+`transformer_2`); with `--model wan2.2-t2v-14b` both experts are quantized by
+default, `transformer` first so the low-noise expert calibrates against the
+already-quantized high-noise expert. The recipe quantizes only the linears under
+`blocks`, keeping the **first 3 and last 3** of the 40 blocks (and everything
+outside `blocks`: text encoder, VAE, embedders, `proj_out`, ...) in original
+precision. The exclusion is applied **before calibration** so that for SVDQuant
+the excluded blocks' weights stay bit-identical to the original. Every quantized
+linear keeps the full SVDQuant recipe (AWQ `pre_quant_scale` + low-rank
+`svdquant_lora_a/b`).
+
+```sh
+python quantize.py \
+    --model wan2.2-t2v-14b \
+    --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \
+    --format fp4 --quant-algo svdquant --lowrank {32|64|128} \
+    --batch-size 1 --calib-size 16 --n-steps 20 \
+    --hf-ckpt-dir ./hf_ckpt
+```
+
 #### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM)
 
 The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/kernels/quantization/conv/README.md`](../../modelopt/torch/kernels/quantization/conv/README.md) for kernel details.

examples/diffusers/quantization/config.py

@@ -41,7 +41,11 @@ def set_quant_config_attr(quant_config, trt_high_precision_dtype, quant_algo, **
     elif quant_algo == "svdquant":
         if "lowrank" in kwargs:
             algo_cfg["lowrank"] = kwargs["lowrank"]
-        # Layers excluded from the SVDQuant algorithm (no AWQ smoothing, no
+        # SmoothQuant-style migration strength applied before the SVD
+        # (1.0 = migrate outliers fully to the weights).
+        if "alpha" in kwargs:
+            algo_cfg["alpha"] = kwargs["alpha"]
+        # Layers excluded from the SVDQuant algorithm (no smoothing, no
         # low-rank branch); they stay quantized with plain max calibration.
         if kwargs.get("skip_layers"):
             algo_cfg["skip_layers"] = kwargs["skip_layers"]

examples/diffusers/quantization/models_utils.py

@@ -204,6 +204,23 @@ def get_model_filter_func(
     },
     ModelType.WAN22_T2V_14b: {
         **_WAN_BASE_CONFIG,
+        # Wan 2.2 A14B is a two-expert pipeline: the high-noise expert
+        # (``transformer``, t >= boundary) and the low-noise expert
+        # (``transformer_2``). A deployable checkpoint quantizes both;
+        # ``transformer`` goes first so the low-noise expert calibrates against
+        # the already-quantized high-noise expert, matching deployment.
+        "backbone": ["transformer", "transformer_2"],
+        # Pre-calibration form of ``filter_func_wan_video``: quantize only the
+        # linears under ``blocks``, excluding the first/last 3 of the 40 blocks
+        # (everything outside ``blocks`` -- patch_embedding, condition_embedder,
+        # proj_out -- stays in original precision). Applied before calibration
+        # via ``build_block_range_quant_cfg`` so SVDQuant never mutates the
+        # excluded blocks' weights. Not applied to VAE backbones (no ``blocks``).
+        "block_range": {
+            "exclude_first_n": 3,
+            "exclude_last_n": 3,
+            "block_module": "blocks",
+        },
         "from_pretrained_extra_args": {
             "boundary_ratio": 0.875,
         },

examples/diffusers/quantization/quantize.py

@@ -108,12 +108,18 @@ def __init__(
         self.model_config = model_config
         self.logger = logger
 
-    def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any:
+    def get_quant_config(
+        self, n_steps: int, backbone: torch.nn.Module, backbone_name: str = "transformer"
+    ) -> Any:
         """
         Build quantization configuration based on format.
 
         Args:
             n_steps: Number of denoising steps
+            backbone: Backbone module the config is built for
+            backbone_name: Name of the backbone in the pipeline; VAE-type
+                backbones ("vae", "video_decoder") skip the transformer-block
+                recipe below.
 
         Returns:
             Quantization configuration object
@@ -168,12 +174,19 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any:
                 }
             )
 
-        # Apply the transformer-block-range recipe (e.g. Qwen-Image) BEFORE
-        # calibration. This restricts quantization to `transformer_blocks` and
-        # excludes the first/last N blocks. It must run before calibration so that
-        # SVDQuant does not mutate the weights of the excluded blocks. The recipe
-        # is format-agnostic (applies to FP8/NVFP4/SVDQuant alike).
+        # Apply the transformer-block-range recipe (e.g. Qwen-Image, Wan 2.2)
+        # BEFORE calibration. This restricts quantization to the transformer's
+        # block list and excludes the first/last N blocks. It must run before
+        # calibration so that SVDQuant does not mutate the weights of the
+        # excluded blocks. The recipe is format-agnostic (applies to
+        # FP8/NVFP4/SVDQuant alike) but transformer-only: VAE backbones have no
+        # block list and keep their dedicated recipe (e.g. Wan `--backbone vae`).
         block_range = MODEL_DEFAULTS.get(self.model_config.model_type, {}).get("block_range")
+        if block_range is not None and backbone_name in ("vae", "video_decoder"):
+            self.logger.info(
+                f"Skipping transformer-block-range recipe for VAE backbone '{backbone_name}'."
+            )
+            block_range = None
         if block_range is not None:
             recipe_rules = build_block_range_quant_cfg(
                 backbone,
@@ -543,7 +556,16 @@ def create_argument_parser() -> argparse.ArgumentParser:
         choices=[c.value for c in CollectMethod],
         help="Calibration collection method, works for INT8, not including smoothquant",
     )
-    quant_group.add_argument("--alpha", type=float, default=1.0, help="SmoothQuant alpha parameter")
+    quant_group.add_argument(
+        "--alpha",
+        type=float,
+        default=1.0,
+        help=(
+            "SmoothQuant/SVDQuant migration strength in [0, 1]. For SVDQuant, 1.0 (default) "
+            "migrates outliers fully from activations to weights, letting the SVD low-rank "
+            "branch absorb them."
+        ),
+    )
     quant_group.add_argument("--lowrank", type=int, default=32, help="SVDQuant lowrank parameter")
     quant_group.add_argument(
         "--quantize-mha", action="store_true", help="Quantizing MHA into FP8 if its True"
@@ -611,7 +633,12 @@ def main() -> None:
 
     model_type = ModelType(args.model)
     if args.backbone is None:
-        args.backbone = [MODEL_DEFAULTS[model_type]["backbone"]]
+        # Model defaults may name a single backbone or several (e.g. Wan 2.2
+        # A14B's two experts).
+        default_backbone = MODEL_DEFAULTS[model_type]["backbone"]
+        args.backbone = (
+            [default_backbone] if isinstance(default_backbone, str) else list(default_backbone)
+        )
     s = time.time()
 
     model_dtype = {"default": DataType(args.model_dtype).torch_dtype}
@@ -696,7 +723,9 @@ def main() -> None:
 
             for backbone_name, backbone in pipeline_manager.iter_backbones():
                 logger.info(f"Quantizing backbone: {backbone_name}")
-                backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone)
+                backbone_quant_config = quantizer.get_quant_config(
+                    calib_config.n_steps, backbone, backbone_name=backbone_name
+                )
 
                 # Calibration runs the full pipeline (not just `mod`), so the
                 # closure intentionally ignores the backbone argument.

examples/diffusers/quantization/quantize_config.py

@@ -75,7 +75,7 @@ class QuantizationConfig:
     algo: QuantAlgo = QuantAlgo.MAX
     percentile: float = 1.0
     collect_method: CollectMethod = CollectMethod.DEFAULT
-    alpha: float = 1.0  # SmoothQuant alpha
+    alpha: float = 1.0  # SmoothQuant/SVDQuant migration strength
     lowrank: int = 32  # SVDQuant lowrank
     quantize_mha: bool = False
     compress: bool = False

modelopt/torch/quantization/config.py

@@ -1154,6 +1154,20 @@ class SVDQuantConfig(QuantizeAlgorithmConfig):
         ),
     )
 
+    alpha: float = ModeloptField(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        title="SmoothQuant-style migration strength",
+        description=(
+            "Per-channel activation amax is collected in a single pass and SmoothQuant-style "
+            "smoothing ``scale = w_amax^(1-alpha) / act_amax^alpha`` is applied before the "
+            "SVD. The default ``alpha=1.0`` migrates outliers fully from activations to "
+            "weights (flat activations), letting the SVD low-rank branch absorb them, as in "
+            "the SVDQuant paper."
+        ),
+    )
+
 
 class GPTQCalibConfig(QuantizeAlgorithmConfig):
     """The config for GPTQ quantization.