From 3becd183ea3d003a8b891874ae7c570c3211411a Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 9 Jul 2026 14:37:18 -0700 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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) From 6c04b1ed288ee5e8388ffd478befe9e228fcdef6 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 11 Aug 2026 23:26:20 -0700 Subject: [PATCH 5/5] Add MiniMax-H3 diffusion quantization support --- .../diffusers/quantization/calibration.py | 41 +++++++ .../diffusers/quantization/models_utils.py | 30 +++++ .../quantization/pipeline_manager.py | 77 ++++++++++++ examples/diffusers/quantization/quantize.py | 83 ++++++++++++- examples/diffusers/quantization/utils.py | 19 +++ modelopt/torch/export/diffusers_utils.py | 25 +++- modelopt/torch/export/unified_export_hf.py | 104 +++++++++++----- .../torch/export/test_export_diffusers.py | 111 ++++++++++++++++++ 8 files changed, 457 insertions(+), 33 deletions(-) diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 27b1ec22436..34f4fe892eb 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import Any +import torch from models_utils import MODEL_DEFAULTS, ModelType from pipeline_manager import PipelineManager from quantize_config import CalibrationConfig @@ -50,6 +51,7 @@ def __init__( self.config = config self.model_type = model_type self.logger = logger + self._minimax_h3_generator: torch.Generator | None = None def load_and_batch_prompts(self) -> list[list[str]]: """ @@ -95,6 +97,8 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: elif self.model_type in [ModelType.WAN22_T2V_14b, ModelType.WAN22_T2V_5b]: # Special handling for WAN video models self._run_wan_video_calibration(prompt_batch, extra_args) + elif self.model_type == ModelType.MINIMAX_H3: + self._run_minimax_h3_calibration(prompt_batch, extra_args) else: common_args = { "prompt": prompt_batch, @@ -105,6 +109,43 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: self.logger.debug(f"Completed calibration batch {i + 1}/{self.config.num_batches}") self.logger.info("Calibration completed successfully") + def _run_minimax_h3_calibration( + self, prompt_batch: list[str], extra_args: dict[str, Any] + ) -> None: + """Calibrate MiniMax-H3 through its prompt-only T2V denoising path.""" + if len(prompt_batch) != 1: + raise ValueError( + "MiniMax-H3's T2V ModularPipeline accepts one prompt string per call; " + "use --batch-size 1 for calibration." + ) + + extra_params = self.pipeline_manager.config.extra_params + height = int(extra_params.get("height", extra_args["height"])) + width = int(extra_params.get("width", extra_args["width"])) + num_frames = int(extra_params.get("num_frames", extra_args["num_frames"])) + seed = int(extra_params.get("seed", 0)) + if self._minimax_h3_generator is None: + self._minimax_h3_generator = torch.Generator(device="cpu").manual_seed(seed) + + self.logger.debug( + "MiniMax-H3 T2V calibration call: height=%d width=%d num_frames=%d " + "num_inference_steps=%d seed=%d", + height, + width, + num_frames, + self.config.n_steps, + seed, + ) + self.pipe( + prompt=prompt_batch[0], + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=self.config.n_steps, + generator=self._minimax_h3_generator, + output=["videos", "audio", "sampling_rate"], + ) + def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index ebbf262c746..3f580f25c26 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -43,6 +43,7 @@ filter_func_flux_dev, filter_func_ltx2_vae, filter_func_ltx_video, + filter_func_minimax_h3, filter_func_qwen_image, filter_func_wan_vae, filter_func_wan_video, @@ -64,6 +65,7 @@ class ModelType(str, Enum): WAN22_T2V_14b = "wan2.2-t2v-14b" WAN22_T2V_5b = "wan2.2-t2v-5b" QWEN_IMAGE = "qwen-image" + MINIMAX_H3 = "minimax-h3" _FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { @@ -74,6 +76,7 @@ class ModelType(str, Enum): ModelType.WAN22_T2V_14b: filter_func_wan_video, ModelType.WAN22_T2V_5b: filter_func_wan_video, ModelType.QWEN_IMAGE: filter_func_qwen_image, + ModelType.MINIMAX_H3: filter_func_minimax_h3, } _VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { @@ -107,6 +110,7 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers", ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ModelType.QWEN_IMAGE: "Qwen/Qwen-Image", + ModelType.MINIMAX_H3: "MiniMaxAI/MiniMax-H3", } MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline] | None] = { @@ -122,6 +126,9 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: WanPipeline, ModelType.WAN22_T2V_5b: WanPipeline, ModelType.QWEN_IMAGE: QwenImagePipeline, + # MiniMax-H3 uses ModularPipeline and is created by PipelineManager's + # text-to-video-only path. Keep it unavailable to the generic TRT loader. + ModelType.MINIMAX_H3: None, } # Shared dataset configurations @@ -288,6 +295,29 @@ def get_model_filter_func( "*.txt_mod.1", ], }, + ModelType.MINIMAX_H3: { + "backbone": "transformer", + "dataset": _OPENVID_DATASET, + "svdquant_alpha": 0.8, + # MiniMax-H3's guidance-distilled T2V workflow consumes only a prompt; + # there is no negative prompt or guidance scale. These are the trained + # 16:9 canvas and minimum valid frame count. Calibration-only overrides + # (for example 544x960) can be supplied through --extra-param. + "inference_extra_args": { + "height": 768, + "width": 1344, + "num_frames": 124, + }, + # Apply this before calibration, especially for SVDQuant: quantize only + # transformer blocks 2..49. This keeps the first two blocks and every + # projection/embedder/final-layer module outside the block list bit- + # identical to the source weights. + "block_range": { + "exclude_first_n": 2, + "exclude_last_n": 0, + "block_module": "transformer_blocks", + }, + }, } diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..8481a28438c 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -16,6 +16,7 @@ import logging import warnings from collections.abc import Iterator +from pathlib import Path from typing import Any import torch @@ -43,6 +44,7 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None self._video_decoder: torch.nn.Module | None = None + self._components_manager: Any | None = None @staticmethod def create_pipeline_from( @@ -100,6 +102,11 @@ def create_pipeline(self) -> Any: self.logger.info("LTX-2 pipeline created successfully") return self.pipe + if self.config.model_type == ModelType.MINIMAX_H3: + self.pipe = self._create_minimax_h3_pipeline() + self.logger.info("MiniMax-H3 T2V modular pipeline created successfully") + return self.pipe + pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: raise ValueError( @@ -143,6 +150,14 @@ def setup_device(self) -> None: self.logger.info("Skipping device setup for LTX-2 pipeline (handled internally)") return + if self.config.model_type == ModelType.MINIMAX_H3: + if self.config.cpu_offloading: + self.logger.info("MiniMax-H3 ComponentsManager auto CPU offloading is enabled") + else: + self.logger.info("Moving MiniMax-H3 T2V components to CUDA") + self.pipe.to("cuda") + return + if self.config.cpu_offloading: self.logger.info("Enabling CPU offloading for memory efficiency") self.pipe.enable_model_cpu_offload() @@ -192,6 +207,68 @@ def iter_backbones(self) -> Iterator[tuple[str, torch.nn.Module]]: raise RuntimeError(f"Pipeline missing backbone module '{name}'.") yield name, module + def _create_minimax_h3_pipeline(self) -> Any: + """Load only MiniMax-H3's prompt-only T2V workflow.""" + try: + from diffusers import ComponentsManager, ModularPipeline + except ImportError as exc: + raise ImportError( + "MiniMax-H3 is not part of a Diffusers release yet. Install Diffusers " + "from source with: pip install git+https://github.com/huggingface/diffusers.git" + ) from exc + + forbidden_conditioning = { + key for key in ("image", "last_image", "references") if key in self.config.extra_params + } + if forbidden_conditioning: + raise ValueError( + "MiniMax-H3 calibration supports text-only T2V. Remove conditioning " + f"parameters: {sorted(forbidden_conditioning)}" + ) + + self._components_manager = ComponentsManager() + local_files_only = Path(self.config.model_path).exists() + pipe = ModularPipeline.from_pretrained( + self.config.model_path, + workflow="t2va", + components_manager=self._components_manager, + local_files_only=local_files_only, + ) + # The downloaded modular index records the Hub repository in every + # component spec. Override it so a local --override-model-path remains + # completely local and cannot silently fetch transformer_ref. + pipe.load_components( + dtype=self.config.model_dtype, + pretrained_model_name_or_path=self.config.model_path, + local_files_only=local_files_only, + ) + # Enable offload only after every component has been registered. Enabling + # it before load_components() makes ComponentsManager rebuild hooks on + # each add and resets the requested reserve margin to its default. + if self.config.cpu_offloading: + memory_reserve_margin = str( + self.config.extra_params.get("memory_reserve_margin", "12GB") + ) + self.logger.info( + "Enabling MiniMax-H3 automatic CPU offload with memory reserve margin %s", + memory_reserve_margin, + ) + self._components_manager.enable_auto_cpu_offload( + device="cuda", memory_reserve_margin=memory_reserve_margin + ) + pipe.set_progress_bar_config(disable=True) + + missing_components = pipe.null_component_names + if missing_components: + raise RuntimeError( + f"MiniMax-H3 T2V failed to load required components: {sorted(missing_components)}" + ) + if getattr(pipe, "transformer_ref", None) is not None: + raise RuntimeError("MiniMax-H3 T2V workflow unexpectedly loaded transformer_ref.") + loaded_components = sorted(pipe.components) + self.logger.info("Loaded MiniMax-H3 T2V components only: %s", loaded_components) + return pipe + def _ensure_ltx2_transformer_cached(self) -> None: if not self.pipe: raise RuntimeError("Pipeline not created. Call create_pipeline() first.") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 4165d657565..2c0423ba36f 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -203,6 +203,30 @@ def get_quant_config( ) quant_cfg_list.extend(recipe_rules) + if self.model_config.model_type == ModelType.MINIMAX_H3: + # The requested H3 recipe quantizes Linear weights/inputs only. The + # Diffusers NVFP4 preset enables its softmax quantizer by default; + # turn every attention-internal quantizer off BEFORE calibration so + # excluded blocks/token_refiner never run any quantized MHA path. + attention_quantizers = ( + "q_bmm_quantizer", + "k_bmm_quantizer", + "v_bmm_quantizer", + "softmax_quantizer", + "bmm2_output_quantizer", + ) + quant_cfg_list.extend( + { + "quantizer_name": f"*{quantizer_name}", + "enable": False, + } + for quantizer_name in attention_quantizers + ) + self.logger.info( + "MiniMax-H3 pre-calibration recipe disables all attention-internal quantizers: %s", + attention_quantizers, + ) + # Per-model SVDQuant exclusions (e.g. Qwen-Image's text-stream linears): # matching layers skip the SVDQuant low-rank branch and AWQ smoothing but # stay quantized with plain max calibration. @@ -559,11 +583,10 @@ def create_argument_parser() -> argparse.ArgumentParser: quant_group.add_argument( "--alpha", type=float, - default=1.0, + default=None, 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." + "SmoothQuant/SVDQuant migration strength in [0, 1]. Defaults to 0.8 for " + "MiniMax-H3 and 1.0 for other models." ), ) quant_group.add_argument("--lowrank", type=int, default=32, help="SVDQuant lowrank parameter") @@ -622,6 +645,49 @@ def create_argument_parser() -> argparse.ArgumentParser: return parser +def validate_minimax_h3_configuration( + model_config: ModelConfig, + quant_config: QuantizationConfig, + calib_config: CalibrationConfig, + export_config: ExportConfig, +) -> None: + """Validate the deliberately narrow MiniMax-H3 T2V export workflow.""" + if model_config.model_type != ModelType.MINIMAX_H3: + return + + if model_config.backbone != ["transformer"]: + raise ValueError( + "MiniMax-H3 T2V supports only --backbone transformer; transformer_ref belongs " + "to the Ref2VA workflow and must not be loaded or quantized." + ) + if quant_config.format != QuantFormat.FP4 or quant_config.algo not in ( + QuantAlgo.MAX, + QuantAlgo.SVDQUANT, + ): + raise ValueError( + "MiniMax-H3 supports only the requested recipes: '--format fp4 " + "--quant-algo max' (NVFP4) or '--format fp4 --quant-algo svdquant'." + ) + if quant_config.quantize_mha: + raise ValueError("MiniMax-H3 quantizes transformer linears only; omit --quantize-mha.") + if export_config.restore_from is None: + if calib_config.batch_size != 1: + raise ValueError( + "MiniMax-H3's T2V ModularPipeline accepts one prompt string per call; " + "set --batch-size 1." + ) + if calib_config.n_steps < 2: + raise ValueError( + "MiniMax-H3 num_inference_steps includes the terminal sigma=0, so --n-steps " + "must be at least 2 to execute a transformer forward." + ) + if export_config.onnx_dir is not None: + raise ValueError( + "MiniMax-H3 support is calibration + unified HF checkpoint export only; " + "--onnx-dir is intentionally unsupported." + ) + + def main() -> None: from diffusers.models.normalization import RMSNorm as DiffuserRMSNorm @@ -664,12 +730,18 @@ def main() -> None: extra_params=extra_params, ) + # Keep model-specific recipe values alongside the block/filter recipe in + # MODEL_DEFAULTS. Models without an override retain the existing 1.0. + alpha = args.alpha + if alpha is None: + alpha = MODEL_DEFAULTS[model_type].get("svdquant_alpha", 1.0) + quant_config = QuantizationConfig( format=QuantFormat(args.format), algo=QuantAlgo(args.quant_algo), percentile=args.percentile, collect_method=CollectMethod(args.collect_method), - alpha=args.alpha, + alpha=alpha, lowrank=args.lowrank, quantize_mha=args.quantize_mha, compress=args.compress, @@ -705,6 +777,7 @@ def main() -> None: export_config.validate() if not export_config.restore_from: calib_config.validate() + validate_minimax_h3_configuration(model_config, quant_config, calib_config, export_config) pipeline_manager = PipelineManager(model_config, logger) pipe = pipeline_manager.create_pipeline() diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..5e0e99ad99a 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -111,6 +111,25 @@ def filter_func_wan_video(name: str) -> bool: return pattern.match(name) is not None +# MiniMax-H3 has 50 ``transformer_blocks``. Only blocks 2..49 are +# quantized: the first two blocks and every module outside the block list stay +# in their original precision. In particular, this excludes the Diffusers +# equivalents of ``video_patch_proj``, ``audio_patch_proj``, +# ``condition_proj``, ``time_embedder``, ``token_refiner`` and +# ``final_layer`` (``proj_in``, ``audio_proj_in``, ``context_embedder``, +# ``time_embedder``, ``token_refiner``, ``norm_out``/``proj_out``/ +# ``audio_proj_out``). +_MINIMAX_H3_BLOCK_RE = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)(?:\.|$)") + + +def filter_func_minimax_h3(name: str) -> bool: + """Return ``True`` for MiniMax-H3 modules that must remain unquantized.""" + match = _MINIMAX_H3_BLOCK_RE.search(name) + if match is None: + return True + return int(match.group(1)) < 2 + + # Qwen-Image's transformer has 60 ``transformer_blocks``. The recipe quantizes # only those blocks while keeping the first two and last two -- and everything # outside ``transformer_blocks`` -- in original precision. The model-agnostic, diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 0309c83e1f6..e96ac4127db 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -28,6 +28,7 @@ from safetensors.torch import load_file, safe_open DiffusionPipeline: type[Any] | None +ModularPipeline: type[Any] | None ModelMixin: type[Any] | None try: # diffusers is optional for LTX-2 export paths from diffusers import DiffusionPipeline as _DiffusionPipeline @@ -35,9 +36,16 @@ DiffusionPipeline = _DiffusionPipeline ModelMixin = _ModelMixin + try: # ModularPipeline is unavailable in older supported diffusers releases + from diffusers import ModularPipeline as _ModularPipeline + + ModularPipeline = _ModularPipeline + except Exception: # pragma: no cover + ModularPipeline = None _HAS_DIFFUSERS = True except Exception: # pragma: no cover DiffusionPipeline = None + ModularPipeline = None ModelMixin = None _HAS_DIFFUSERS = False @@ -69,6 +77,8 @@ def is_diffusers_object(model: Any) -> bool: diffusers_types: tuple[type, ...] = () if DiffusionPipeline is not None: diffusers_types = (*diffusers_types, DiffusionPipeline) + if ModularPipeline is not None: + diffusers_types = (*diffusers_types, ModularPipeline) if ModelMixin is not None: diffusers_types = (*diffusers_types, ModelMixin) if TI2VidTwoStagesPipeline is not None: @@ -80,6 +90,18 @@ def is_diffusers_object(model: Any) -> bool: return isinstance(model, diffusers_types) +def is_modular_pipeline(model: Any) -> bool: + """Return whether *model* is a diffusers ``ModularPipeline``.""" + return ModularPipeline is not None and isinstance(model, ModularPipeline) + + +def is_diffusers_pipeline(model: Any) -> bool: + """Return whether *model* is a standard or modular diffusers pipeline.""" + return (DiffusionPipeline is not None and isinstance(model, DiffusionPipeline)) or ( + is_modular_pipeline(model) + ) + + def generate_diffusion_dummy_inputs( model: nn.Module, device: torch.device, dtype: torch.dtype ) -> dict[str, torch.Tensor] | None: @@ -654,6 +676,7 @@ def get_diffusion_components( Supports: - diffusers `DiffusionPipeline`: returns `pipeline.components` + - diffusers `ModularPipeline`: returns loaded `pipeline.components` - diffusers component `nn.Module` (e.g., UNet / transformer) - LTX-2 pipeline (duck-typed): returns stage-1 transformer only as `stage_1_transformer` @@ -680,7 +703,7 @@ def get_diffusion_components( return all_components # diffusers pipeline - if _HAS_DIFFUSERS and DiffusionPipeline is not None and isinstance(model, DiffusionPipeline): + if _HAS_DIFFUSERS and is_diffusers_pipeline(model): # Get all components from the pipeline all_components = {name: comp for name, comp in model.components.items() if comp is not None} diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a903adfc846..3e47544c98d 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -44,6 +44,8 @@ hide_quantizers_from_state_dict, infer_dtype_from_model, is_diffusers_object, + is_diffusers_pipeline, + is_modular_pipeline, is_qkv_projection, merge_diffusion_checkpoint, ) @@ -1261,15 +1263,9 @@ def _export_diffusers_checkpoint( name: comp for name, comp in all_components.items() if isinstance(comp, nn.Module) } - # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior) - is_diffusers_pipe = False - if HAS_DIFFUSERS: - try: - from diffusers import DiffusionPipeline as _DiffusionPipeline - - is_diffusers_pipe = isinstance(pipe, _DiffusionPipeline) - except Exception: - is_diffusers_pipe = False + # Pipeline checks control component subfolders and the pipeline index format. + is_diffusers_pipe = is_diffusers_pipeline(pipe) + is_modular_pipe = is_modular_pipeline(pipe) # Export each nn.Module component with quantization handling for component_name, component in module_components.items(): @@ -1408,40 +1404,94 @@ def _export_diffusers_checkpoint( print(f" Saved to: {component_export_dir}") - # For pipelines, also save model_index.json + # For pipelines, also save the pipeline index. if is_diffusers_pipe: - model_index_path = export_dir / "model_index.json" + model_index_filename = "modular_model_index.json" if is_modular_pipe else "model_index.json" + model_index_path = export_dir / model_index_filename is_partial_export = components is not None - # For full export, preserve original model_index.json when possible. - # For partial export, skip this to avoid listing non-exported components. - if not is_partial_export: + # For a standard pipeline, preserve the original index when possible. A + # modular pipeline must serialize its active config instead: workflows can + # prune components (for example MiniMax-H3 T2V omits transformer_ref), while + # the source index still lists components used by other workflows. + if not is_partial_export and not is_modular_pipe: source_path = getattr(pipe, "name_or_path", None) or getattr( getattr(pipe, "config", None), "_name_or_path", None ) if source_path: - candidate_model_index = Path(source_path) / "model_index.json" + candidate_model_index = Path(source_path) / model_index_filename if candidate_model_index.exists(): with open(candidate_model_index) as file: model_index = json.load(file) with open(model_index_path, "w") as file: json.dump(model_index, file, indent=4) - # Full-export fallback to Diffusers-native config serialization. - # Partial export skips this for the same reason as above. - if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): + # Diffusers-native config serialization preserves workflow metadata. Always + # refresh a modular index from the active pipeline, even when retrying into + # an existing directory: a stale source/workflow index could otherwise keep + # components (for example transformer_ref) that were not loaded or exported. + # Modular partial exports are pruned below after serializing the full config. + if ( + is_modular_pipe or (not is_partial_export and not model_index_path.exists()) + ) and hasattr(pipe, "save_config"): pipe.save_config(export_dir) - # Last resort: synthesize a minimal model_index.json from exported components. + # Last resort: synthesize an index from the pipeline config/components. if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: - model_index = { - "_class_name": type(pipe).__name__, - "_diffusers_version": diffusers.__version__, - } - for name, comp in all_components.items(): - module = type(comp).__module__ - library = module.split(".")[0] - model_index[name] = [library, type(comp).__name__] + if is_modular_pipe: + model_index = dict(pipe.config) + else: + model_index = { + "_class_name": type(pipe).__name__, + "_diffusers_version": diffusers.__version__, + } + for name, comp in all_components.items(): + module = type(comp).__module__ + library = module.split(".")[0] + model_index[name] = [library, type(comp).__name__] + + with open(model_index_path, "w") as file: + json.dump(model_index, file, indent=4) + + if is_modular_pipe and model_index_path.exists(): + with open(model_index_path) as file: + model_index = json.load(file) + + pipeline_component_names = set(getattr(pipe, "components", {})) + if is_partial_export: + for component_name in pipeline_component_names - set(all_components): + model_index.pop(component_name, None) + + # A modular component entry carries its own loading location. Point every + # exported component at the newly written local subfolder so reloading the + # unified checkpoint cannot silently fetch the original unquantized model. + for component_name, component in all_components.items(): + module = type(component).__module__ + default_library = module.split(".")[0] + entry = model_index.get(component_name) + if not (isinstance(entry, (list, tuple)) and len(entry) == 3): + # Components created from config (for example H3's + # video_processor) intentionally have no modular-index entry. + # Adding one would change them to from_pretrained components, + # which can make load_components() call a nonexistent loader. + continue + + library, class_name, loading_spec = entry + loading_spec = dict(loading_spec or {}) + library = library or default_library + class_name = class_name or type(component).__name__ + + # diffusers 0.35 used ``repo``; newer versions use the more explicit + # ``pretrained_model_name_or_path`` field. + location_key = ( + "repo" + if "repo" in loading_spec + and "pretrained_model_name_or_path" not in loading_spec + else "pretrained_model_name_or_path" + ) + loading_spec[location_key] = str(export_dir) + loading_spec["subfolder"] = component_name + model_index[component_name] = [library, class_name, loading_spec] with open(model_index_path, "w") as file: json.dump(model_index, file, indent=4) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 753c81a4b0e..0d62f53b709 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -15,6 +15,7 @@ import copy import json +from pathlib import Path import pytest import torch @@ -31,12 +32,15 @@ from safetensors import safe_open from safetensors.torch import save_file +import modelopt.torch.export.diffusers_utils as diffusers_utils import modelopt.torch.export.unified_export_hf as unified_export_hf import modelopt.torch.quantization as mtq from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.diffusers_utils import ( generate_diffusion_dummy_inputs, + get_diffusion_components, hide_quantizers_from_state_dict, + is_diffusers_object, ) from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint @@ -72,6 +76,25 @@ def _read_safetensors_metadata(path): return dict(file.metadata() or {}) +class _FakeConfigComponent: + def save_config(self, save_directory): + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + with open(save_directory / "scheduler_config.json", "w") as file: + json.dump({"_class_name": type(self).__name__}, file) + + +class _FakeModularPipeline: + def __init__(self, components, config, source_path=None): + self.components = components + self.config = config + self._pretrained_model_name_or_path = source_path + + def save_config(self, save_directory): + with open(Path(save_directory) / "modular_model_index.json", "w") as file: + json.dump(self.config, file, indent=4) + + @pytest.mark.parametrize( "model_factory", [get_tiny_unet, get_tiny_dit, get_tiny_flux, get_tiny_flux2] ) @@ -88,6 +111,94 @@ def test_export_diffusers_models_non_quantized(tmp_path, model_factory): assert "quantization_config" not in config_data +@pytest.mark.parametrize("location_key", ["pretrained_model_name_or_path", "repo"]) +def test_export_modular_pipeline_components_and_index(tmp_path, monkeypatch, location_key): + """A modular pipeline exports local components and keeps a modular loading index.""" + monkeypatch.setattr(diffusers_utils, "ModularPipeline", _FakeModularPipeline) + + transformer = get_tiny_dit() + scheduler = _FakeConfigComponent() + video_processor = _FakeConfigComponent() + original_location = "upstream/original-model" + config = { + "_blocks_class_name": "FakeTextToVideoBlocks", + "_class_name": "FakeModularPipeline", + "workflow": "text2video", + "transformer": [ + "diffusers", + type(transformer).__name__, + { + location_key: original_location, + "subfolder": "transformer", + "type_hint": ["diffusers", type(transformer).__name__], + }, + ], + "scheduler": [ + "diffusers", + type(scheduler).__name__, + { + location_key: original_location, + "subfolder": "scheduler", + "type_hint": ["diffusers", type(scheduler).__name__], + }, + ], + } + source_dir = tmp_path / "source" + source_dir.mkdir() + source_index = { + **config, + # The source repository can describe a component belonging to another + # workflow even though it is absent from this workflow's active config. + "transformer_ref": [ + "diffusers", + type(transformer).__name__, + { + location_key: original_location, + "subfolder": "transformer_ref", + "type_hint": ["diffusers", type(transformer).__name__], + }, + ], + } + with open(source_dir / "modular_model_index.json", "w") as file: + json.dump(source_index, file) + + pipe = _FakeModularPipeline( + { + "transformer": transformer, + "scheduler": scheduler, + # Config-created components have no modular-index loading entry. + "video_processor": video_processor, + }, + config=config, + source_path=source_dir, + ) + export_dir = tmp_path / f"modular_export_{location_key}" + export_dir.mkdir() + # A retry must replace a stale index rather than retain entries from another + # workflow/source model. + with open(export_dir / "modular_model_index.json", "w") as file: + json.dump({**source_index, "workflow": "stale-workflow"}, file) + + assert is_diffusers_object(pipe) + assert get_diffusion_components(pipe) == pipe.components + + export_hf_checkpoint(pipe, export_dir=export_dir) + + assert (export_dir / "transformer" / "config.json").exists() + assert (export_dir / "scheduler" / "scheduler_config.json").exists() + assert not (export_dir / "model_index.json").exists() + + modular_index = _load_config(export_dir / "modular_model_index.json") + assert modular_index["workflow"] == "text2video" + assert modular_index["_blocks_class_name"] == "FakeTextToVideoBlocks" + assert "transformer_ref" not in modular_index + assert "video_processor" not in modular_index + for component_name in ("transformer", "scheduler"): + loading_spec = modular_index[component_name][2] + assert loading_spec[location_key] == str(export_dir) + assert loading_spec["subfolder"] == component_name + + def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatch): model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant"