From d12cd5798f53ebc3b172f945bc290e091f99906e Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 9 Oct 2026 17:26:48 +0000 Subject: [PATCH 01/10] Add imatrix calibration for the GGML formats and use it for fused MoE experts algorithm: {method: imatrix} encodes each weight once, its codebook search weighted by the square root of its inputs' per-column E[x_j^2] -- llama.cpp's importance matrix. After max calibration, one forward with weight fake-quant disabled collects the statistic for every weight whose quantizer backend registered an imatrix encoder; the GGML encoder then encodes the weight with it, writes the decoded weight back and pins the payload, so later forwards and export use those codes. Dense quantized linears and HF fused-MoE experts are both covered, with one statistic per expert over the tokens routed to it. GPTQ does not reach fused experts, so a weight that already holds a pinned payload is left alone and [{method: gptq}, {method: imatrix}] runs GPTQ on the dense linears and imatrix on the experts. The built-in iq2_xxs, iq2_xs and iq2_s recipes now do this; on a dense model the imatrix pass finds nothing left to encode. The square root, as for GPTQ, measured better than the raw statistic on Qwen3.8-27B (MLP) and Qwen3.5-35B-A3B (routed experts): the raw one is heavy-tailed, and noisy for an expert that sees few tokens. On Qwen3.5-35B-A3B IQ2_XXS, wikitext-2 KL divergence from BF16 goes from 0.219 (GPTQ, experts plain) to 0.212 (GPTQ + imatrix experts). The input hooks local_hessian used for dense linears and fused experts are shared as _register_weight_input_hooks, and the floor on the second moments GPTQ applies before its square root is shared as floored_second_moments. hf_ptq's --moe_calib_experts_ratio now applies to every algorithm of a list. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + examples/hf_ptq/example_utils.py | 6 + modelopt/torch/models/nemotron_h/mtp.py | 7 +- modelopt/torch/quantization/config.py | 14 ++ modelopt/torch/quantization/ggml/__init__.py | 4 +- modelopt/torch/quantization/ggml/common.py | 21 ++ modelopt/torch/quantization/ggml/gptq.py | 14 +- modelopt/torch/quantization/ggml/imatrix.py | 70 ++++++ modelopt/torch/quantization/mode.py | 14 ++ modelopt/torch/quantization/model_calib.py | 114 ++++++++-- .../torch/quantization/utils/calib_utils.py | 33 +++ .../configs/ptq/presets/model/iq2_s.yaml | 26 ++- .../configs/ptq/presets/model/iq2_xs.yaml | 26 ++- .../configs/ptq/presets/model/iq2_xxs.yaml | 26 ++- modelopt_recipes/general/ptq/iq2_s.yaml | 5 +- modelopt_recipes/general/ptq/iq2_xs.yaml | 5 +- modelopt_recipes/general/ptq/iq2_xxs.yaml | 5 +- tests/examples/hf_ptq/test_example_utils.py | 19 ++ tests/unit/recipe/test_presets.py | 10 +- .../torch/quantization/test_ggml_imatrix.py | 207 ++++++++++++++++++ 20 files changed, 557 insertions(+), 70 deletions(-) create mode 100644 modelopt/torch/quantization/ggml/imatrix.py create mode 100644 tests/unit/torch/quantization/test_ggml_imatrix.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 10e57cb2f85..aa272ca7bec 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -39,6 +39,7 @@ Changelog - Add ``iq1_m`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 1.75 bits per weight between ``iq1_s`` and ``iq2_xxs``. The same 256-value block constraint applies. - GPTQ (``algorithm: {method: gptq}``) now supports the GGML formats (``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs``, ``iq2_s``, ``q8_0``). Each GGML block is quantized as a whole and its error compensated in the later columns, and export writes exactly the codes GPTQ chose. ``block_size`` must be a multiple of the format's block size (256 for the IQ formats). The built-in ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` recipes now quantize only the MLP/MoE weights and calibrate them with GPTQ, so they need calibration data. - The ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` encoders accept an optional per-input-column ``importance`` that weights each value's squared error in the codebook search, as llama.cpp's imatrix does. GPTQ's new ``importance_weighted`` option passes them the square root of its Hessian diagonal, and the built-in IQ1 and IQ2 recipes enable it, so each block's search favours the columns its outputs depend on most. +- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` recipes now do this. - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it. - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index bab9c1140cb..041fa7f1a17 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -299,6 +299,12 @@ def build_quant_cfg( } elif isinstance(quant_cfg["algorithm"], dict): quant_cfg["algorithm"]["moe_calib_experts_ratio"] = moe_calib_experts_ratio + elif isinstance(quant_cfg["algorithm"], list) and all( + isinstance(algorithm, dict) for algorithm in quant_cfg["algorithm"] + ): + # e.g. [gptq, imatrix]: every calibration pass routes tokens to the same experts. + for algorithm in quant_cfg["algorithm"]: + algorithm["moe_calib_experts_ratio"] = moe_calib_experts_ratio else: warnings.warn( f"Quantization algorithm: {quant_cfg['algorithm']} does not support setting moe_calib_experts_ratio" diff --git a/modelopt/torch/models/nemotron_h/mtp.py b/modelopt/torch/models/nemotron_h/mtp.py index b3a234ac21f..69fd9422521 100644 --- a/modelopt/torch/models/nemotron_h/mtp.py +++ b/modelopt/torch/models/nemotron_h/mtp.py @@ -33,8 +33,8 @@ from transformers.masking_utils import create_causal_mask from modelopt.torch.quantization.model_calib import ( - _LocalHessianInputHook, _needs_activation_forward_for_max_calib, + _WeightInputHook, ) from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.utils import warn_rank_0 @@ -248,14 +248,13 @@ def prepare_for_calibration(full_model) -> bool: @wraps(original_forward) def forward_with_mtp(*args, **kwargs): - # AWQ, GPTQ, and local-Hessian need inputs even with activation quantizers disabled. + # AWQ, GPTQ, local-Hessian and imatrix need inputs with activation quantizers disabled. # Their wrappers/hooks exist only during calibration, unlike retained debug statistics. if not _needs_activation_forward_for_max_calib(mtp) and not any( hasattr(module, "_forward_no_awq") or hasattr(module, GPTQHelper.CACHE_NAME) or any( - isinstance(hook, _LocalHessianInputHook) - for hook in module._forward_pre_hooks.values() + isinstance(hook, _WeightInputHook) for hook in module._forward_pre_hooks.values() ) for module in mtp.modules() ): diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6a538e16d5b..5b00ce2e898 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1252,6 +1252,20 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig): ) +class ImatrixCalibConfig(QuantizeAlgorithmConfig): + """The config for imatrix calibration. + + Encodes each weight once, its codebook search weighted by the square root of its inputs' + per-column second moment ``E[x_j^2]`` -- llama.cpp's importance matrix. Covers the GGML formats + with a weighted search, + on dense linears and on each expert of HF fused-MoE experts, which GPTQ does not reach. A + weight that already holds final codes, such as GPTQ's, is left alone, so + ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest. + """ + + method: Literal["imatrix"] = ModeloptField("imatrix") + + _ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index c4c5da3cba2..b67dc841a21 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -15,9 +15,11 @@ """GGML-compatible block quantization formats.""" -# Importing these installs the TensorQuantizer dispatch entry and the GPTQ helper. +# Importing these installs the TensorQuantizer dispatch entry, the GPTQ helper and the imatrix +# encoder. from . import backend as _backend from . import gptq as _gptq +from . import imatrix as _imatrix from .iq1_m import * from .iq1_m import __all__ as _iq1_m_all from .iq1_s import * diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 5e83b9c79c8..ddedd453983 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -383,6 +383,27 @@ def validate_weight( raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") +# Floor on the per-column second moments, relative to their mean, so a column calibration never +# activated still has some weight in the search. +_IMPORTANCE_FLOOR = 1e-4 + + +def floored_second_moments(second_moments: torch.Tensor | None) -> torch.Tensor | None: + """The per-input-column ``E[x_j^2]``, floored at 1e-4 of its mean. + + ``None`` when there is no statistic or it has no positive mean: a module no calibration token + reached (an idle MoE expert) would make every code score zero error, so it keeps the plain + search. + """ + if second_moments is None: + return None + second_moments = second_moments.float() + mean = second_moments.mean() + if not mean > 0: + return None + return second_moments.clamp_min(_IMPORTANCE_FLOOR * mean) + + def importance_blocks( importance: torch.Tensor | None, weight: torch.Tensor, format_name: str ) -> torch.Tensor | None: diff --git a/modelopt/torch/quantization/ggml/gptq.py b/modelopt/torch/quantization/ggml/gptq.py index 7feece2f95b..763a42b8f99 100644 --- a/modelopt/torch/quantization/ggml/gptq.py +++ b/modelopt/torch/quantization/ggml/gptq.py @@ -20,13 +20,9 @@ import torch from ..utils.calib_utils import GPTQHelper, register_gptq_helper -from .common import pin_packed_weight +from .common import floored_second_moments, pin_packed_weight from .registry import GGML_FORMAT_REGISTRY -# Floor on the Hessian diagonal, relative to its mean, so a column calibration never activated -# still has some weight in the search. -_IMPORTANCE_FLOOR = 1e-4 - __all__ = ["GGMLGPTQHelper", "gptq_group_update"] @@ -99,12 +95,8 @@ def _blockwise_update(self, block_size): assert self.weight is not None, "_blockwise_update called before update_weights()" importance = None if self.importance_weighted and ggml_format.weighted: - diagonal = self.hessian.diagonal().to(self.weight.device, torch.float32) - mean = diagonal.mean() - # A module no calibration token reached (an idle MoE expert) has a zero Hessian, and - # weighting by it would make every code score zero error: keep the plain search. - if mean > 0: - importance = diagonal.clamp_min(_IMPORTANCE_FLOOR * mean).sqrt() + diagonal = floored_second_moments(self.hessian.diagonal().to(self.weight.device)) + importance = None if diagonal is None else diagonal.sqrt() def quantize_group(group, columns): kwargs = {} if importance is None else {"importance": importance[columns]} diff --git a/modelopt/torch/quantization/ggml/imatrix.py b/modelopt/torch/quantization/ggml/imatrix.py new file mode 100644 index 00000000000..599fe19eab8 --- /dev/null +++ b/modelopt/torch/quantization/ggml/imatrix.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""imatrix calibration for GGML formats: one encode per weight, weighted by its inputs' ``E[x_j^2]``.""" + +import torch + +from ..utils.calib_utils import ImatrixEncoder, register_imatrix_encoder +from .common import floored_second_moments, pin_packed_weight, pinned_packed_weight +from .registry import GGML_FORMAT_REGISTRY + +__all__ = ["GGMLImatrixEncoder"] + + +class GGMLImatrixEncoder(ImatrixEncoder): + """Encodes a ``ggml``-backend weight with llama.cpp's imatrix statistic and pins its codes. + + The search is weighted by the square root of ``E[x_j^2]``, as GPTQ's is. Without error + compensation the statistic itself is the exact diagonal of the output error, but it is + heavy-tailed and, for an MoE expert that sees few calibration tokens, noisy; its root did better + on both dense and MoE models. The pinned payload is what later forwards and export use, and it + marks the weight as final: a quantizer that already holds one, such as GPTQ's, is left alone. + """ + + def accepts(self, quantizer) -> bool: + """Enabled formats with a weighted search that do not hold final codes yet.""" + num_bits = getattr(quantizer, "num_bits", None) + if not isinstance(num_bits, str) or not quantizer.is_enabled: + return False + ggml_format = GGML_FORMAT_REGISTRY.get(num_bits) + return ( + ggml_format is not None + and ggml_format.weighted + and pinned_packed_weight(quantizer, num_bits) is None + ) + + def encode(self, weight, quantizer, second_moments): + """Encode ``weight`` once, weighted by ``second_moments``, and pin the payload.""" + ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] + extra_args = quantizer.backend_extra_args or {} + floored = floored_second_moments(second_moments) + kwargs = {} if floored is None else {"importance": floored.sqrt()} + packed, shape = ggml_format.quantize( + weight.contiguous(), + block_chunk_size=extra_args.get("block_chunk_size", ggml_format.block_chunk_size), + **kwargs, + ) + pin_packed_weight(quantizer, quantizer.num_bits, packed) + decoded = ggml_format.dequantize( + packed, + shape, + dtype=torch.float32, + block_chunk_size=extra_args.get("decode_chunk_size", ggml_format.decode_chunk_size), + ) + return decoded.to(weight.dtype) + + +register_imatrix_encoder("ggml", GGMLImatrixEncoder()) diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 1b4fb0dcbe3..47b188ab9fc 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -38,6 +38,7 @@ AWQLiteCalibConfig, CompressConfig, GPTQCalibConfig, + ImatrixCalibConfig, LocalHessianCalibConfig, LSQConfig, MaxCalibConfig, @@ -62,6 +63,7 @@ from .model_calib import ( awq, gptq, + imatrix_calibrate, layerwise_calibrate, local_hessian_calibrate, lsq, @@ -558,6 +560,18 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: _calib_func = gptq +@CalibrateModeRegistry.register_mode +class ImatrixModeDescriptor(BaseCalibrateModeDescriptor): + """Mode for imatrix calibration algorithm.""" + + @property + def config_class(self) -> type[QuantizeAlgorithmConfig]: + """Specifies the config class for the mode.""" + return ImatrixCalibConfig + + _calib_func = imatrix_calibrate + + @CalibrateModeRegistry.register_mode class LSQModeDescriptor(BaseCalibrateModeDescriptor): """Mode for LSQ (Learned Scale Quantization) algorithm.""" diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index ba46d38c153..0e5e2e22a51 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -69,7 +69,7 @@ persistent_materialization, promote_static_block_weight_quantizers, ) -from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper +from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, _IMATRIX_ENCODER_REGISTRY, GPTQHelper __all__ = [ "CalibratorFactory", @@ -941,7 +941,7 @@ def _is_quant_fused_experts(module: nn.Module) -> bool: ) -class _LocalHessianInputHook: +class _WeightInputHook: """Identify active input collectors so auxiliary modules can request a calibration forward.""" def __init__(self, hook: Callable): @@ -951,13 +951,14 @@ def __call__(self, module, args): self.hook(module, args) -def _register_local_hessian_input_hooks(model, names, capture, block_size, warned): +def _register_weight_input_hooks(model, names, capture, on_weight=None): """Register forward hooks feeding each weight's input activations to ``capture``. - Local-Hessian-specific (kept here rather than as a general ``QuantModule`` API): dense - quantized linears hook the layer input; HF fused-MoE experts hook the shared input quantizers, - keyed by the active expert (``_current_expert_idx``). Weights without a hook (conv, - SequentialQuantizer, non-eager experts) fall back to plain MSE. Returns removable handles. + ``capture(weight_quantizer, weight, inputs)`` runs on every forward of a weight whose quantizer + is enabled: dense quantized linears hook the layer input; HF fused-MoE experts hook the shared + input quantizers, keyed by the active expert (``_current_expert_idx``). Weights without a hook + (conv, SequentialQuantizer, non-eager experts) are never captured. ``on_weight(name, weight, + weight_quantizer)``, if given, is called once per hooked weight. Returns removable handles. """ handles: list = [] @@ -980,15 +981,14 @@ def _expert_hook(_input_quantizer, args): weight = getattr(module, "weight", None) if weight is None or weight.dim() != 2 or not module.weight_quantizer.is_enabled: continue - _warn_local_hessian_fallback( - name, weight, module.weight_quantizer, block_size, warned - ) + if on_weight is not None: + on_weight(name, weight, module.weight_quantizer) def _dense_hook(linear, args): if args: capture(linear.weight_quantizer, linear.weight, args[0]) - handles.append(module.register_forward_pre_hook(_LocalHessianInputHook(_dense_hook))) + handles.append(module.register_forward_pre_hook(_WeightInputHook(_dense_hook))) elif _is_quant_fused_experts(module): with enable_weight_access_and_writeback(module, model, names): first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") @@ -1005,9 +1005,8 @@ def _dense_hook(linear, args): input_quantizer = getattr(module, input_q_name, None) if weight is None or quantizers is None or input_quantizer is None: continue - _warn_local_hessian_fallback( - f"{name}.{weight_name}", weight[0], quantizers[0], block_size, warned - ) + if on_weight is not None: + on_weight(f"{name}.{weight_name}", weight[0], quantizers[0]) # Snapshot which experts are enabled now, before the caching forward silences # all weight quantizers — so we don't capture (and discard) disabled experts. enabled = {i for i, q in enumerate(quantizers) if q.is_enabled} @@ -1015,7 +1014,7 @@ def _dense_hook(linear, args): continue handles.append( input_quantizer.register_forward_pre_hook( - _LocalHessianInputHook( + _WeightInputHook( _make_expert_hook(module, weight_name, quantizers, enabled) ) ) @@ -1091,7 +1090,14 @@ def capture(weight_quantizer, weight, input_tensor): # Phase 2: capture each weight's input activations during a forward with weight fake-quant # disabled (so H = ΣXᵀX reflects full-precision weights); input quantizers are left as-is. warned: set = set() - handles = _register_local_hessian_input_hooks(model, names, capture, block_size, warned) + handles = _register_weight_input_hooks( + model, + names, + capture, + on_weight=lambda name, weight, quantizer: _warn_local_hessian_fallback( + name, weight, quantizer, block_size, warned + ), + ) print_rank_0("local_hessian: Caching activations and computing local Hessian...") try: with set_quantizer_by_cfg_context( @@ -2358,6 +2364,82 @@ def _make_gptq_handle(name, m): print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") +@torch.no_grad() +def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None): + """Encode weights once with their inputs' per-column second moments, llama.cpp's imatrix. + + After ``max_calibrate``, one forward with weight fake-quant disabled accumulates + ``E[x_j^2]`` over the inputs of every weight whose quantizer backend registered an imatrix + encoder (see :func:`~modelopt.torch.quantization.utils.calib_utils.register_imatrix_encoder`) + and accepts the quantizer. Each such weight is then encoded once, its search weighted by that + statistic, and replaced by its fake-quantized value. Dense quantized linears and HF fused-MoE + experts (one statistic per expert, over the tokens routed to it) are both covered, so this + reaches the experts GPTQ skips. Encoders skip a quantizer that already holds final codes, so + ``[gptq, imatrix]`` runs GPTQ where it applies and imatrix on the rest. + """ + if forward_loop is None: + warn_rank_0("imatrix: no forward_loop given; skipping imatrix calibration") + return + # The forward below collects what imatrix needs; max only runs one if activations need it. + max_calibrate(model, forward_loop, skip_forward_without_activation_calib=True) + names = module_name_maps(model) + + # Pick the targets before the forward disables every weight quantizer. + def _accepted(quantizer): + backend = getattr(quantizer, "backend", None) + encoder = _IMATRIX_ENCODER_REGISTRY.get(backend) if isinstance(backend, str) else None + return encoder is not None and encoder.accepts(quantizer) + + targets: dict[int, list] = {} # id(weight quantizer) -> [sum of x^2 per column, token count] + target_modules = [] + for module in names.name_to_module.values(): + if not isinstance(module, QuantModule): + continue + accepted = [q for q in module.iter_weight_quantizers_for_calibration() if _accepted(q)] + if accepted: + target_modules.append(module) + targets.update((id(q), [None, 0]) for q in accepted) + if not targets: + print_rank_0("imatrix: no weight quantizer has an imatrix encoder, skipping") + return + + def capture(quantizer, _weight, inputs): + entry = targets.get(id(quantizer)) + if entry is None: + return + inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs + x = inputs.detach().reshape(-1, inputs.shape[-1]).float() + if x.shape[0] == 0: + return + sum_sq = x.square().sum(dim=0) + entry[0] = sum_sq if entry[0] is None else entry[0] + sum_sq + entry[1] += x.shape[0] + + handles = _register_weight_input_hooks(model, names, capture) + try: + with set_quantizer_by_cfg_context( + model, [{"quantizer_name": "*weight_quantizer", "enable": False}] + ): + forward_loop(model) + finally: + for handle in handles: + handle.remove() + + encoded = 0 + for module in target_modules: + with enable_weight_access_and_writeback(module, model, names): + for weight, quantizer in module.iter_weights_for_calibration(): + entry = targets.get(id(quantizer)) + if entry is None: + continue + sum_sq, count = entry + second_moments = None if count == 0 else sum_sq.to(weight.device) / count + encoder = _IMATRIX_ENCODER_REGISTRY[quantizer.backend] + weight.copy_(encoder.encode(weight, quantizer, second_moments)) + encoded += 1 + print_rank_0(f"imatrix: encoded {encoded} weights") + + def _run_weight_scale_calibration(model, forward_loop, scale_algorithm): """Run the weight-scale calibration algorithm. diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index d1736f41272..6a79eaac90d 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -327,3 +327,36 @@ def register_gptq_helper(backend: str, factory: type[GPTQHelper]) -> None: construct ``factory`` instead of the default ``GPTQHelper``. """ _GPTQ_HELPER_REGISTRY[backend] = factory + + +class ImatrixEncoder: + """Encodes a weight once, its search weighted by its inputs' per-column second moments. + + :func:`modelopt.torch.quantization.model_calib.imatrix_calibrate` collects ``E[x_j^2]`` for + every weight whose quantizer an encoder accepts, then calls :meth:`encode` with it. + """ + + def accepts(self, quantizer) -> bool: + """Whether this encoder handles ``quantizer``; decided before calibration runs.""" + raise NotImplementedError + + def encode( + self, weight: torch.Tensor, quantizer, second_moments: torch.Tensor | None + ) -> torch.Tensor: + """Return ``weight`` fake-quantized by ``quantizer``, weighted by ``second_moments``. + + ``second_moments`` is ``None`` when no calibration input reached the weight. + """ + raise NotImplementedError + + +_IMATRIX_ENCODER_REGISTRY: dict[str, ImatrixEncoder] = {} + + +def register_imatrix_encoder(backend: str, encoder: ImatrixEncoder) -> None: + """Register the :class:`ImatrixEncoder` for a quantizer backend. + + :func:`modelopt.torch.quantization.model_calib.imatrix_calibrate` only encodes weights whose + ``weight_quantizer.backend`` has a registered encoder. + """ + _IMATRIX_ENCODER_REGISTRY[backend] = encoder diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml index 608d476d9eb..f2510f5eb22 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizeConfig preset for IQ2_S weight-only quantization of MLP/MoE layers with GPTQ. +# QuantizeConfig preset for IQ2_S weight-only quantization of MLP/MoE layers with GPTQ, +# and imatrix for the fused MoE experts GPTQ does not reach. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -22,15 +23,20 @@ imports: iq2_s: configs/numerics/iq2_s algorithm: - method: gptq - # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. - block_size: 256 - # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. - perc_damp: 0.3 - # Weight each block's codebook search by the square root of the Hessian diagonal. - importance_weighted: true - layerwise: - enable: true + - method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. + perc_damp: 0.3 + # Weight each block's codebook search by the square root of the Hessian diagonal. + importance_weighted: true + layerwise: + enable: true + # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the + # weights GPTQ already encoded. + - method: imatrix + layerwise: + enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml index 9eea88fe4cf..d1e6c7493ca 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizeConfig preset for IQ2_XS weight-only quantization of MLP/MoE layers with GPTQ. +# QuantizeConfig preset for IQ2_XS weight-only quantization of MLP/MoE layers with GPTQ, +# and imatrix for the fused MoE experts GPTQ does not reach. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -22,15 +23,20 @@ imports: iq2_xs: configs/numerics/iq2_xs algorithm: - method: gptq - # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. - block_size: 256 - # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. - perc_damp: 0.3 - # Weight each block's codebook search by the square root of the Hessian diagonal. - importance_weighted: true - layerwise: - enable: true + - method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. + perc_damp: 0.3 + # Weight each block's codebook search by the square root of the Hessian diagonal. + importance_weighted: true + layerwise: + enable: true + # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the + # weights GPTQ already encoded. + - method: imatrix + layerwise: + enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml index d5bcea814f4..f66bcdb6b02 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizeConfig preset for IQ2_XXS weight-only quantization of MLP/MoE layers with GPTQ. +# QuantizeConfig preset for IQ2_XXS weight-only quantization of MLP/MoE layers with GPTQ, +# and imatrix for the fused MoE experts GPTQ does not reach. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -22,15 +23,20 @@ imports: iq2_xxs: configs/numerics/iq2_xxs algorithm: - method: gptq - # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. - block_size: 256 - # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. - perc_damp: 0.3 - # Weight each block's codebook search by the square root of the Hessian diagonal. - importance_weighted: true - layerwise: - enable: true + - method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. + perc_damp: 0.3 + # Weight each block's codebook search by the square root of the Hessian diagonal. + importance_weighted: true + layerwise: + enable: true + # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the + # weights GPTQ already encoded. + - method: imatrix + layerwise: + enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/general/ptq/iq2_s.yaml b/modelopt_recipes/general/ptq/iq2_s.yaml index cbcb314d59b..3623f3c2261 100644 --- a/modelopt_recipes/general/ptq/iq2_s.yaml +++ b/modelopt_recipes/general/ptq/iq2_s.yaml @@ -21,7 +21,8 @@ imports: metadata: description: >- - Applies GGML-compatible IQ2_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ. - Attention and the other linears stay in BF16. Requires calibration data. + Applies GGML-compatible IQ2_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ; + fused MoE experts, which GPTQ does not reach, get an imatrix-weighted encode. Attention and the + other linears stay in BF16. Requires calibration data. quantize: $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml index ef7beb85cb0..4926d9622b3 100644 --- a/modelopt_recipes/general/ptq/iq2_xs.yaml +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -21,7 +21,8 @@ imports: metadata: description: >- - Applies GGML-compatible IQ2_XS weight-only quantization to MLP/MoE layers, calibrated with GPTQ. - Attention and the other linears stay in BF16. Requires calibration data. + Applies GGML-compatible IQ2_XS weight-only quantization to MLP/MoE layers, calibrated with GPTQ; + fused MoE experts, which GPTQ does not reach, get an imatrix-weighted encode. Attention and the + other linears stay in BF16. Requires calibration data. quantize: $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xxs.yaml b/modelopt_recipes/general/ptq/iq2_xxs.yaml index ad6c16705cb..e43854ed0a3 100644 --- a/modelopt_recipes/general/ptq/iq2_xxs.yaml +++ b/modelopt_recipes/general/ptq/iq2_xxs.yaml @@ -21,7 +21,8 @@ imports: metadata: description: >- - Applies GGML-compatible IQ2_XXS weight-only quantization to MLP/MoE layers, calibrated with GPTQ. - Attention and the other linears stay in BF16. Requires calibration data. + Applies GGML-compatible IQ2_XXS weight-only quantization to MLP/MoE layers, calibrated with GPTQ; + fused MoE experts, which GPTQ does not reach, get an imatrix-weighted encode. Attention and the + other linears stay in BF16. Requires calibration data. quantize: $import: preset diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index a6e44a0c363..e25799a5b69 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -115,6 +115,25 @@ def prepare(*args, model_class): # them dequantized to BF16 GptOssExperts (so ModelOpt can quantize/export the experts). +@pytest.mark.parametrize( + ("algorithm", "expected"), + [ + ("max", {"method": "max", "moe_calib_experts_ratio": 0.5}), + ({"method": "gptq"}, {"method": "gptq", "moe_calib_experts_ratio": 0.5}), + ( + [{"method": "gptq"}, {"method": "imatrix"}], + [ + {"method": "gptq", "moe_calib_experts_ratio": 0.5}, + {"method": "imatrix", "moe_calib_experts_ratio": 0.5}, + ], + ), + ], +) +def test_build_quant_cfg_sets_moe_calib_experts_ratio_on_every_algorithm(algorithm, expected): + quant_cfg = example_utils.build_quant_cfg({"quant_cfg": [], "algorithm": algorithm}, None, 0.5) + assert quant_cfg["algorithm"] == expected + + def test_get_original_hf_quant_method_mxfp4_dict(): # gpt-oss layout: quantization_config is a plain dict carrying quant_method. cfg = SimpleNamespace( diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index f97b2d59bd4..62bddcbcfda 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -34,6 +34,7 @@ from modelopt.torch.opt.config_loader import BUILTIN_CONFIG_ROOT from modelopt.torch.quantization.config import LocalHessianCalibConfig, QuantizeConfig from modelopt.torch.quantization.ggml import ( + GGML_FORMAT_REGISTRY, IQ1_M_BLOCK_SIZE, IQ1_M_EFFECTIVE_BITS, IQ1_S_BLOCK_SIZE, @@ -173,8 +174,13 @@ def test_iq_recipe_matches_packing_contract(qformat, block_size, effective_bits) if qformat == "q8_0": assert quantize["algorithm"] is None else: - assert quantize["algorithm"]["method"] == "gptq" - assert quantize["algorithm"]["block_size"] % block_size == 0 + algorithms = quantize["algorithm"] + algorithms = algorithms if isinstance(algorithms, list) else [algorithms] + assert algorithms[0]["method"] == "gptq" + assert algorithms[0]["block_size"] % block_size == 0 + # Formats with a weighted search add imatrix for the fused MoE experts GPTQ does not reach. + expected = ["gptq", "imatrix"] if GGML_FORMAT_REGISTRY[qformat].weighted else ["gptq"] + assert [algorithm["method"] for algorithm in algorithms] == expected # --- RecipeSupersededAction: the flags --recipe replaces ---------------------------------------- diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py new file mode 100644 index 00000000000..e622baae6e7 --- /dev/null +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""imatrix calibration for the GGML formats.""" + +import copy +import dataclasses + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml.common import pinned_packed_weight +from modelopt.torch.quantization.nn import QuantModuleRegistry + +IQ2_XXS = GGML_FORMAT_REGISTRY["iq2_xxs"] +IMATRIX = {"method": "imatrix"} +GPTQ = {"method": "gptq", "block_size": 256, "perc_damp": 0.3} + + +def _config(algorithm, pattern="*weight_quantizer"): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": pattern, + "cfg": {"num_bits": "iq2_xxs", "block_sizes": {-1: 256}, "backend": "ggml"}, + "enable": True, + }, + ], + "algorithm": algorithm, + } + + +def _inputs(rows=1024, columns=512): + # Columns of very different scale, so weighting the search by them matters. + generator = torch.Generator().manual_seed(0) + scale = torch.randn(columns, generator=generator).mul(1.5).exp() + return torch.randn(rows, columns, generator=generator) * scale + + +def _importance(second_moments): + return second_moments.clamp_min(1e-4 * second_moments.mean()).sqrt() + + +def _record_importance(monkeypatch): + calls = [] + + def recording_quantize(weight, **kwargs): + calls.append(kwargs.get("importance")) + return IQ2_XXS.quantize(weight, **kwargs) + + monkeypatch.setitem( + GGML_FORMAT_REGISTRY, "iq2_xxs", dataclasses.replace(IQ2_XXS, quantize=recording_quantize) + ) + return calls + + +def _decoded(quantizer, shape): + pinned = pinned_packed_weight(quantizer, "iq2_xxs") + assert pinned is not None + return IQ2_XXS.dequantize(pinned, torch.tensor(shape), dtype=torch.float32) + + +def test_imatrix_weights_the_search_by_the_root_of_the_inputs_second_moments(monkeypatch): + calls = _record_importance(monkeypatch) + inputs = _inputs() + mtq.quantize(nn.Linear(512, 8, bias=False), _config(IMATRIX), forward_loop=lambda m: m(inputs)) + + torch.testing.assert_close(calls[-1], _importance(inputs.square().mean(dim=0))) + + +def test_imatrix_pins_the_codes_it_chose(): + torch.manual_seed(0) + model = nn.Linear(512, 8, bias=False) + original = model.weight.detach().clone() + plain = copy.deepcopy(model) + inputs = _inputs() + mtq.quantize(model, _config(IMATRIX), forward_loop=lambda m: m(inputs)) + mtq.quantize(plain, _config("max"), forward_loop=lambda m: m(inputs)) + + decoded = _decoded(model.weight_quantizer, model.weight.shape) + assert torch.equal(model.weight.detach(), decoded) + torch.testing.assert_close(model(inputs), inputs @ decoded.T) + + def output_error(weight): + return ((weight - original) @ inputs.T).square().sum() + + assert output_error(decoded) < output_error(plain.weight_quantizer(plain.weight)) + + +def test_imatrix_leaves_codes_gptq_chose(): + torch.manual_seed(0) + gptq_only = nn.Linear(512, 8, bias=False) + both = copy.deepcopy(gptq_only) + inputs = _inputs() + mtq.quantize(gptq_only, _config(GPTQ), forward_loop=lambda m: m(inputs)) + mtq.quantize(both, _config([GPTQ, IMATRIX]), forward_loop=lambda m: m(inputs)) + + assert torch.equal( + pinned_packed_weight(both.weight_quantizer, "iq2_xxs"), + pinned_packed_weight(gptq_only.weight_quantizer, "iq2_xxs"), + ) + assert torch.equal(both.weight, gptq_only.weight) + + +def test_a_module_calibration_never_reached_gets_the_plain_search(): + torch.manual_seed(0) + model = nn.ModuleDict( + {"used": nn.Linear(512, 8, bias=False), "unused": nn.Linear(512, 8, bias=False)} + ) + original = model["unused"].weight.detach().clone() + inputs = _inputs() + mtq.quantize(model, _config(IMATRIX), forward_loop=lambda m: m["used"](inputs)) + + packed, shape = IQ2_XXS.quantize(original) + plain = IQ2_XXS.dequantize(packed, shape, dtype=torch.float32) + assert torch.equal(model["unused"].weight.detach(), plain) + assert torch.equal(_decoded(model["unused"].weight_quantizer, original.shape), plain) + + +NUM_EXPERTS = 4 +HIDDEN = 256 +INTERMEDIATE = 256 + + +class _FusedExperts(nn.Module): + """HF transformers 5 fused-experts layout: one ``F.linear`` pair per routed expert.""" + + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.gate_up_proj = nn.Parameter(torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN) * 0.02) + self.down_proj = nn.Parameter(torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE) * 0.02) + + def forward(self, hidden_states, expert_of_token): + output = torch.zeros_like(hidden_states) + for idx in range(self.num_experts): + tokens = (expert_of_token == idx).nonzero().squeeze(-1) + if len(tokens) == 0: + continue + gate, up = F.linear(hidden_states[tokens], self.gate_up_proj[idx]).chunk(2, dim=-1) + output[tokens] = F.linear(F.silu(gate) * up, self.down_proj[idx]) + return output + + +class _MoE(nn.Module): + """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token.""" + + def __init__(self): + super().__init__() + self.experts = _FusedExperts() + + def forward(self, hidden_states): + return self.experts(hidden_states, torch.arange(len(hidden_states)) % 3) + + +@pytest.fixture +def moe(): + torch.manual_seed(0) + model = _MoE() + yield model + if QuantModuleRegistry.get(_FusedExperts) is not None: + QuantModuleRegistry.unregister(_FusedExperts) + + +def test_imatrix_weights_each_fused_expert_by_its_own_tokens(monkeypatch, moe): + original = copy.deepcopy(moe.experts) + inputs = _inputs(rows=96, columns=HIDDEN) + calls = _record_importance(monkeypatch) + mtq.quantize(moe, _config(IMATRIX, "*experts*weight_quantizer"), lambda m: m(inputs)) + + # One encode per expert and projection: gate_up for experts 0..3, then down for 0..3. + gate_up_calls, down_calls = calls[-8:-4], calls[-4:] + experts = moe.experts + for idx in range(NUM_EXPERTS): + tokens = inputs[idx::3] if idx < 3 else inputs[:0] + if len(tokens) == 0: + # The idle expert keeps the plain search, and its codes are still pinned. + assert gate_up_calls[idx] is None + assert down_calls[idx] is None + else: + gate, up = F.linear(tokens, original.gate_up_proj[idx]).chunk(2, dim=-1) + torch.testing.assert_close(gate_up_calls[idx], _importance(tokens.square().mean(dim=0))) + torch.testing.assert_close( + down_calls[idx], _importance((F.silu(gate) * up).square().mean(dim=0)) + ) + for weight, quantizers in ( + (experts.gate_up_proj, experts.gate_up_proj_weight_quantizers), + (experts.down_proj, experts.down_proj_weight_quantizers), + ): + assert torch.equal(weight[idx].detach(), _decoded(quantizers[idx], weight[idx].shape)) From 8013d53b9e682534d12b0392c35a62d8944c8a35 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sat, 10 Oct 2026 23:12:30 +0000 Subject: [PATCH 02/10] Add an IQ2_XXS experts-only recipe calibrated with imatrix general/ptq/iq2_xxs_experts_only quantizes only the MoE routed experts (the scope of nvfp4_experts_only) to IQ2_XXS and calibrates them with imatrix alone, since GPTQ does not reach fused experts. Shared experts, dense layers and attention stay in BF16. The imatrix test now runs the recipe's quantizers on a toy MoE block with a shared expert, which they leave alone. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../presets/model/iq2_xxs_experts_only.yaml | 38 +++++++++++++++++++ .../general/ptq/iq2_xxs_experts_only.yaml | 28 ++++++++++++++ modelopt_recipes/ptq.md | 18 ++++++--- .../torch/quantization/test_ggml_imatrix.py | 31 ++++++++++++--- 4 files changed, 103 insertions(+), 12 deletions(-) create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_xxs_experts_only.yaml diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml new file mode 100644 index 00000000000..fc339260c89 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ2_XXS weight-only quantization of MoE routed experts only, calibrated +# with imatrix. GPTQ does not reach fused MoE experts, so imatrix weights each expert's search by the +# inputs routed to it. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq2_xxs: configs/numerics/iq2_xxs + +algorithm: + method: imatrix + layerwise: + enable: true +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: iq2_xxs + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: iq2_xxs + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq2_xxs_experts_only.yaml b/modelopt_recipes/general/ptq/iq2_xxs_experts_only.yaml new file mode 100644 index 00000000000..075fa92488b --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xxs_experts_only.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ2_XXS weight-only PTQ of MoE routed experts only, with imatrix. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_xxs_experts_only + +metadata: + description: >- + Applies GGML-compatible IQ2_XXS weight-only quantization to MoE routed experts only, each + expert's search weighted by imatrix statistics of the tokens routed to it. Shared experts, dense + layers and attention stay in BF16. Requires calibration data. +quantize: + $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index d7eb46775ea..81b9504e388 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 32 general/ptq/ recipes (click to expand) +All 33 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -60,9 +60,10 @@ supported combinations. | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | | `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise) | | `iq1_m` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | -| `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | -| `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | -| `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | +| `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | +| `iq2_xxs_experts_only` | IQ2_XXS W2A16 (2.06 bpw), MoE routed experts only | none | imatrix (layerwise) | +| `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | +| `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) |
@@ -148,11 +149,16 @@ activations and tensor-core math are what deliver the throughput. `conv1d` and the vision branch stay in BF16. The formats trade size against accuracy in order: 1.56, 1.75, 2.06, 2.31 and 2.56 bits per weight. Weights are calibrated with layerwise GPTQ (`block_size: 256`, `perc_damp: 0.3`), so calibration - data is required; fused-MoE experts are quantized without the GPTQ update. Quantized - weights must have a final dimension divisible by 256. + data is required. GPTQ does not reach fused-MoE experts: the IQ2 recipes follow it with + imatrix, which weights each expert's search by the tokens routed to it, while the IQ1 + recipes still encode them without calibration. Quantized weights must have a final + dimension divisible by 256. Unified HF export writes the packed GGML blocks; Megatron export additionally requires tensor and pipeline parallel sizes of 1, and does not support fused-MoE experts. +- **`iq2_xxs_experts_only`** — IQ2_XXS on the MoE routed experts only (`*.experts.*`, + `*block_sparse_moe*`), calibrated with imatrix alone; shared experts, dense layers and + attention stay in BF16. - **`q8_0`** — GGML-compatible weight-only quantization on eligible linear layers, with BF16 activations; `lm_head`, MoE routers, `conv1d` and the vision branch stay in BF16. Q8_0 uses 8.5 bits per weight and requires no calibration data. Quantized diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index e622baae6e7..1a61b003022 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -24,6 +24,7 @@ from torch import nn import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.common import pinned_packed_weight from modelopt.torch.quantization.nn import QuantModuleRegistry @@ -159,15 +160,26 @@ def forward(self, hidden_states, expert_of_token): return output -class _MoE(nn.Module): +class _MoEBlock(nn.Module): """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token.""" def __init__(self): super().__init__() self.experts = _FusedExperts() + self.shared_expert = nn.Linear(HIDDEN, HIDDEN, bias=False) def forward(self, hidden_states): - return self.experts(hidden_states, torch.arange(len(hidden_states)) % 3) + routed = self.experts(hidden_states, torch.arange(len(hidden_states)) % 3) + return routed + self.shared_expert(hidden_states) + + +class _MoE(nn.Module): + def __init__(self): + super().__init__() + self.mlp = _MoEBlock() + + def forward(self, hidden_states): + return self.mlp(hidden_states) @pytest.fixture @@ -180,14 +192,21 @@ def moe(): def test_imatrix_weights_each_fused_expert_by_its_own_tokens(monkeypatch, moe): - original = copy.deepcopy(moe.experts) + original = copy.deepcopy(moe.mlp) inputs = _inputs(rows=96, columns=HIDDEN) calls = _record_importance(monkeypatch) - mtq.quantize(moe, _config(IMATRIX, "*experts*weight_quantizer"), lambda m: m(inputs)) + # The shipped experts-only recipe's quantizers, calibrated in one pass (no decoder layers here). + recipe = load_recipe("general/ptq/iq2_xxs_experts_only").quantize.model_dump() + assert recipe["algorithm"]["method"] == "imatrix" + mtq.quantize(moe, {**recipe, "algorithm": IMATRIX}, lambda m: m(inputs)) + + # The shared expert stays in BF16. + assert not moe.mlp.shared_expert.weight_quantizer.is_enabled + assert torch.equal(moe.mlp.shared_expert.weight, original.shared_expert.weight) # One encode per expert and projection: gate_up for experts 0..3, then down for 0..3. gate_up_calls, down_calls = calls[-8:-4], calls[-4:] - experts = moe.experts + experts = moe.mlp.experts for idx in range(NUM_EXPERTS): tokens = inputs[idx::3] if idx < 3 else inputs[:0] if len(tokens) == 0: @@ -195,7 +214,7 @@ def test_imatrix_weights_each_fused_expert_by_its_own_tokens(monkeypatch, moe): assert gate_up_calls[idx] is None assert down_calls[idx] is None else: - gate, up = F.linear(tokens, original.gate_up_proj[idx]).chunk(2, dim=-1) + gate, up = F.linear(tokens, original.experts.gate_up_proj[idx]).chunk(2, dim=-1) torch.testing.assert_close(gate_up_calls[idx], _importance(tokens.square().mean(dim=0))) torch.testing.assert_close( down_calls[idx], _importance((F.silu(gate) * up).square().mean(dim=0)) From a2d5a166676b2ba40bfc699ca10ad25cf3b1e842 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 11 Oct 2026 18:27:32 +0000 Subject: [PATCH 03/10] Test that layerwise imatrix sees the earlier layers' quantization On a tiny Llama, layer 0's MLP statistics match between the single-pass and layerwise modes, and layer 1's differ only when layerwise, since its inputs then pass through layer 0's quantized MLP. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../torch/quantization/test_ggml_imatrix.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index 1a61b003022..8c6591ed3c2 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -21,6 +21,7 @@ import pytest import torch import torch.nn.functional as F +from _test_utils.torch.transformers_models import get_tiny_llama from torch import nn import modelopt.torch.quantization as mtq @@ -224,3 +225,28 @@ def test_imatrix_weights_each_fused_expert_by_its_own_tokens(monkeypatch, moe): (experts.down_proj, experts.down_proj_weight_quantizers), ): assert torch.equal(weight[idx].detach(), _decoded(quantizers[idx], weight[idx].shape)) + + +@pytest.mark.parametrize("layerwise", [False, True]) +def test_layerwise_imatrix_sees_the_earlier_layers_quantization(monkeypatch, layerwise): + calls = _record_importance(monkeypatch) + + def run(layerwise): + model = get_tiny_llama( + dtype=torch.float32, hidden_size=256, intermediate_size=256, num_attention_heads=4 + ) + input_ids = torch.randint(0, 32, (4, 16), generator=torch.Generator().manual_seed(0)) + algorithm = {"method": "imatrix", "layerwise": {"enable": layerwise}} + start = len(calls) + mtq.quantize(model, _config(algorithm, "*mlp*weight_quantizer"), lambda m: m(input_ids)) + weighted = [c for c in calls[start:] if c is not None] + assert len(weighted) == 6 # gate, up and down of both layers + return weighted + + single_pass, other = run(False), run(layerwise) + # Layer 0's MLP inputs come only from unquantized weights either way. + for a, b in zip(single_pass[:3], other[:3]): + torch.testing.assert_close(a, b) + # Layer 1's do not when layerwise: they pass through layer 0's quantized MLP first. + changed = any(not torch.allclose(a, b) for a, b in zip(single_pass[3:], other[3:])) + assert changed == layerwise From 5f6799ae39884e292029c2978b32e8e2774bf015 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 11 Oct 2026 21:29:09 +0000 Subject: [PATCH 04/10] Collect imatrix statistics in a single pass in the IQ2 recipes The recipes' imatrix steps drop layerwise and use the default single forward pass, as llama.cpp collects its imatrix. GPTQ stays layerwise. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/config.py | 2 ++ modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml | 4 +--- modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml | 4 +--- modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml | 4 +--- .../configs/ptq/presets/model/iq2_xxs_experts_only.yaml | 2 -- modelopt_recipes/ptq.md | 8 ++++---- tests/unit/torch/quantization/test_ggml_imatrix.py | 4 ++-- 7 files changed, 11 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 5b00ce2e898..b025bb359d9 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1261,6 +1261,8 @@ class ImatrixCalibConfig(QuantizeAlgorithmConfig): on dense linears and on each expert of HF fused-MoE experts, which GPTQ does not reach. A weight that already holds final codes, such as GPTQ's, is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest. + By default the statistics come from a single forward pass, as llama.cpp collects them; with + ``layerwise`` enabled, each layer's are collected after the earlier layers are quantized. """ method: Literal["imatrix"] = ModeloptField("imatrix") diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml index f2510f5eb22..6cbd2f50633 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml @@ -33,10 +33,8 @@ algorithm: layerwise: enable: true # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the - # weights GPTQ already encoded. + # weights GPTQ already encoded. It collects its statistics in a single forward pass. - method: imatrix - layerwise: - enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml index d1e6c7493ca..dfa285ed590 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -33,10 +33,8 @@ algorithm: layerwise: enable: true # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the - # weights GPTQ already encoded. + # weights GPTQ already encoded. It collects its statistics in a single forward pass. - method: imatrix - layerwise: - enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml index f66bcdb6b02..66a8b3ce8cb 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml @@ -33,10 +33,8 @@ algorithm: layerwise: enable: true # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the - # weights GPTQ already encoded. + # weights GPTQ already encoded. It collects its statistics in a single forward pass. - method: imatrix - layerwise: - enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml index fc339260c89..88c597d7bc7 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml @@ -25,8 +25,6 @@ imports: algorithm: method: imatrix - layerwise: - enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*.experts.*weight_quantizer' diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 81b9504e388..6bdbbf3ce75 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -60,10 +60,10 @@ supported combinations. | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | | `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise) | | `iq1_m` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | -| `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | -| `iq2_xxs_experts_only` | IQ2_XXS W2A16 (2.06 bpw), MoE routed experts only | none | imatrix (layerwise) | -| `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | -| `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ, then imatrix for fused MoE experts (layerwise) | +| `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq2_xxs_experts_only` | IQ2_XXS W2A16 (2.06 bpw), MoE routed experts only | none | imatrix | +| `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) | diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index 8c6591ed3c2..66100b07c0c 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -196,10 +196,10 @@ def test_imatrix_weights_each_fused_expert_by_its_own_tokens(monkeypatch, moe): original = copy.deepcopy(moe.mlp) inputs = _inputs(rows=96, columns=HIDDEN) calls = _record_importance(monkeypatch) - # The shipped experts-only recipe's quantizers, calibrated in one pass (no decoder layers here). + # The shipped experts-only recipe, which calibrates in a single pass. recipe = load_recipe("general/ptq/iq2_xxs_experts_only").quantize.model_dump() assert recipe["algorithm"]["method"] == "imatrix" - mtq.quantize(moe, {**recipe, "algorithm": IMATRIX}, lambda m: m(inputs)) + mtq.quantize(moe, recipe, lambda m: m(inputs)) # The shared expert stays in BF16. assert not moe.mlp.shared_expert.weight_quantizer.is_enabled From 8ebb377755bbba7c13274e30139f03652b2f96e3 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 11 Oct 2026 22:07:26 +0000 Subject: [PATCH 05/10] Add the imatrix step to the IQ1 recipes #2723 gave IQ1_S and IQ1_M a weighted search, so their recipes now follow GPTQ with a single imatrix pass for the fused MoE experts, as the IQ2 recipes do and as the preset test requires. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- .../configs/ptq/presets/model/iq1_m.yaml | 24 +++++++++++-------- .../configs/ptq/presets/model/iq1_s.yaml | 24 +++++++++++-------- modelopt_recipes/general/ptq/iq1_m.yaml | 5 ++-- modelopt_recipes/general/ptq/iq1_s.yaml | 5 ++-- modelopt_recipes/ptq.md | 11 ++++----- 6 files changed, 40 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aa272ca7bec..30583b7ac08 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -39,7 +39,7 @@ Changelog - Add ``iq1_m`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 1.75 bits per weight between ``iq1_s`` and ``iq2_xxs``. The same 256-value block constraint applies. - GPTQ (``algorithm: {method: gptq}``) now supports the GGML formats (``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs``, ``iq2_s``, ``q8_0``). Each GGML block is quantized as a whole and its error compensated in the later columns, and export writes exactly the codes GPTQ chose. ``block_size`` must be a multiple of the format's block size (256 for the IQ formats). The built-in ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` recipes now quantize only the MLP/MoE weights and calibrate them with GPTQ, so they need calibration data. - The ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` encoders accept an optional per-input-column ``importance`` that weights each value's squared error in the codebook search, as llama.cpp's imatrix does. GPTQ's new ``importance_weighted`` option passes them the square root of its Hessian diagonal, and the built-in IQ1 and IQ2 recipes enable it, so each block's search favours the columns its outputs depend on most. -- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` recipes now do this. +- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this. - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it. - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml index 6c965cfef70..7c8da2ced35 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizeConfig preset for IQ1_M weight-only quantization of MLP/MoE layers with GPTQ. +# QuantizeConfig preset for IQ1_M weight-only quantization of MLP/MoE layers with GPTQ, +# and imatrix for the fused MoE experts GPTQ does not reach. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -22,15 +23,18 @@ imports: iq1_m: configs/numerics/iq1_m algorithm: - method: gptq - # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. - block_size: 256 - # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. - perc_damp: 0.3 - # Weight each block's codebook search by the square root of the Hessian diagonal. - importance_weighted: true - layerwise: - enable: true + - method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. + perc_damp: 0.3 + # Weight each block's codebook search by the square root of the Hessian diagonal. + importance_weighted: true + layerwise: + enable: true + # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the + # weights GPTQ already encoded. It collects its statistics in a single forward pass. + - method: imatrix quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml index f8476d290b8..874b6aa29ba 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# QuantizeConfig preset for IQ1_S weight-only quantization of MLP/MoE layers with GPTQ. +# QuantizeConfig preset for IQ1_S weight-only quantization of MLP/MoE layers with GPTQ, +# and imatrix for the fused MoE experts GPTQ does not reach. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -22,15 +23,18 @@ imports: iq1_s: configs/numerics/iq1_s algorithm: - method: gptq - # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. - block_size: 256 - # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. - perc_damp: 0.3 - # Weight each block's codebook search by the square root of the Hessian diagonal. - importance_weighted: true - layerwise: - enable: true + - method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + # Heavier than GPTQ's usual 0.01: low damping overfits the calibration Hessian at these bit widths. + perc_damp: 0.3 + # Weight each block's codebook search by the square root of the Hessian diagonal. + importance_weighted: true + layerwise: + enable: true + # GPTQ only reaches 2-D linears; imatrix encodes the fused MoE experts it leaves, and skips the + # weights GPTQ already encoded. It collects its statistics in a single forward pass. + - method: imatrix quant_cfg: - $import: base_disable_all - quantizer_name: '*mlp*weight_quantizer' diff --git a/modelopt_recipes/general/ptq/iq1_m.yaml b/modelopt_recipes/general/ptq/iq1_m.yaml index 902345854ee..09dd24941a3 100644 --- a/modelopt_recipes/general/ptq/iq1_m.yaml +++ b/modelopt_recipes/general/ptq/iq1_m.yaml @@ -21,7 +21,8 @@ imports: metadata: description: >- - Applies GGML-compatible IQ1_M weight-only quantization to MLP/MoE layers, calibrated with GPTQ. - Attention and the other linears stay in BF16. Requires calibration data. + Applies GGML-compatible IQ1_M weight-only quantization to MLP/MoE layers, calibrated with GPTQ; + fused MoE experts, which GPTQ does not reach, get an imatrix-weighted encode. Attention and the + other linears stay in BF16. Requires calibration data. quantize: $import: preset diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml index 5cd09531323..ac961fefe86 100644 --- a/modelopt_recipes/general/ptq/iq1_s.yaml +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -21,7 +21,8 @@ imports: metadata: description: >- - Applies GGML-compatible IQ1_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ. - Attention and the other linears stay in BF16. Requires calibration data. + Applies GGML-compatible IQ1_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ; + fused MoE experts, which GPTQ does not reach, get an imatrix-weighted encode. Attention and the + other linears stay in BF16. Requires calibration data. quantize: $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 6bdbbf3ce75..ac73623994f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -58,8 +58,8 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | -| `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise) | -| `iq1_m` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise) | +| `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq1_m` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | | `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | | `iq2_xxs_experts_only` | IQ2_XXS W2A16 (2.06 bpw), MoE routed experts only | none | imatrix | | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | @@ -149,10 +149,9 @@ activations and tensor-core math are what deliver the throughput. `conv1d` and the vision branch stay in BF16. The formats trade size against accuracy in order: 1.56, 1.75, 2.06, 2.31 and 2.56 bits per weight. Weights are calibrated with layerwise GPTQ (`block_size: 256`, `perc_damp: 0.3`), so calibration - data is required. GPTQ does not reach fused-MoE experts: the IQ2 recipes follow it with - imatrix, which weights each expert's search by the tokens routed to it, while the IQ1 - recipes still encode them without calibration. Quantized weights must have a final - dimension divisible by 256. + data is required. GPTQ does not reach fused-MoE experts, so imatrix follows it and + weights each expert's search by the tokens routed to it. Quantized weights must have a + final dimension divisible by 256. Unified HF export writes the packed GGML blocks; Megatron export additionally requires tensor and pipeline parallel sizes of 1, and does not support fused-MoE experts. From 7fdbba663425a2043690a29f3f242efb73bfa052 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 12 Oct 2026 06:28:44 +0000 Subject: [PATCH 06/10] Add experts-only recipes for every IQ format iq1_s, iq1_m, iq2_xs and iq2_s get an experts_only recipe like iq2_xxs_experts_only: the format on the MoE routed experts only, calibrated with imatrix alone. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../ptq/presets/model/iq1_m_experts_only.yaml | 36 +++++++++++++++++++ .../ptq/presets/model/iq1_s_experts_only.yaml | 36 +++++++++++++++++++ .../ptq/presets/model/iq2_s_experts_only.yaml | 36 +++++++++++++++++++ .../presets/model/iq2_xs_experts_only.yaml | 36 +++++++++++++++++++ .../general/ptq/iq1_m_experts_only.yaml | 28 +++++++++++++++ .../general/ptq/iq1_s_experts_only.yaml | 28 +++++++++++++++ .../general/ptq/iq2_s_experts_only.yaml | 28 +++++++++++++++ .../general/ptq/iq2_xs_experts_only.yaml | 28 +++++++++++++++ modelopt_recipes/ptq.md | 13 ++++--- tests/unit/recipe/test_presets.py | 30 ++++++++++++++++ 10 files changed, 295 insertions(+), 4 deletions(-) create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml create mode 100644 modelopt_recipes/general/ptq/iq1_m_experts_only.yaml create mode 100644 modelopt_recipes/general/ptq/iq1_s_experts_only.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_s_experts_only.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_xs_experts_only.yaml diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml new file mode 100644 index 00000000000..e63c426add2 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ1_M weight-only quantization of MoE routed experts only, calibrated +# with imatrix. GPTQ does not reach fused MoE experts, so imatrix weights each expert's search by the +# inputs routed to it. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq1_m: configs/numerics/iq1_m + +algorithm: + method: imatrix +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: iq1_m + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: iq1_m + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml new file mode 100644 index 00000000000..d2b5c861421 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ1_S weight-only quantization of MoE routed experts only, calibrated +# with imatrix. GPTQ does not reach fused MoE experts, so imatrix weights each expert's search by the +# inputs routed to it. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq1_s: configs/numerics/iq1_s + +algorithm: + method: imatrix +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: iq1_s + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: iq1_s + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml new file mode 100644 index 00000000000..0140a56f65a --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ2_S weight-only quantization of MoE routed experts only, calibrated +# with imatrix. GPTQ does not reach fused MoE experts, so imatrix weights each expert's search by the +# inputs routed to it. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq2_s: configs/numerics/iq2_s + +algorithm: + method: imatrix +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: iq2_s + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: iq2_s + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml new file mode 100644 index 00000000000..b8718268fc3 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ2_XS weight-only quantization of MoE routed experts only, calibrated +# with imatrix. GPTQ does not reach fused MoE experts, so imatrix weights each expert's search by the +# inputs routed to it. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq2_xs: configs/numerics/iq2_xs + +algorithm: + method: imatrix +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: iq2_xs + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq1_m_experts_only.yaml b/modelopt_recipes/general/ptq/iq1_m_experts_only.yaml new file mode 100644 index 00000000000..f8641838056 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_m_experts_only.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ1_M weight-only PTQ of MoE routed experts only, with imatrix. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq1_m_experts_only + +metadata: + description: >- + Applies GGML-compatible IQ1_M weight-only quantization to MoE routed experts only, each + expert's search weighted by imatrix statistics of the tokens routed to it. Shared experts, dense + layers and attention stay in BF16. Requires calibration data. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq1_s_experts_only.yaml b/modelopt_recipes/general/ptq/iq1_s_experts_only.yaml new file mode 100644 index 00000000000..e121c6a92a6 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_s_experts_only.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ1_S weight-only PTQ of MoE routed experts only, with imatrix. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq1_s_experts_only + +metadata: + description: >- + Applies GGML-compatible IQ1_S weight-only quantization to MoE routed experts only, each + expert's search weighted by imatrix statistics of the tokens routed to it. Shared experts, dense + layers and attention stay in BF16. Requires calibration data. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_s_experts_only.yaml b/modelopt_recipes/general/ptq/iq2_s_experts_only.yaml new file mode 100644 index 00000000000..df8c1bb6e9d --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_s_experts_only.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ2_S weight-only PTQ of MoE routed experts only, with imatrix. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_s_experts_only + +metadata: + description: >- + Applies GGML-compatible IQ2_S weight-only quantization to MoE routed experts only, each + expert's search weighted by imatrix statistics of the tokens routed to it. Shared experts, dense + layers and attention stay in BF16. Requires calibration data. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs_experts_only.yaml b/modelopt_recipes/general/ptq/iq2_xs_experts_only.yaml new file mode 100644 index 00000000000..d1977c1e62d --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xs_experts_only.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ2_XS weight-only PTQ of MoE routed experts only, with imatrix. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_xs_experts_only + +metadata: + description: >- + Applies GGML-compatible IQ2_XS weight-only quantization to MoE routed experts only, each + expert's search weighted by imatrix statistics of the tokens routed to it. Shared experts, dense + layers and attention stay in BF16. Requires calibration data. +quantize: + $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index ac73623994f..4bda575ae85 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 33 general/ptq/ recipes (click to expand) +All 37 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -59,11 +59,15 @@ supported combinations. | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | | `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq1_s_experts_only` | IQ1_S W1A16, MoE routed experts only | none | imatrix | | `iq1_m` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq1_m_experts_only` | IQ1_M W1A16 (1.75 bpw), MoE routed experts only | none | imatrix | | `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | | `iq2_xxs_experts_only` | IQ2_XXS W2A16 (2.06 bpw), MoE routed experts only | none | imatrix | | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq2_xs_experts_only` | IQ2_XS W2A16 (2.31 bpw), MoE routed experts only | none | imatrix | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for fused MoE experts | +| `iq2_s_experts_only` | IQ2_S W2A16 (2.56 bpw), MoE routed experts only | none | imatrix | | `q8_0` | Q8_0 W8A16 (8.5 bpw), eligible linears | none | none (no calibration) |
@@ -155,9 +159,10 @@ activations and tensor-core math are what deliver the throughput. Unified HF export writes the packed GGML blocks; Megatron export additionally requires tensor and pipeline parallel sizes of 1, and does not support fused-MoE experts. -- **`iq2_xxs_experts_only`** — IQ2_XXS on the MoE routed experts only (`*.experts.*`, - `*block_sparse_moe*`), calibrated with imatrix alone; shared experts, dense layers and - attention stay in BF16. +- **`iq1_s_experts_only` / `iq1_m_experts_only` / `iq2_xxs_experts_only` / + `iq2_xs_experts_only` / `iq2_s_experts_only`** — the same IQ formats on the MoE routed + experts only (`*.experts.*`, `*block_sparse_moe*`), calibrated with imatrix alone; shared + experts, dense layers and attention stay in BF16. - **`q8_0`** — GGML-compatible weight-only quantization on eligible linear layers, with BF16 activations; `lm_head`, MoE routers, `conv1d` and the vision branch stay in BF16. Q8_0 uses 8.5 bits per weight and requires no calibration data. Quantized diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 62bddcbcfda..e776fe032a7 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -183,6 +183,36 @@ def test_iq_recipe_matches_packing_contract(qformat, block_size, effective_bits) assert [algorithm["method"] for algorithm in algorithms] == expected +@pytest.mark.parametrize( + ("qformat", "block_size", "effective_bits"), + [ + ("iq1_s", IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS), + ("iq1_m", IQ1_M_BLOCK_SIZE, IQ1_M_EFFECTIVE_BITS), + ("iq2_xxs", IQ2_XXS_BLOCK_SIZE, IQ2_XXS_EFFECTIVE_BITS), + ("iq2_xs", IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS), + ("iq2_s", IQ2_S_BLOCK_SIZE, IQ2_S_EFFECTIVE_BITS), + ], +) +def test_iq_experts_only_recipe_quantizes_the_routed_experts_with_imatrix( + qformat, block_size, effective_bits +): + quantize = load_recipe(f"general/ptq/{qformat}_experts_only").quantize.model_dump( + exclude_unset=True + ) + weight_cfgs = { + entry["quantizer_name"]: entry["cfg"] + for entry in quantize["quant_cfg"] + if isinstance(entry.get("cfg"), dict) and entry["cfg"].get("num_bits") == qformat + } + + assert set(weight_cfgs) == {"*.experts.*weight_quantizer", "*block_sparse_moe*weight_quantizer"} + for weight_cfg in weight_cfgs.values(): + assert weight_cfg["backend"] == "ggml" + assert weight_cfg["block_sizes"][-1] == block_size + assert weight_cfg["effective_bits"] == effective_bits + assert quantize["algorithm"] == {"method": "imatrix"} + + # --- RecipeSupersededAction: the flags --recipe replaces ---------------------------------------- From 690c5a1fa1f5ac664a714f7a2e1178c65032f8cd Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 12 Oct 2026 22:05:53 +0000 Subject: [PATCH 07/10] Check imatrix targets before calibrating and bound the hook's memory - A weight the encoder cannot take, such as an IQ weight whose width is not a multiple of 256, is now refused before the calibration forward and before any weight changes, naming its module. It used to fail only after the forward, with the earlier weights already encoded. - The input hook squares 1024 rows at a time: for a 4096-token batch into a 17408-wide projection, the calibration forward's peak fell from 689 to 204 MiB above resident. - imatrix warns when run on several ranks, since its statistics are not synced across them, as local_hessian does. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/imatrix.py | 7 +++- modelopt/torch/quantization/model_calib.py | 32 ++++++++++++++++--- .../torch/quantization/utils/calib_utils.py | 3 ++ .../torch/quantization/test_ggml_imatrix.py | 27 +++++++++++++++- 4 files changed, 62 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/ggml/imatrix.py b/modelopt/torch/quantization/ggml/imatrix.py index 599fe19eab8..b49d2c25455 100644 --- a/modelopt/torch/quantization/ggml/imatrix.py +++ b/modelopt/torch/quantization/ggml/imatrix.py @@ -18,7 +18,7 @@ import torch from ..utils.calib_utils import ImatrixEncoder, register_imatrix_encoder -from .common import floored_second_moments, pin_packed_weight, pinned_packed_weight +from .common import floored_second_moments, pin_packed_weight, pinned_packed_weight, validate_weight from .registry import GGML_FORMAT_REGISTRY __all__ = ["GGMLImatrixEncoder"] @@ -46,6 +46,11 @@ def accepts(self, quantizer) -> bool: and pinned_packed_weight(quantizer, num_bits) is None ) + def validate(self, weight, quantizer): + """Refuse a weight the format cannot encode, such as one whose width is not a block multiple.""" + ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] + validate_weight(weight, ggml_format.name.upper(), block_size=ggml_format.block_size) + def encode(self, weight, quantizer, second_moments): """Encode ``weight`` once, weighted by ``second_moments``, and pin the payload.""" ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 0e5e2e22a51..0df514098d1 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2364,6 +2364,11 @@ def _make_gptq_handle(name, m): print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") +# Calibration rows squared at a time, bounding imatrix's float32 temporaries to this slice of an +# input rather than twice the whole of it. +_IMATRIX_ROW_CHUNK = 1024 + + @torch.no_grad() def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None): """Encode weights once with their inputs' per-column second moments, llama.cpp's imatrix. @@ -2392,26 +2397,43 @@ def _accepted(quantizer): targets: dict[int, list] = {} # id(weight quantizer) -> [sum of x^2 per column, token count] target_modules = [] - for module in names.name_to_module.values(): + for name, module in names.name_to_module.items(): if not isinstance(module, QuantModule): continue accepted = [q for q in module.iter_weight_quantizers_for_calibration() if _accepted(q)] if accepted: - target_modules.append(module) + target_modules.append((name, module)) targets.update((id(q), [None, 0]) for q in accepted) if not targets: print_rank_0("imatrix: no weight quantizer has an imatrix encoder, skipping") return + # Refuse a weight the encoder cannot take before the forward and before any weight changes; + # otherwise it would fail only after calibration, with the earlier weights already encoded. + for name, module in target_modules: + for weight, quantizer in module.iter_weights_for_calibration(): + if id(quantizer) in targets: + try: + _IMATRIX_ENCODER_REGISTRY[quantizer.backend].validate(weight, quantizer) + except (ValueError, TypeError) as e: + raise type(e)(f"imatrix cannot encode a weight of {name}: {e}") from e + + # TODO: sum the statistics across data-parallel ranks, as max_calibrate syncs amax. + if dist_is_initialized() and dist_size() > 1: + warn_rank_0( + "imatrix: statistics are not synced across ranks, so ranks that calibrate on " + "different data may encode different codes. Treat imatrix as single-rank for now." + ) + def capture(quantizer, _weight, inputs): entry = targets.get(id(quantizer)) if entry is None: return inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs - x = inputs.detach().reshape(-1, inputs.shape[-1]).float() + x = inputs.detach().reshape(-1, inputs.shape[-1]) if x.shape[0] == 0: return - sum_sq = x.square().sum(dim=0) + sum_sq = sum(rows.float().square().sum(dim=0) for rows in x.split(_IMATRIX_ROW_CHUNK)) entry[0] = sum_sq if entry[0] is None else entry[0] + sum_sq entry[1] += x.shape[0] @@ -2426,7 +2448,7 @@ def capture(quantizer, _weight, inputs): handle.remove() encoded = 0 - for module in target_modules: + for _, module in target_modules: with enable_weight_access_and_writeback(module, model, names): for weight, quantizer in module.iter_weights_for_calibration(): entry = targets.get(id(quantizer)) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 6a79eaac90d..f5b630d7a88 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -340,6 +340,9 @@ def accepts(self, quantizer) -> bool: """Whether this encoder handles ``quantizer``; decided before calibration runs.""" raise NotImplementedError + def validate(self, weight: torch.Tensor, quantizer) -> None: + """Raise if :meth:`encode` would refuse ``weight``; checked before calibration runs.""" + def encode( self, weight: torch.Tensor, quantizer, second_moments: torch.Tensor | None ) -> torch.Tensor: diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index 66100b07c0c..4ea77d4630d 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -26,6 +26,7 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe +from modelopt.torch.quantization import model_calib from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.common import pinned_packed_weight from modelopt.torch.quantization.nn import QuantModuleRegistry @@ -81,7 +82,7 @@ def _decoded(quantizer, shape): def test_imatrix_weights_the_search_by_the_root_of_the_inputs_second_moments(monkeypatch): calls = _record_importance(monkeypatch) - inputs = _inputs() + inputs = _inputs(rows=2500) # more rows than the hook squares at a time mtq.quantize(nn.Linear(512, 8, bias=False), _config(IMATRIX), forward_loop=lambda m: m(inputs)) torch.testing.assert_close(calls[-1], _importance(inputs.square().mean(dim=0))) @@ -121,6 +122,30 @@ def test_imatrix_leaves_codes_gptq_chose(): assert torch.equal(both.weight, gptq_only.weight) +def test_imatrix_refuses_a_width_it_cannot_encode_before_calibrating(): + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(512, 640, bias=False), nn.Linear(640, 8, bias=False)) + original = copy.deepcopy(model) + forwards = [] + + with pytest.raises( + ValueError, match=r"imatrix cannot encode a weight of 1: .*divisible by 256" + ): + mtq.quantize(model, _config(IMATRIX), forward_loop=forwards.append) + + assert not forwards + for module, before in zip(model, original): + assert torch.equal(module.weight, before.weight) + + +def test_imatrix_warns_that_ranks_are_not_synced(monkeypatch): + monkeypatch.setattr(model_calib, "dist_is_initialized", lambda: True) + monkeypatch.setattr(model_calib, "dist_size", lambda: 2) + inputs = _inputs() + with pytest.warns(UserWarning, match="not synced across ranks"): + mtq.quantize(nn.Linear(512, 8, bias=False), _config(IMATRIX), lambda m: m(inputs)) + + def test_a_module_calibration_never_reached_gets_the_plain_search(): torch.manual_seed(0) model = nn.ModuleDict( From ad8fe161b81edd11abbeb0115741060724672cce Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 13 Oct 2026 06:22:38 +0000 Subject: [PATCH 08/10] Check imatrix targets before max calibration runs its forward With activation quantizers enabled, max_calibrate runs a forward and sets their amax, and that ran before imatrix refused a weight it cannot encode. Target discovery, the width check and the multi-rank warning now come first; with no targets, max still calibrates the activations. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 15 ++++++++------- .../torch/quantization/test_ggml_imatrix.py | 17 +++++++++++++++-- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 0df514098d1..6e6e8ccd3f5 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2385,8 +2385,6 @@ def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None) if forward_loop is None: warn_rank_0("imatrix: no forward_loop given; skipping imatrix calibration") return - # The forward below collects what imatrix needs; max only runs one if activations need it. - max_calibrate(model, forward_loop, skip_forward_without_activation_calib=True) names = module_name_maps(model) # Pick the targets before the forward disables every weight quantizer. @@ -2404,11 +2402,8 @@ def _accepted(quantizer): if accepted: target_modules.append((name, module)) targets.update((id(q), [None, 0]) for q in accepted) - if not targets: - print_rank_0("imatrix: no weight quantizer has an imatrix encoder, skipping") - return - # Refuse a weight the encoder cannot take before the forward and before any weight changes; + # Refuse a weight the encoder cannot take before any forward or calibration state changes; # otherwise it would fail only after calibration, with the earlier weights already encoded. for name, module in target_modules: for weight, quantizer in module.iter_weights_for_calibration(): @@ -2419,12 +2414,18 @@ def _accepted(quantizer): raise type(e)(f"imatrix cannot encode a weight of {name}: {e}") from e # TODO: sum the statistics across data-parallel ranks, as max_calibrate syncs amax. - if dist_is_initialized() and dist_size() > 1: + if targets and dist_is_initialized() and dist_size() > 1: warn_rank_0( "imatrix: statistics are not synced across ranks, so ranks that calibrate on " "different data may encode different codes. Treat imatrix as single-rank for now." ) + # The forward below collects what imatrix needs; max only runs one if activations need it. + max_calibrate(model, forward_loop, skip_forward_without_activation_calib=True) + if not targets: + print_rank_0("imatrix: no weight quantizer has an imatrix encoder, skipping") + return + def capture(quantizer, _weight, inputs): entry = targets.get(id(quantizer)) if entry is None: diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index 4ea77d4630d..ebb01ef052e 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -122,20 +122,33 @@ def test_imatrix_leaves_codes_gptq_chose(): assert torch.equal(both.weight, gptq_only.weight) -def test_imatrix_refuses_a_width_it_cannot_encode_before_calibrating(): +@pytest.mark.parametrize("calibrate_inputs", [False, True]) +def test_imatrix_refuses_a_width_it_cannot_encode_before_calibrating(calibrate_inputs): torch.manual_seed(0) model = nn.Sequential(nn.Linear(512, 640, bias=False), nn.Linear(640, 8, bias=False)) original = copy.deepcopy(model) + config = _config(IMATRIX) + if calibrate_inputs: + # max_calibrate then runs a forward of its own, which must not run either. + config["quant_cfg"].append( + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": (4, 3)}, "enable": True} + ) + inputs = _inputs(rows=64) forwards = [] + def forward_loop(model): + forwards.append(model) + model(inputs) + with pytest.raises( ValueError, match=r"imatrix cannot encode a weight of 1: .*divisible by 256" ): - mtq.quantize(model, _config(IMATRIX), forward_loop=forwards.append) + mtq.quantize(model, config, forward_loop=forward_loop) assert not forwards for module, before in zip(model, original): assert torch.equal(module.weight, before.weight) + assert module.input_quantizer.amax is None def test_imatrix_warns_that_ranks_are_not_synced(monkeypatch): From 9e33d08de8445a0ad2502ed7a69c3bdbeb3dbcc6 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 13 Oct 2026 06:23:47 +0000 Subject: [PATCH 09/10] Move the imatrix changelog entry to 0.49.0 0.48.0 has been branched, so the entry moves to the new 0.49.0 section, and it now mentions the experts-only recipes. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 30583b7ac08..820c29aecd5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Quantization* + +- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. + **Backward Breaking Changes** **Deprecations** @@ -39,7 +43,6 @@ Changelog - Add ``iq1_m`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 1.75 bits per weight between ``iq1_s`` and ``iq2_xxs``. The same 256-value block constraint applies. - GPTQ (``algorithm: {method: gptq}``) now supports the GGML formats (``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs``, ``iq2_s``, ``q8_0``). Each GGML block is quantized as a whole and its error compensated in the later columns, and export writes exactly the codes GPTQ chose. ``block_size`` must be a multiple of the format's block size (256 for the IQ formats). The built-in ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` recipes now quantize only the MLP/MoE weights and calibrate them with GPTQ, so they need calibration data. - The ``iq1_s``, ``iq1_m``, ``iq2_xxs``, ``iq2_xs`` and ``iq2_s`` encoders accept an optional per-input-column ``importance`` that weights each value's squared error in the codebook search, as llama.cpp's imatrix does. GPTQ's new ``importance_weighted`` option passes them the square root of its Hessian diagonal, and the built-in IQ1 and IQ2 recipes enable it, so each block's search favours the columns its outputs depend on most. -- Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this. - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it. - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. From 7a53b58f60c854b0e6f435c4f03a33ce81d98ea7 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 15 Oct 2026 17:04:24 +0000 Subject: [PATCH 10/10] Say why build_quant_cfg sets the MoE ratio on every algorithm The comment claimed every calibration pass routes tokens to the same experts, which the passes neither do nor need. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- examples/hf_ptq/example_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 041fa7f1a17..f6e7d74f211 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -302,7 +302,7 @@ def build_quant_cfg( elif isinstance(quant_cfg["algorithm"], list) and all( isinstance(algorithm, dict) for algorithm in quant_cfg["algorithm"] ): - # e.g. [gptq, imatrix]: every calibration pass routes tokens to the same experts. + # e.g. [gptq, imatrix]: each algorithm runs its own calibration, so each gets the ratio. for algorithm in quant_cfg["algorithm"]: algorithm["moe_calib_experts_ratio"] = moe_calib_experts_ratio else: