diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 10e57cb2f85..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** diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index bab9c1140cb..f6e7d74f211 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]: 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" 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..b025bb359d9 100644 --- a/modelopt/torch/quantization/config.py +++ b/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 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..b49d2c25455 --- /dev/null +++ b/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()) 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..6e6e8ccd3f5 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,105 @@ 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. + + 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 + 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 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((name, module)) + targets.update((id(q), [None, 0]) for q in accepted) + + # 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(): + 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 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: + return + inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs + x = inputs.detach().reshape(-1, inputs.shape[-1]) + if x.shape[0] == 0: + return + 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] + + 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..f5b630d7a88 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -327,3 +327,39 @@ 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 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: + """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/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_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.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/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.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml index 608d476d9eb..6cbd2f50633 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,18 @@ 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. 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/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.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml index 9eea88fe4cf..dfa285ed590 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,18 @@ 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. 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/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/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml index d5bcea814f4..66a8b3ce8cb 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,18 @@ 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. 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/iq2_xxs_experts_only.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_experts_only.yaml new file mode 100644 index 00000000000..88c597d7bc7 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs_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_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 +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/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_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.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/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.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_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.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_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/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/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..4bda575ae85 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 37 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -58,11 +58,16 @@ 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) | -| `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) | +| `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) |
@@ -148,11 +153,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, 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. +- **`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/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..e776fe032a7 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,43 @@ 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 + + +@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 ---------------------------------------- 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..ebb01ef052e --- /dev/null +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -0,0 +1,290 @@ +# 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 _test_utils.torch.transformers_models import get_tiny_llama +from torch import nn + +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 + +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(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))) + + +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) + + +@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, 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): + 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( + {"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 _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): + 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 +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.mlp) + inputs = _inputs(rows=96, columns=HIDDEN) + calls = _record_importance(monkeypatch) + # 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, 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.mlp.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.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)) + ) + 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)) + + +@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