From 3becd183ea3d003a8b891874ae7c570c3211411a Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 9 Jul 2026 14:37:18 -0700 Subject: [PATCH 1/4] Wan 2.2 T2V A14B diffusers PTQ: two-expert NVFP4-SVDQuant HF checkpoint 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 Signed-off-by: Jingyu Xin --- CHANGELOG.rst | 1 + examples/diffusers/README.md | 22 ++++ .../diffusers/quantization/models_utils.py | 17 +++ examples/diffusers/quantization/quantize.py | 36 ++++-- tests/_test_utils/torch/diffusers_models.py | 11 +- .../test_export_diffusers_hf_ckpt.py | 104 +++++++++++++++++- 6 files changed, 176 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 48aed5eb8b9..653d462ac90 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,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 `_) 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/``. diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index a9efb5fc3a3..7e1305cc51c 100644 --- a/examples/diffusers/README.md +++ b/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. diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 4d1bd803305..ebbf262c746 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/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, }, diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..7fb0349f980 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/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, @@ -611,7 +624,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 +714,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. diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index c680c64bc31..145b8981ea0 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -222,8 +222,11 @@ def get_tiny_wan22_vae(**config_kwargs): def create_tiny_wan22_pipeline_dir(tmp_path: Path) -> Path: """Create and save a tiny Wan 2.2 (14B-style) pipeline to a directory. - Uses the same tiny config as diffusers' own Wan 2.2 tests: - - Transformer: 2 heads, 12 head_dim, 2 layers (hidden_dim=24) + Scaled down from diffusers' own tiny Wan 2.2 test config: + - Transformer: 4 heads, 12 head_dim (hidden_dim=48, divisible by the NVFP4 + block size of 16; head_dim stays 12 so the RoPE axis split 4/4/4 remains + even); ``num_layers=8`` so the first-3/last-3 block-range recipe (which + needs >= 8 blocks) leaves blocks {3, 4} quantized - VAE: base_dim=3, z_dim=16 - Text encoder: hf-internal-testing/tiny-random-t5 (hidden_size=32) - Dual transformer (14B style) with boundary_ratio=0.875 @@ -237,10 +240,10 @@ def create_tiny_wan22_pipeline_dir(tmp_path: Path) -> Path: vae = get_tiny_wan22_vae() torch.manual_seed(0) - transformer = get_tiny_wan22_transformer() + transformer = get_tiny_wan22_transformer(num_layers=8, num_attention_heads=4) torch.manual_seed(0) - transformer_2 = get_tiny_wan22_transformer() + transformer_2 = get_tiny_wan22_transformer(num_layers=8, num_attention_heads=4) scheduler = UniPCMultistepScheduler( prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index ede9693cffc..08b791a528b 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -175,13 +175,13 @@ def _module_prefixes(keys: set[str], suffix: str) -> set[str]: return {k[: -len(suffix)] for k in keys if k.endswith(suffix)} -def _block_indices(prefixes: set[str]) -> set[int]: - """transformer_blocks indices referenced by a set of module prefixes.""" +def _block_indices(prefixes: set[str], block_re: str = r"transformer_blocks\.(\d+)\.") -> set[int]: + """Block indices referenced by a set of module prefixes.""" import re indices = set() for prefix in prefixes: - match = re.search(r"transformer_blocks\.(\d+)\.", prefix) + match = re.search(block_re, prefix) if match: indices.add(int(match.group(1))) return indices @@ -310,12 +310,34 @@ def test_qwen_image_hf_ckpt_export( assert any(k.endswith(".weight_scale_2") for k in keys) +# Tiny Wan 2.2 fixture has 8 blocks; the recipe excludes the first 3 and last 3, +# so only blocks 3 and 4 are quantized. +_WAN22_QUANTIZED_BLOCKS = {3, 4} +_WAN22_LORA_RANK = 8 +_WAN22_BLOCK_RE = r"^blocks\.(\d+)\." +# Wan has no svdquant_skip_layers: every quantized linear in a block (self-attn, +# cross-attn, FFN) keeps the full SVDQuant recipe (low-rank branch + pre_quant_scale). +_WAN22_SVDQUANT_PROMOTED_SUFFIXES = ( + ".attn1.to_q", + ".attn1.to_k", + ".attn1.to_v", + ".attn1.to_out.0", + ".attn2.to_q", + ".attn2.to_k", + ".attn2.to_v", + ".attn2.to_out.0", + ".ffn.net.0.proj", + ".ffn.net.2", +) + + class Wan22HfExportModel(NamedTuple): model: str backbone: str | None format_type: str quant_algo: str collect_method: str + lowrank: int | None = None def _suffix(self) -> str: stem = self.model.replace("wan2.2-t2v-", "") @@ -359,6 +381,8 @@ def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path: ] if self.backbone is not None: cmd_args.extend(["--backbone", self.backbone]) + if self.lowrank is not None: + cmd_args.extend(["--lowrank", str(self.lowrank)]) run_example_command(cmd_args, "diffusers/quantization") return hf_ckpt_dir @@ -371,10 +395,17 @@ def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path: Wan22HfExportModel("wan2.2-t2v-14b", None, "fp8", "max", "default"), marks=minimum_sm(89), ), + pytest.param( + Wan22HfExportModel( + "wan2.2-t2v-14b", None, "fp4", "svdquant", "default", lowrank=_WAN22_LORA_RANK + ), + marks=minimum_sm(89), + ), ], ids=[ "wan22_14b_transformer_int8_smoothquant", "wan22_14b_transformer_fp8_max", + "wan22_14b_nvfp4_svdquant", ], ) def test_wan22_hf_ckpt_export( @@ -389,3 +420,70 @@ def test_wan22_hf_ckpt_export( weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin")) assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" + + if wan_model.quant_algo != "svdquant": + return + + from safetensors import safe_open + + # Wan 2.2 A14B is a two-expert pipeline (high-noise `transformer` + + # low-noise `transformer_2`); the default backbone recipe quantizes both. + for expert in ("transformer", "transformer_2"): + expert_dir = hf_ckpt_dir / expert + config_path = expert_dir / "config.json" + assert config_path.exists(), f"no {expert}/config.json in {hf_ckpt_dir}" + quant_config = json.loads(config_path.read_text()).get("quantization_config") + assert quant_config is not None, f"{expert}: missing quantization_config" + assert quant_config.get("quant_method") == "modelopt" + assert quant_config.get("quant_algo") == "NVFP4_SVD" + group = next(iter(quant_config.get("config_groups", {}).values()), {}) + assert group.get("lora_rank") == _WAN22_LORA_RANK + assert group.get("pre_quant_scale") is True + assert quant_config.get("ignore"), f"{expert}: expected excluded modules in 'ignore'" + + keys: set[str] = set() + lora_tensors: dict[str, object] = {} + safetensors_files = sorted(expert_dir.rglob("*.safetensors")) + assert safetensors_files, f"no safetensors in {expert_dir}" + for path in safetensors_files: + with safe_open(str(path), framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not iterable + keys.add(key) + if key.endswith((".svdquant_lora_a", ".svdquant_lora_b")): + lora_tensors[key] = handle.get_tensor(key) + + # No live quantizer state should leak into the exported checkpoint. + assert not any("weight_quantizer" in k for k in keys), f"{expert}: quantizer keys leaked" + assert not any("input_quantizer._amax" in k for k in keys) + + # Recipe: only the middle `blocks` are quantized — first-3/last-3 are + # excluded, and nothing outside `blocks`. + weight_scale_prefixes = _module_prefixes(keys, ".weight_scale") + assert weight_scale_prefixes, f"{expert}: no quantized linears found in export" + assert all(p.startswith("blocks.") for p in weight_scale_prefixes), ( + f"{expert}: a non-blocks module was quantized: {weight_scale_prefixes}" + ) + assert _block_indices(weight_scale_prefixes, _WAN22_BLOCK_RE) == _WAN22_QUANTIZED_BLOCKS, ( + f"{expert}: expected only blocks {_WAN22_QUANTIZED_BLOCKS} quantized" + ) + + # Every quantized linear keeps the full SVDQuant recipe (no skip patterns): + # low-rank factors + pre_quant_scale on self-attn, cross-attn, and FFN. + expected_promoted = { + f"blocks.{block}{suffix}" + for block in _WAN22_QUANTIZED_BLOCKS + for suffix in _WAN22_SVDQUANT_PROMOTED_SUFFIXES + } + a_prefixes = _module_prefixes(keys, ".svdquant_lora_a") + b_prefixes = _module_prefixes(keys, ".svdquant_lora_b") + pqs_prefixes = _module_prefixes(keys, ".pre_quant_scale") + assert a_prefixes == b_prefixes == pqs_prefixes == expected_promoted + assert weight_scale_prefixes == expected_promoted + # Rank-consistent shapes; lora_a=[rank, in], lora_b=[out, rank], rank == --lowrank. + for key, tensor in lora_tensors.items(): + if key.endswith(".svdquant_lora_a"): + assert tensor.shape[0] == _WAN22_LORA_RANK + else: + assert tensor.shape[1] == _WAN22_LORA_RANK + # NVFP4 secondary scales are present. + assert any(k.endswith(".weight_scale_2") for k in keys), f"{expert}: missing weight_scale_2" From 9c9e168e1f414444fd23babf1bb4afe0b395bea3 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 9 Jul 2026 15:22:36 -0700 Subject: [PATCH 2/4] SVDQuant: optional fixed migration strength (skips the AWQ-Lite alpha 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 Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/config.py | 4 + examples/diffusers/quantization/quantize.py | 13 ++ .../diffusers/quantization/quantize_config.py | 2 + modelopt/torch/quantization/config.py | 16 +++ modelopt/torch/quantization/model_calib.py | 119 +++++++++++++----- tests/unit/torch/quantization/test_calib.py | 36 ++++++ 6 files changed, 160 insertions(+), 30 deletions(-) diff --git a/examples/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py index cb8fdf3a5da..5bd6c0c5ea5 100644 --- a/examples/diffusers/quantization/config.py +++ b/examples/diffusers/quantization/config.py @@ -41,6 +41,10 @@ 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"] + # Fixed migration strength: skip the AWQ-Lite alpha search and smooth + # SmoothQuant-style at this alpha before the SVD. + if kwargs.get("svdquant_alpha") is not None: + algo_cfg["alpha"] = kwargs["svdquant_alpha"] # Layers excluded from the SVDQuant algorithm (no AWQ smoothing, no # low-rank branch); they stay quantized with plain max calibration. if kwargs.get("skip_layers"): diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 7fb0349f980..be9f2ed0105 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -224,6 +224,7 @@ def get_quant_config( self.config.algo.value, alpha=self.config.alpha, lowrank=self.config.lowrank, + svdquant_alpha=self.config.svdquant_alpha, skip_layers=svdquant_skip_layers, ) self.logger.info(f"Quant config {quant_config}") @@ -558,6 +559,17 @@ def create_argument_parser() -> argparse.ArgumentParser: ) quant_group.add_argument("--alpha", type=float, default=1.0, help="SmoothQuant alpha parameter") quant_group.add_argument("--lowrank", type=int, default=32, help="SVDQuant lowrank parameter") + quant_group.add_argument( + "--svdquant-alpha", + type=float, + default=None, + help=( + "Fixed SVDQuant migration strength in [0, 1]; skips the AWQ-Lite alpha search " + "(one SmoothQuant-style stats pass instead). 1.0 migrates outliers fully from " + "activations to weights, letting the SVD low-rank branch absorb them. " + "Default: None (AWQ-Lite search)." + ), + ) quant_group.add_argument( "--quantize-mha", action="store_true", help="Quantizing MHA into FP8 if its True" ) @@ -662,6 +674,7 @@ def main() -> None: collect_method=CollectMethod(args.collect_method), alpha=args.alpha, lowrank=args.lowrank, + svdquant_alpha=args.svdquant_alpha, quantize_mha=args.quantize_mha, compress=args.compress, block_size=args.block_size, diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index a92dd4e8147..6d02d48150b 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -77,6 +77,8 @@ class QuantizationConfig: collect_method: CollectMethod = CollectMethod.DEFAULT alpha: float = 1.0 # SmoothQuant alpha lowrank: int = 32 # SVDQuant lowrank + # Fixed SVDQuant migration strength; None keeps the AWQ-Lite alpha search. + svdquant_alpha: float | None = None quantize_mha: bool = False compress: bool = False block_size: int = 16 # NVFP4 block size diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 0ca30d18448..02d54b0ec16 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1154,6 +1154,22 @@ class SVDQuantConfig(QuantizeAlgorithmConfig): ), ) + alpha: float | None = ModeloptField( + default=None, + ge=0.0, + le=1.0, + title="Fixed SmoothQuant-style migration strength (skips the AWQ-Lite search)", + description=( + "When set, the AWQ-Lite alpha search (an extra forward pass with one quantized " + "GEMM per alpha candidate per linear) is skipped. Instead, 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. " + "``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. ``None`` (default) keeps the AWQ-Lite output-MSE search." + ), + ) + class GPTQCalibConfig(QuantizeAlgorithmConfig): """The config for GPTQ quantization. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7e5bb85c09b..a526e41dae9 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1097,34 +1097,6 @@ def smoothquant(model: nn.Module, forward_loop: ForwardLoop | None = None, alpha max_calibrate(model, forward_loop) - def postprocess(module): - # It is important to keep scaling math in fp32 to be numerically safe - act_amax = module.input_quantizer.amax.float() - weight_scale = module.weight.abs().amax(dim=0, keepdim=True) - device, dtype = module.weight.device, module.weight.dtype - - parallel_group = module.parallel_state.tensor_parallel_group - if is_quantized_column_parallel_linear(module) and parallel_group.is_initialized(): - dist.all_reduce(act_amax, op=dist.ReduceOp.MAX, group=parallel_group.group) - dist.all_reduce(weight_scale, op=dist.ReduceOp.MAX, group=parallel_group.group) - - scale_a = (weight_scale.pow(1 - alpha) / act_amax.pow(alpha)).squeeze() - - # Now that activation per-channel amax have been collected, use per-tensor quantization for activation - # TODO: make this a buffer after we support only heterogeneous checkpointing for MCore - module.input_quantizer._amax_for_smoothing = act_amax.cpu() - module.input_quantizer.reset_amax() - module.input_quantizer.axis = None - module.input_quantizer.amax = act_amax.amax().to(dtype=dtype, device=device) - - # Some channel could have 0 amax which causes scale_a to overflow. Explicitly mask them out here - epsilon = 1.0 / (1 << 31) - if scale_a.min() <= epsilon: - zero_mask = act_amax <= epsilon - scale_a[zero_mask] = 1 - scale_a = scale_a.clamp(min=1e-4, max=1e4) - apply_pre_quant_scale_and_smooth(module, scale_a) - name_to_module = dict(model.named_modules()) smoothed_modules = 0 for name, module in name_to_module.items(): @@ -1144,12 +1116,90 @@ def postprocess(module): ) with enable_weight_access_and_writeback(module, model, name_to_module): - postprocess(module) + _smoothquant_postprocess(module, alpha) smoothed_modules += 1 print_rank_0(f"Smoothed {smoothed_modules} modules") +def _smoothquant_postprocess(module: nn.Module, alpha: float): + """Apply SmoothQuant smoothing to one quantized linear from collected per-channel act amax. + + Computes ``scale = w_amax^(1-alpha) / act_amax^alpha``, restores the input quantizer to + per-tensor amax, and folds the scales via :func:`apply_pre_quant_scale_and_smooth`. + """ + # It is important to keep scaling math in fp32 to be numerically safe + act_amax = module.input_quantizer.amax.float() + weight_scale = module.weight.abs().amax(dim=0, keepdim=True) + device, dtype = module.weight.device, module.weight.dtype + + parallel_group = module.parallel_state.tensor_parallel_group + if is_quantized_column_parallel_linear(module) and parallel_group.is_initialized(): + dist.all_reduce(act_amax, op=dist.ReduceOp.MAX, group=parallel_group.group) + dist.all_reduce(weight_scale, op=dist.ReduceOp.MAX, group=parallel_group.group) + + scale_a = (weight_scale.pow(1 - alpha) / act_amax.pow(alpha)).squeeze() + + # Now that activation per-channel amax have been collected, use per-tensor quantization for activation + # TODO: make this a buffer after we support only heterogeneous checkpointing for MCore + module.input_quantizer._amax_for_smoothing = act_amax.cpu() + module.input_quantizer.reset_amax() + module.input_quantizer.axis = None + module.input_quantizer.amax = act_amax.amax().to(dtype=dtype, device=device) + + # Some channel could have 0 amax which causes scale_a to overflow. Explicitly mask them out here + epsilon = 1.0 / (1 << 31) + if scale_a.min() <= epsilon: + zero_mask = act_amax <= epsilon + scale_a[zero_mask] = 1 + scale_a = scale_a.clamp(min=1e-4, max=1e4) + apply_pre_quant_scale_and_smooth(module, scale_a) + + +@torch.no_grad() +def _smooth_fixed_alpha(model: nn.Module, forward_loop: ForwardLoop, alpha: float): + """Single-pass SmoothQuant-style smoothing at a fixed migration strength. + + Unlike :func:`smoothquant`, this applies to any quantizer format. Used by SVDQuant's + fixed-``alpha`` mode, where the SVD low-rank branch absorbs the weight outliers created + by the migration, so no per-layer alpha search is needed. + """ + for name, module in model.named_modules(): + if ( + is_quantized_linear(module) + and module.input_quantizer.is_enabled + and module.input_quantizer.axis is None + ): + module.input_quantizer.axis = -1 + + max_calibrate(model, forward_loop) + + name_to_module = dict(model.named_modules()) + smoothed_modules = 0 + for name, module in name_to_module.items(): + if ( + is_quantized_linear(module) + and module.weight_quantizer.is_enabled + and module.input_quantizer.is_enabled + ): + if not hasattr(module.input_quantizer, "_amax"): + warnings.warn(f"{name} is not calibrated, skip smoothing") + continue + if module.input_quantizer.axis != -1: + warnings.warn(f"Only per-channel smoothing is supported, skip {name}") + continue + + assert module.input_quantizer._amax.numel() > 1, ( + f"Error: {name} has only one channel to smooth" + ) + + with enable_weight_access_and_writeback(module, model, name_to_module): + _smoothquant_postprocess(module, alpha) + + smoothed_modules += 1 + print_rank_0(f"Smoothed {smoothed_modules} modules (fixed alpha={alpha})") + + def awq( model: nn.Module, forward_loop: ForwardLoop | None = None, @@ -1773,6 +1823,7 @@ def svdquant( forward_loop: ForwardLoop | None = None, lowrank: int = 32, skip_layers: list[str] | None = None, + alpha: float | None = None, **kwargs, ): """Lite version of SVDQuant. @@ -1821,7 +1872,15 @@ def postprocess(module, name): quantizer.disable() skipped_quantizers.append(quantizer) - awq(model, forward_loop, "awq_lite", **kwargs) + if alpha is not None: + # Fixed migration strength: one SmoothQuant-style stats pass instead of + # AWQ-Lite's cache + candidate-search passes. The search's objective + # (plain weight-quant output MSE, activations unquantized) predates the + # SVD residual anyway; with a low-rank absorber, aggressive fixed + # migration (alpha ~ 1.0) follows the SVDQuant paper. + _smooth_fixed_alpha(model, forward_loop, alpha) + else: + awq(model, forward_loop, "awq_lite", **kwargs) for quantizer in skipped_quantizers: quantizer.enable() diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 64d89141bcd..7c083d40f26 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -419,6 +419,42 @@ def test_svdquant_lora_weights(): assert lora_residual.shape == module.weight.shape +def test_svdquant_fixed_alpha_skips_search(): + """SVDQuant with a fixed migration strength: one SmoothQuant-style stats pass + (no AWQ-Lite alpha-search pass), scales following the SmoothQuant formula.""" + torch.manual_seed(0) + model = _SimpleMLP(64, 64, 64, 64) + + quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() + quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8, "alpha": 1.0} + + x = torch.randn(2, 64, 64) + calls: list[int] = [] + + def counting_loop(model): + calls.append(1) + model(x) + + mtq.quantize(model, quant_config, counting_loop) + + # Fixed alpha needs one stats pass + the final max-calibrate pass; the + # AWQ-Lite candidate-search pass (a third forward loop) is skipped. + assert len(calls) == 2 + + # alpha=1.0 migrates fully to the weights: pre_quant_scale = 1 / act_amax + # (SmoothQuant convention), so the smoothed activations are flat. + first_linear = model.net[0] + act_amax = x.abs().amax(dim=(0, 1)).float() + expected = (1.0 / act_amax).clamp(min=1e-4, max=1e4) + pre_quant_scale = first_linear.input_quantizer.pre_quant_scale.float().squeeze() + assert torch.allclose(pre_quant_scale, expected, rtol=1e-3) + + for module in model.modules(): + if isinstance(module, torch.nn.Linear): + assert module.weight_quantizer.svdquant_lora_a is not None + assert module.weight_quantizer.svdquant_lora_b is not None + + def test_layerwise_calibrate_support_gate(): class _UnsupportedModel(nn.Module): def __init__(self): From c0b8daf98faabc458ec23de7de075d8231474221 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 9 Jul 2026 15:23:20 -0700 Subject: [PATCH 3/4] Changelog: SVDQuant fixed-alpha option Co-Authored-By: Claude Fable 5 Signed-off-by: Jingyu Xin --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 653d462ac90..6f155ddfc31 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -20,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 ``SVDQuantConfig.alpha`` (diffusers example: ``--svdquant-alpha``): a fixed SmoothQuant-style migration strength for SVDQuant that skips the AWQ-Lite alpha search — one per-channel act-amax stats pass instead of AWQ-Lite's cache + candidate-search passes. ``alpha=1.0`` migrates activation outliers fully into the weights (which the SVD low-rank branch absorbs), following the SVDQuant paper. Default ``None`` keeps the AWQ-Lite search. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv: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/``. From 1e44cbf11bfcba337ac1c9b51a5e9afd3ef21fac Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 9 Jul 2026 15:48:32 -0700 Subject: [PATCH 4/4] SVDQuant: replace the AWQ-Lite search with fixed SmoothQuant-style migration 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 Signed-off-by: Jingyu Xin --- CHANGELOG.rst | 2 +- examples/diffusers/quantization/config.py | 10 +++--- examples/diffusers/quantization/quantize.py | 16 ++++----- .../diffusers/quantization/quantize_config.py | 4 +-- modelopt/torch/quantization/config.py | 18 +++++----- modelopt/torch/quantization/model_calib.py | 33 ++++++++++--------- tests/unit/torch/quantization/test_calib.py | 16 ++++----- 7 files changed, 46 insertions(+), 53 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6f155ddfc31..2346c3dd954 100755 --- a/CHANGELOG.rst +++ b/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 `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. @@ -20,7 +21,6 @@ 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 ``SVDQuantConfig.alpha`` (diffusers example: ``--svdquant-alpha``): a fixed SmoothQuant-style migration strength for SVDQuant that skips the AWQ-Lite alpha search — one per-channel act-amax stats pass instead of AWQ-Lite's cache + candidate-search passes. ``alpha=1.0`` migrates activation outliers fully into the weights (which the SVD low-rank branch absorbs), following the SVDQuant paper. Default ``None`` keeps the AWQ-Lite search. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv: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/``. diff --git a/examples/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py index 5bd6c0c5ea5..a1603d3e2c7 100644 --- a/examples/diffusers/quantization/config.py +++ b/examples/diffusers/quantization/config.py @@ -41,11 +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"] - # Fixed migration strength: skip the AWQ-Lite alpha search and smooth - # SmoothQuant-style at this alpha before the SVD. - if kwargs.get("svdquant_alpha") is not None: - algo_cfg["alpha"] = kwargs["svdquant_alpha"] - # 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"] diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index be9f2ed0105..4165d657565 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -224,7 +224,6 @@ def get_quant_config( self.config.algo.value, alpha=self.config.alpha, lowrank=self.config.lowrank, - svdquant_alpha=self.config.svdquant_alpha, skip_layers=svdquant_skip_layers, ) self.logger.info(f"Quant config {quant_config}") @@ -557,19 +556,17 @@ 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("--lowrank", type=int, default=32, help="SVDQuant lowrank parameter") quant_group.add_argument( - "--svdquant-alpha", + "--alpha", type=float, - default=None, + default=1.0, help=( - "Fixed SVDQuant migration strength in [0, 1]; skips the AWQ-Lite alpha search " - "(one SmoothQuant-style stats pass instead). 1.0 migrates outliers fully from " - "activations to weights, letting the SVD low-rank branch absorb them. " - "Default: None (AWQ-Lite search)." + "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" ) @@ -674,7 +671,6 @@ def main() -> None: collect_method=CollectMethod(args.collect_method), alpha=args.alpha, lowrank=args.lowrank, - svdquant_alpha=args.svdquant_alpha, quantize_mha=args.quantize_mha, compress=args.compress, block_size=args.block_size, diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index 6d02d48150b..ca51b91b7c2 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -75,10 +75,8 @@ 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 - # Fixed SVDQuant migration strength; None keeps the AWQ-Lite alpha search. - svdquant_alpha: float | None = None quantize_mha: bool = False compress: bool = False block_size: int = 16 # NVFP4 block size diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 02d54b0ec16..e560142b484 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1154,19 +1154,17 @@ class SVDQuantConfig(QuantizeAlgorithmConfig): ), ) - alpha: float | None = ModeloptField( - default=None, + alpha: float = ModeloptField( + default=1.0, ge=0.0, le=1.0, - title="Fixed SmoothQuant-style migration strength (skips the AWQ-Lite search)", + title="SmoothQuant-style migration strength", description=( - "When set, the AWQ-Lite alpha search (an extra forward pass with one quantized " - "GEMM per alpha candidate per linear) is skipped. Instead, 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. " - "``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. ``None`` (default) keeps the AWQ-Lite output-MSE search." + "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." ), ) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index a526e41dae9..591e5fd8e71 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1160,9 +1160,9 @@ def _smoothquant_postprocess(module: nn.Module, alpha: float): def _smooth_fixed_alpha(model: nn.Module, forward_loop: ForwardLoop, alpha: float): """Single-pass SmoothQuant-style smoothing at a fixed migration strength. - Unlike :func:`smoothquant`, this applies to any quantizer format. Used by SVDQuant's - fixed-``alpha`` mode, where the SVD low-rank branch absorbs the weight outliers created - by the migration, so no per-layer alpha search is needed. + Unlike :func:`smoothquant`, this applies to any quantizer format. It is SVDQuant's + calibration: the SVD low-rank branch absorbs the weight outliers created by the + migration, so no per-layer alpha search is needed. """ for name, module in model.named_modules(): if ( @@ -1823,11 +1823,16 @@ def svdquant( forward_loop: ForwardLoop | None = None, lowrank: int = 32, skip_layers: list[str] | None = None, - alpha: float | None = None, + alpha: float = 1.0, **kwargs, ): """Lite version of SVDQuant. + Calibration is a fixed SmoothQuant-style outlier migration (strength ``alpha``, + default 1.0 = migrate fully to the weights) followed by the SVD low-rank + decomposition that absorbs the migrated outliers, then max calibration of the + residual — two forward-loop passes total. + Args: model: Model to be calibrated. forward_loop: A callable which takes the model as argument and @@ -1855,9 +1860,10 @@ def postprocess(module, name): create_and_replace_svdquant_linear_on_the_fly(model=model) # Modules matching `skip_layers` opt out of the SVDQuant algorithm but stay - # quantized: temporarily disable their quantizers so awq_lite neither smooths - # their weights nor attaches a pre_quant_scale, then re-enable them so the - # final max calibration collects their amax like a plain max recipe. + # quantized: temporarily disable their quantizers so the smoothing pass + # neither smooths their weights nor attaches a pre_quant_scale, then + # re-enable them so the final max calibration collects their amax like a + # plain max recipe. skipped_quantizers = [] if skip_layers: for name, module in model.named_modules(): @@ -1872,15 +1878,10 @@ def postprocess(module, name): quantizer.disable() skipped_quantizers.append(quantizer) - if alpha is not None: - # Fixed migration strength: one SmoothQuant-style stats pass instead of - # AWQ-Lite's cache + candidate-search passes. The search's objective - # (plain weight-quant output MSE, activations unquantized) predates the - # SVD residual anyway; with a low-rank absorber, aggressive fixed - # migration (alpha ~ 1.0) follows the SVDQuant paper. - _smooth_fixed_alpha(model, forward_loop, alpha) - else: - awq(model, forward_loop, "awq_lite", **kwargs) + # Fixed migration strength: one SmoothQuant-style stats pass. With the SVD + # low-rank branch absorbing the migrated weight outliers, no per-layer + # search is needed (alpha ~ 1.0 follows the SVDQuant paper). + _smooth_fixed_alpha(model, forward_loop, alpha) for quantizer in skipped_quantizers: quantizer.enable() diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 7c083d40f26..847da73b325 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -419,14 +419,14 @@ def test_svdquant_lora_weights(): assert lora_residual.shape == module.weight.shape -def test_svdquant_fixed_alpha_skips_search(): - """SVDQuant with a fixed migration strength: one SmoothQuant-style stats pass - (no AWQ-Lite alpha-search pass), scales following the SmoothQuant formula.""" +def test_svdquant_smoothquant_calibration(): + """SVDQuant calibrates with a single SmoothQuant-style stats pass (no AWQ-Lite + search pass), migrating fully to the weights by default (alpha=1.0).""" torch.manual_seed(0) model = _SimpleMLP(64, 64, 64, 64) quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() - quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8, "alpha": 1.0} + quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8} x = torch.randn(2, 64, 64) calls: list[int] = [] @@ -437,12 +437,12 @@ def counting_loop(model): mtq.quantize(model, quant_config, counting_loop) - # Fixed alpha needs one stats pass + the final max-calibrate pass; the - # AWQ-Lite candidate-search pass (a third forward loop) is skipped. + # One SmoothQuant-style stats pass + the final max-calibrate pass. There is + # no AWQ-Lite candidate-search pass (a third forward loop) in SVDQuant. assert len(calls) == 2 - # alpha=1.0 migrates fully to the weights: pre_quant_scale = 1 / act_amax - # (SmoothQuant convention), so the smoothed activations are flat. + # The default alpha=1.0 migrates fully to the weights: pre_quant_scale = + # 1 / act_amax (SmoothQuant convention), so the smoothed activations are flat. first_linear = model.net[0] act_amax = x.abs().amax(dim=(0, 1)).float() expected = (1.0 / act_amax).clamp(min=1e-4, max=1e4)