Add imatrix calibration for the GGML formats and use it for fused MoE experts

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 ``<format>_experts_only`` recipes apply imatrix to the routed MoE experts alone.
+
 **Backward Breaking Changes**
 
 **Deprecations**

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]: 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:
             warnings.warn(
                 f"Quantization algorithm: {quant_cfg['algorithm']} does not support setting moe_calib_experts_ratio"

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()
         ):

modelopt/torch/quantization/config.py

@@ -1252,6 +1252,22 @@ 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.
+    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")
+
+
 _ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig
 
 

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 *

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:

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]}

modelopt/torch/quantization/ggml/imatrix.py

@@ -0,0 +1,75 @@
+# 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, validate_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 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]
+        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())

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."""