From 0c0508b33caa8c8ee57c87de0aa3c1dd98d422e1 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 16 Oct 2026 18:02:42 +0000 Subject: [PATCH 1/3] Run GPTQ on HF fused MoE experts GPTQ only collected Hessians for modules with a 2-D weight and one weight quantizer, so HF fused MoE experts (one 3-D parameter, a quantizer per expert) were skipped and left to round-to-nearest or imatrix. gptq() now builds one helper per expert and projection. Each Hessian comes from a forward hook on the projection's shared input quantizer, keyed by the expert currently running, so it holds exactly the (input-quantized) tokens routed to that expert. GPTQHelper takes an optional weight_quantizer and weight for that, and writes its result back in place, so an expert's slice is updated inside the 3-D parameter. An expert no calibration token reaches gets the plain search, as a dense layer calibration never reaches does. The IQ recipes, docs and changelog no longer say GPTQ skips the experts. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 3 +- modelopt/torch/quantization/config.py | 6 +- modelopt/torch/quantization/ggml/gptq.py | 8 +- modelopt/torch/quantization/model_calib.py | 130 ++++++++++++++---- .../torch/quantization/utils/calib_utils.py | 70 +++++++--- .../configs/ptq/presets/model/iq1_m.yaml | 6 +- .../ptq/presets/model/iq1_m_experts_only.yaml | 4 +- .../configs/ptq/presets/model/iq1_s.yaml | 6 +- .../ptq/presets/model/iq1_s_experts_only.yaml | 4 +- .../configs/ptq/presets/model/iq2_s.yaml | 6 +- .../ptq/presets/model/iq2_s_experts_only.yaml | 4 +- .../configs/ptq/presets/model/iq2_xs.yaml | 6 +- .../presets/model/iq2_xs_experts_only.yaml | 4 +- .../configs/ptq/presets/model/iq2_xxs.yaml | 6 +- .../presets/model/iq2_xxs_experts_only.yaml | 4 +- modelopt_recipes/general/ptq/iq1_m.yaml | 6 +- modelopt_recipes/general/ptq/iq1_s.yaml | 6 +- modelopt_recipes/general/ptq/iq2_s.yaml | 6 +- modelopt_recipes/general/ptq/iq2_xs.yaml | 6 +- modelopt_recipes/general/ptq/iq2_xxs.yaml | 6 +- modelopt_recipes/ptq.md | 14 +- .../_test_utils/torch/quantization/models.py | 51 +++++++ tests/unit/torch/quantization/conftest.py | 30 ++++ .../unit/torch/quantization/test_ggml_gptq.py | 26 ++++ .../torch/quantization/test_ggml_imatrix.py | 59 +------- 25 files changed, 313 insertions(+), 164 deletions(-) create mode 100644 tests/unit/torch/quantization/conftest.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f20a33f1c17..c248590f258 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,7 +9,8 @@ Changelog *Quantization* - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. -- 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. +- 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 where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. +- GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. *Speculative Decoding* diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index b025bb359d9..a43e031f135 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1258,9 +1258,9 @@ class ImatrixCalibConfig(QuantizeAlgorithmConfig): 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. + on dense linears and on each expert of HF fused-MoE experts. 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 any weight it leaves. 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. """ diff --git a/modelopt/torch/quantization/ggml/gptq.py b/modelopt/torch/quantization/ggml/gptq.py index 763a42b8f99..92c56ed5003 100644 --- a/modelopt/torch/quantization/ggml/gptq.py +++ b/modelopt/torch/quantization/ggml/gptq.py @@ -78,15 +78,15 @@ class GGMLGPTQHelper(GPTQHelper): those exact codes rather than encoding the GPTQ'd weight again, which would not return them. """ - def update_weights(self, block_size, perc_damp): + def update_weights(self, block_size, perc_damp, weight=None): """Run the GPTQ update, then pin the chosen payload to the weight quantizer.""" - super().update_weights(block_size, perc_damp) - quantizer = self.module.weight_quantizer + super().update_weights(block_size, perc_damp, weight) + quantizer = self.weight_quantizer pin_packed_weight(quantizer, quantizer.num_bits, self._packed) self._packed = None def _blockwise_update(self, block_size): - quantizer = self.module.weight_quantizer + quantizer = self.weight_quantizer ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] extra_args = quantizer.backend_extra_args or {} block_chunk_size = extra_args.get("block_chunk_size", ggml_format.block_chunk_size) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 2e5a2d65393..bd5dadb0dc0 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -941,6 +941,29 @@ def _is_quant_fused_experts(module: nn.Module) -> bool: ) +def _fused_expert_projections(module: nn.Module): + """Yield ``(weight_name, weight, weight_quantizers, input_quantizer)`` per fused-experts projection. + + ``weight`` is the 3-D ``(num_experts, out, in)`` parameter, ``weight_quantizers`` holds one + quantizer per expert, and ``input_quantizer`` is shared by the experts and runs once per expert + forward with ``module._current_expert_idx`` already set to that expert. + """ + first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") + for weight_name, quantizers_name, input_q_name in ( + ( + first_proj_attr, + f"{first_proj_attr}_weight_quantizers", + f"{first_proj_attr}_input_quantizer", + ), + ("down_proj", "down_proj_weight_quantizers", "down_proj_input_quantizer"), + ): + weight = getattr(module, weight_name, None) + quantizers = getattr(module, quantizers_name, None) + input_quantizer = getattr(module, input_q_name, None) + if weight is not None and quantizers is not None and input_quantizer is not None: + yield weight_name, weight, quantizers, input_quantizer + + class _WeightInputHook: """Identify active input collectors so auxiliary modules can request a calibration forward.""" @@ -991,20 +1014,9 @@ def _dense_hook(linear, args): 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") - for weight_name, quantizers_name, input_q_name in ( - ( - first_proj_attr, - f"{first_proj_attr}_weight_quantizers", - f"{first_proj_attr}_input_quantizer", - ), - ("down_proj", "down_proj_weight_quantizers", "down_proj_input_quantizer"), + for weight_name, weight, quantizers, input_quantizer in _fused_expert_projections( + module ): - weight = getattr(module, weight_name, None) - quantizers = getattr(module, quantizers_name, None) - input_quantizer = getattr(module, input_q_name, None) - if weight is None or quantizers is None or input_quantizer is None: - continue 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 @@ -2299,9 +2311,14 @@ def gptq( 1. ``max_calibrate`` to set amax values from the current activations. 2. Promote eligible quantizers to ``StaticBlockScaleQuantizer`` (two-level scaling). - 3. Collect per-linear-layer Hessian matrices via forward hooks. + 3. Collect Hessian matrices via forward hooks: one per quantized linear layer, and one per + expert and projection of HF fused-MoE experts, over the tokens routed to that expert. 4. Blockwise weight updates using the inverse Hessian to compensate for - rounding error (the core GPTQ column-wise update). + rounding error (the core GPTQ column-wise update). An expert no calibration token + reached has a zero Hessian and gets the plain search, as a dense layer would. + + Each expert Hessian is ``in_features x in_features`` float32, so a large MoE needs + layerwise calibration to hold only one layer's expert Hessians at a time. Args: model: The module to quantize — either the full model or a single decoder @@ -2323,42 +2340,95 @@ def gptq( for n, m in model.named_modules() if is_quantized_linear(m) and m.weight_quantizer.is_enabled ] - if not quantized_layers: - print_rank_0("No quantized linear layers found, skipping GPTQ") + fused_experts = [(n, m) for n, m in model.named_modules() if _is_quant_fused_experts(m)] + if not quantized_layers and not fused_experts: + print_rank_0("No quantized linear layers or fused experts found, skipping GPTQ") return - def _make_gptq_handle(name, m): - backend = getattr(m.weight_quantizer, "backend", None) + def _make_gptq_handle(name, m, weight_quantizer=None, weight=None): + quantizer = m.weight_quantizer if weight_quantizer is None else weight_quantizer + backend = getattr(quantizer, "backend", None) if backend is None: cls = GPTQHelper else: cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) return cls( - m, name, offload_to_cpu=True, fused=fused, importance_weighted=importance_weighted + m, + name, + offload_to_cpu=True, + fused=fused, + importance_weighted=importance_weighted, + weight_quantizer=weight_quantizer, + weight=weight, ) gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} for handle in gptq_handles.values(): handle.setup() - print_rank_0(f"Computing Hessians for {len(gptq_handles)} linear layers...") + # HF fused-MoE experts: one helper per expert and projection, fed only the tokens routed to that + # expert. The projection's shared input quantizer runs once per expert forward, so its output is + # exactly what the expert's F.linear consumes, the input-quantized input dense GPTQ also uses. + names = module_name_maps(model) + expert_handles: dict[int, GPTQHelper] = {} # id(expert weight quantizer) -> helper + hooks = [] - with set_quantizer_by_cfg_context( - model, [{"quantizer_name": "*weight_quantizer", "enable": False}] - ): - forward_loop(model) + def _make_expert_hook(experts, per_expert): + def _expert_hook(_input_quantizer, _args, output): + handle = per_expert.get(experts._current_expert_idx) + if handle is not None: + handle.accumulate(output) + + return _expert_hook + + for name, module in fused_experts: + with enable_weight_access_and_writeback(module, model, names): + for weight_name, weight, quantizers, input_quantizer in _fused_expert_projections( + module + ): + per_expert = { + idx: _make_gptq_handle(f"{name}.{weight_name}.{idx}", module, q, weight[idx]) + for idx, q in enumerate(quantizers) + if q.is_enabled + } + if per_expert: + expert_handles.update((id(quantizers[i]), h) for i, h in per_expert.items()) + hooks.append( + input_quantizer.register_forward_hook(_make_expert_hook(module, per_expert)) + ) + + print_rank_0( + f"Computing Hessians for {len(gptq_handles)} linear layers and " + f"{len(expert_handles)} fused-expert weights..." + ) + + try: + with set_quantizer_by_cfg_context( + model, [{"quantizer_name": "*weight_quantizer", "enable": False}] + ): + forward_loop(model) + finally: + for hook in hooks: + hook.remove() for handle in gptq_handles.values(): handle.cleanup() + for handle in [*gptq_handles.values(), *expert_handles.values()]: handle.sync_hessian_across_dp() print_rank_0("Updating weights using GPTQ algorithm...") - names = module_name_maps(model) for handle in gptq_handles.values(): with enable_weight_access_and_writeback(handle.module, model, names): handle.update_weights(block_size, perc_damp) handle.free() - del gptq_handles + for _, module in fused_experts: + with enable_weight_access_and_writeback(module, model, names): + for weight, quantizer in module.iter_weights_for_calibration(): + handle = expert_handles.pop(id(quantizer), None) + if handle is not None: + handle.update_weights(block_size, perc_damp, weight) + handle.free() + del gptq_handles, expert_handles if torch.cuda.is_available(): torch.cuda.empty_cache() @@ -2379,9 +2449,9 @@ def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None) 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. + experts (one statistic per expert, over the tokens routed to it) are both covered. Encoders + skip a quantizer that already holds final codes, so ``[gptq, imatrix]`` runs GPTQ where it + applies and imatrix on any weight it leaves. """ if forward_loop is None: warn_rank_0("imatrix: no forward_loop given; skipping imatrix calibration") diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 8e828d87ca5..c0c7b557fc8 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -130,14 +130,30 @@ class GPTQHelper: CACHE_NAME = "_forward_no_gptq_hessian" - def __init__(self, module, name, offload_to_cpu=False, fused=False, importance_weighted=False): - """Initialize GPTQHelper with module state and Hessian storage.""" + def __init__( + self, + module, + name, + offload_to_cpu=False, + fused=False, + importance_weighted=False, + weight_quantizer=None, + weight=None, + ): + """Initialize GPTQHelper with module state and Hessian storage. + + ``weight_quantizer`` and ``weight`` select one of several weights a module holds, such as + one expert's slice of an HF fused-MoE parameter; by default the helper owns + ``module.weight`` and ``module.weight_quantizer``. + """ self.module = module self.name = name self.fused = fused self.importance_weighted = importance_weighted - in_features = module.weight.shape[-1] - device = module.weight.device + self._weight_quantizer = weight_quantizer + weight = module.weight if weight is None else weight + in_features = weight.shape[-1] + device = weight.device if device.type == "meta" or ( offload_to_cpu and device.type == "cuda" and get_used_gpu_mem_fraction(device) > 0.65 ): @@ -148,6 +164,13 @@ def __init__(self, module, name, offload_to_cpu=False, fused=False, importance_w self.weight: torch.Tensor | None = None self.h_inv: torch.Tensor | None = None + @property + def weight_quantizer(self): + """The quantizer GPTQ quantizes this helper's weight with.""" + if self._weight_quantizer is not None: + return self._weight_quantizer + return self.module.weight_quantizer + def setup(self): """Patch the module's forward to accumulate Hessian during the collection pass.""" gptq_helper = self @@ -158,9 +181,7 @@ def hessian_forward(self, input, *args, **kwargs): hessian_input = self.input_quantizer(inp) else: hessian_input = inp - gptq_helper.hessian, gptq_helper.n_samples = update_hessian( - hessian_input, gptq_helper.hessian, gptq_helper.n_samples - ) + gptq_helper.accumulate(hessian_input) out = self._forward_no_gptq_hessian(input, *args, **kwargs) @@ -168,6 +189,11 @@ def hessian_forward(self, input, *args, **kwargs): bind_forward_method(self.module, hessian_forward, self.CACHE_NAME) + def accumulate(self, inputs): + """Add a batch of this weight's inputs, after input quantization, to the Hessian.""" + inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs + self.hessian, self.n_samples = update_hessian(inputs, self.hessian, self.n_samples) + def cleanup(self): """Unpatch the module's forward method.""" unpatch_forward_method(self.module, self.CACHE_NAME) @@ -197,23 +223,23 @@ def free(self): self.weight = None self.h_inv = None - def update_weights(self, block_size, perc_damp): - """Run GPTQ blockwise weight update on this module. + def update_weights(self, block_size, perc_damp, weight=None): + """Run GPTQ blockwise weight update on this helper's weight. - Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, - logs MSE, and writes the result back to the module. + Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and + overwrites ``weight`` (default ``module.weight``) in place, so an expert's slice of a + fused 3-D parameter is written back into that parameter. """ - hessian = self.hessian.to(self.module.weight.device) - self.weight = self.module.weight.data.float().clone() + weight = self.module.weight if weight is None else weight + hessian = self.hessian.to(weight.device) + self.weight = weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) self._blockwise_update(block_size) - self._print_mse_error(hessian) - self.module.weight.data = self.weight.reshape(self.module.weight.shape).to( - self.module.weight.data.dtype - ) + self._print_mse_error(hessian, weight) + weight.data.copy_(self.weight.reshape(weight.shape)) # ------------------------------------------------------------------ - # Quantize helpers — all read from self.module, self.weight, self.h_inv + # Quantize helpers — all read from self.weight_quantizer, self.weight, self.h_inv # ------------------------------------------------------------------ def _prepare_hessian_inverse(self, hessian, perc_damp): @@ -232,7 +258,7 @@ def _blockwise_update(self, block_size): assert self.weight is not None and self.h_inv is not None, ( "_blockwise_update called before _prepare_hessian_inverse()" ) - quantizer = self.module.weight_quantizer + quantizer = self.weight_quantizer if self.fused and getattr(quantizer, "_is_nvfp4_static_quantizer", False): block_sizes = quantizer.block_sizes @@ -252,9 +278,9 @@ def _blockwise_update(self, block_size): else: gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) - def _print_mse_error(self, hessian): - """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" - w_orig = self.module.weight.float() + def _print_mse_error(self, hessian, weight): + """Log Hessian-weighted relative MSE between ``self.weight`` and original ``weight``.""" + w_orig = weight.float() delta = self.weight - w_orig mse = (delta).mm(hessian).mul(delta).mean() / (w_orig.mm(hessian).mul(w_orig).mean() + 1e-6) suffix = f", n_hessian_samples: {self.n_samples}" if self.n_samples else "" diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml index 7c8da2ced35..85858d988c1 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml @@ -14,7 +14,7 @@ # limitations under the License. # 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. +# and imatrix for any weight GPTQ leaves. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -32,8 +32,8 @@ algorithm: 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. + # imatrix encodes any weight GPTQ 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 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 index e63c426add2..78a1ac85a7e 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_m_experts_only.yaml @@ -14,8 +14,8 @@ # 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. +# with imatrix, which weights each expert's search by the inputs routed to it, collected in a +# single forward pass. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml index 874b6aa29ba..10ab2adad5e 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml @@ -14,7 +14,7 @@ # limitations under the License. # 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. +# and imatrix for any weight GPTQ leaves. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -32,8 +32,8 @@ algorithm: 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. + # imatrix encodes any weight GPTQ 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 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 index d2b5c861421..126be6b447c 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s_experts_only.yaml @@ -14,8 +14,8 @@ # 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. +# with imatrix, which weights each expert's search by the inputs routed to it, collected in a +# single forward pass. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml index 6cbd2f50633..fc8ac9a7e7d 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml @@ -14,7 +14,7 @@ # limitations under the License. # 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. +# and imatrix for any weight GPTQ leaves. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -32,8 +32,8 @@ algorithm: 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. + # imatrix encodes any weight GPTQ 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 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 index 0140a56f65a..4602236b083 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s_experts_only.yaml @@ -14,8 +14,8 @@ # 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. +# with imatrix, which weights each expert's search by the inputs routed to it, collected in a +# single forward pass. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml index dfa285ed590..b9fd20b478c 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -14,7 +14,7 @@ # limitations under the License. # 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. +# and imatrix for any weight GPTQ leaves. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -32,8 +32,8 @@ algorithm: 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. + # imatrix encodes any weight GPTQ 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 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 index b8718268fc3..9659c0e2dc1 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs_experts_only.yaml @@ -14,8 +14,8 @@ # 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. +# with imatrix, which weights each expert's search by the inputs routed to it, collected in a +# single forward pass. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml index 66a8b3ce8cb..a54c292fa8b 100644 --- a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml @@ -14,7 +14,7 @@ # limitations under the License. # 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. +# and imatrix for any weight GPTQ leaves. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: @@ -32,8 +32,8 @@ algorithm: 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. + # imatrix encodes any weight GPTQ 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 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 88c597d7bc7..4f401bfd431 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 @@ -14,8 +14,8 @@ # 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. +# with imatrix, which weights each expert's search by the inputs routed to it, collected in a +# single forward pass. # modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig imports: diff --git a/modelopt_recipes/general/ptq/iq1_m.yaml b/modelopt_recipes/general/ptq/iq1_m.yaml index 09dd24941a3..39ca54e0c3e 100644 --- a/modelopt_recipes/general/ptq/iq1_m.yaml +++ b/modelopt_recipes/general/ptq/iq1_m.yaml @@ -21,8 +21,8 @@ imports: metadata: description: >- - 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. + Applies GGML-compatible IQ1_M weight-only quantization to MLP/MoE layers, calibrated with GPTQ, + including each fused MoE expert; any weight GPTQ leaves gets 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 ac961fefe86..d50ba1879c0 100644 --- a/modelopt_recipes/general/ptq/iq1_s.yaml +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -21,8 +21,8 @@ imports: metadata: description: >- - 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. + Applies GGML-compatible IQ1_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ, + including each fused MoE expert; any weight GPTQ leaves gets 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.yaml b/modelopt_recipes/general/ptq/iq2_s.yaml index 3623f3c2261..5abb8d4206e 100644 --- a/modelopt_recipes/general/ptq/iq2_s.yaml +++ b/modelopt_recipes/general/ptq/iq2_s.yaml @@ -21,8 +21,8 @@ imports: metadata: description: >- - 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. + Applies GGML-compatible IQ2_S weight-only quantization to MLP/MoE layers, calibrated with GPTQ, + including each fused MoE expert; any weight GPTQ leaves gets 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 4926d9622b3..ae5d7d9d476 100644 --- a/modelopt_recipes/general/ptq/iq2_xs.yaml +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -21,8 +21,8 @@ imports: metadata: description: >- - 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. + Applies GGML-compatible IQ2_XS weight-only quantization to MLP/MoE layers, calibrated with GPTQ, + including each fused MoE expert; any weight GPTQ leaves gets 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 e43854ed0a3..c3195c60565 100644 --- a/modelopt_recipes/general/ptq/iq2_xxs.yaml +++ b/modelopt_recipes/general/ptq/iq2_xxs.yaml @@ -21,8 +21,8 @@ imports: metadata: description: >- - 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. + Applies GGML-compatible IQ2_XXS weight-only quantization to MLP/MoE layers, calibrated with GPTQ, + including each fused MoE expert; any weight GPTQ leaves gets 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 ae010dc79fc..552f4aa8a8b 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -58,15 +58,15 @@ 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), then imatrix for fused MoE experts | +| `iq1_s` | IQ1_S W1A16, MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for weights GPTQ left | | `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` | IQ1_M W1A16 (1.75 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for weights GPTQ left | | `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` | IQ2_XXS W2A16 (2.06 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for weights GPTQ left | | `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` | IQ2_XS W2A16 (2.31 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for weights GPTQ left | | `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` | IQ2_S W2A16 (2.56 bpw), MLP + MoE weights only | none | GPTQ (layerwise), then imatrix for weights GPTQ left | | `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) | @@ -153,8 +153,8 @@ 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, so imatrix follows it and - weights each expert's search by the tokens routed to it. Quantized weights must have a + data is required. GPTQ covers fused-MoE experts too, one Hessian per expert over the tokens + routed to it; imatrix follows it for any weight GPTQ leaves. 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 diff --git a/tests/_test_utils/torch/quantization/models.py b/tests/_test_utils/torch/quantization/models.py index f0c78bb1c19..2e1875337df 100644 --- a/tests/_test_utils/torch/quantization/models.py +++ b/tests/_test_utils/torch/quantization/models.py @@ -164,3 +164,54 @@ def forward(self, x): def get_input(self, batch_size=2): return torch.randn(batch_size, self.in_features) + + +TOY_MOE_EXPERTS = 4 +TOY_MOE_HIDDEN = 256 +TOY_MOE_INTERMEDIATE = 256 + + +class ToyFusedExperts(nn.Module): + """HF transformers 5 fused-experts layout: one ``F.linear`` pair per routed expert.""" + + def __init__(self): + super().__init__() + self.num_experts = TOY_MOE_EXPERTS + self.gate_up_proj = nn.Parameter( + torch.randn(TOY_MOE_EXPERTS, 2 * TOY_MOE_INTERMEDIATE, TOY_MOE_HIDDEN) * 0.02 + ) + self.down_proj = nn.Parameter( + torch.randn(TOY_MOE_EXPERTS, TOY_MOE_HIDDEN, TOY_MOE_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 ToyMoEBlock(nn.Module): + """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token.""" + + def __init__(self): + super().__init__() + self.experts = ToyFusedExperts() + self.shared_expert = nn.Linear(TOY_MOE_HIDDEN, TOY_MOE_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 ToyMoE(nn.Module): + def __init__(self): + super().__init__() + self.mlp = ToyMoEBlock() + + def forward(self, hidden_states): + return self.mlp(hidden_states) diff --git a/tests/unit/torch/quantization/conftest.py b/tests/unit/torch/quantization/conftest.py new file mode 100644 index 00000000000..68b9f1d0f44 --- /dev/null +++ b/tests/unit/torch/quantization/conftest.py @@ -0,0 +1,30 @@ +# 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. + +import pytest +import torch +from _test_utils.torch.quantization.models import ToyFusedExperts, ToyMoE + +from modelopt.torch.quantization.nn import QuantModuleRegistry + + +@pytest.fixture +def moe(): + """A one-layer toy MoE with HF fused experts; unregisters its quantized expert class after.""" + torch.manual_seed(0) + model = ToyMoE() + yield model + if QuantModuleRegistry.get(ToyFusedExperts) is not None: + QuantModuleRegistry.unregister(ToyFusedExperts) diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index bc14dd51fdf..ecbc648b235 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -21,6 +21,8 @@ import pytest import torch +import torch.nn.functional as F +from _test_utils.torch.quantization.models import TOY_MOE_EXPERTS, TOY_MOE_HIDDEN import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -205,6 +207,30 @@ def test_a_module_calibration_never_reached_gets_the_plain_search(): torch.testing.assert_close(model["unused"].weight.detach(), plain, rtol=0, atol=0) +def test_gptq_quantizes_each_fused_expert_as_a_linear_fed_only_its_tokens(moe): + # The expert no token reaches is compared against a linear calibration never reaches. + experts = moe.mlp.experts + original = copy.deepcopy(experts) + inputs = _inputs()[:96, :TOY_MOE_HIDDEN] + config = _iq_config(GPTQ, "iq2_xxs") + mtq.quantize(moe, config, forward_loop=lambda m: m(inputs)) + + for idx in range(TOY_MOE_EXPERTS): + tokens = inputs[idx::3] if idx < 3 else inputs[:0] + gate, up = F.linear(tokens, original.gate_up_proj[idx]).chunk(2, dim=-1) + for name, expert_inputs in (("gate_up_proj", tokens), ("down_proj", F.silu(gate) * up)): + reference = torch.nn.Linear(*getattr(original, name)[idx].shape[::-1], bias=False) + reference.weight.data.copy_(getattr(original, name)[idx]) + mtq.quantize(reference, config, forward_loop=lambda m, x=expert_inputs: m(x)) + + assert torch.equal(getattr(experts, name)[idx], reference.weight) + quantizer = getattr(experts, f"{name}_weight_quantizers")[idx] + assert torch.equal( + pinned_packed_weight(quantizer, "iq2_xxs"), + pinned_packed_weight(reference.weight_quantizer, "iq2_xxs"), + ) + + def test_pin_outlives_the_weight_tensor(): model, _ = _gptq_model() # Same values in a new tensor, as an offload round trip or a layerwise resume leaves them. diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index ebb01ef052e..3955dd222b9 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -21,6 +21,8 @@ import pytest import torch import torch.nn.functional as F +from _test_utils.torch.quantization.models import TOY_MOE_EXPERTS as NUM_EXPERTS +from _test_utils.torch.quantization.models import TOY_MOE_HIDDEN as HIDDEN from _test_utils.torch.transformers_models import get_tiny_llama from torch import nn @@ -29,7 +31,6 @@ 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"} @@ -174,62 +175,6 @@ def test_a_module_calibration_never_reached_gets_the_plain_search(): 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) From d60c4ca4955bd3e4419cfd0c1134be43c7baa6c2 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 21 Oct 2026 06:18:50 +0000 Subject: [PATCH 2/3] Address review: GPTQ fused-experts tests and Hessian memory warning - Skip the fused-experts tests when transformers is missing: the shared moe fixture now calls pytest.importorskip, as the HF plugin wraps them. - Run the fused-expert GPTQ test with INT4 too, so the base GPTQHelper's weight plumbing and in-place write-back are covered. At the toy width a GGML weight is a single 256-column group whose GPTQ encode ignores the Hessian, so INT4 is also the case that checks per-expert routing. - Add a layerwise test on a two-layer toy MoE stack: every layer's experts are written, layer 0 matches a single pass, and layer 1 is calibrated on layer 0's quantized output. - Report the fused-expert Hessian size and warn before allocating when it exceeds the free memory of the GPUs holding the experts, pointing to layerwise calibration. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 19 ++++- tests/unit/torch/quantization/conftest.py | 1 + .../unit/torch/quantization/test_ggml_gptq.py | 76 +++++++++++++++++-- 3 files changed, 87 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index bd5dadb0dc0..3dd2bc96bea 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2345,6 +2345,23 @@ def gptq( print_rank_0("No quantized linear layers or fused experts found, skipping GPTQ") return + # Every expert Hessian lives from now until its update; a whole large MoE's does not fit. + expert_hessian_bytes = 0 + expert_devices = set() + for _, module in fused_experts: + for _, weight, quantizers, _ in _fused_expert_projections(module): + expert_hessian_bytes += ( + 4 * weight.shape[-1] ** 2 * sum(q.is_enabled for q in quantizers) + ) + expert_devices.add(weight.device) + gpus = [d for d in expert_devices if d.type == "cuda"] + if gpus and expert_hessian_bytes > sum(torch.cuda.mem_get_info(d)[0] for d in gpus): + warn_rank_0( + f"GPTQ: the fused-MoE expert Hessians need {expert_hessian_bytes / 2**30:.1f} GiB, more " + "than the GPUs holding the experts have free, so they spill to host memory and can " + "exhaust it. Enable layerwise calibration to hold one layer's Hessians at a time." + ) + def _make_gptq_handle(name, m, weight_quantizer=None, weight=None): quantizer = m.weight_quantizer if weight_quantizer is None else weight_quantizer backend = getattr(quantizer, "backend", None) @@ -2399,7 +2416,7 @@ def _expert_hook(_input_quantizer, _args, output): print_rank_0( f"Computing Hessians for {len(gptq_handles)} linear layers and " - f"{len(expert_handles)} fused-expert weights..." + f"{len(expert_handles)} fused-expert weights ({expert_hessian_bytes / 2**30:.1f} GiB)..." ) try: diff --git a/tests/unit/torch/quantization/conftest.py b/tests/unit/torch/quantization/conftest.py index 68b9f1d0f44..78d0ddd4463 100644 --- a/tests/unit/torch/quantization/conftest.py +++ b/tests/unit/torch/quantization/conftest.py @@ -23,6 +23,7 @@ @pytest.fixture def moe(): """A one-layer toy MoE with HF fused experts; unregisters its quantized expert class after.""" + pytest.importorskip("transformers") # the fused-experts quant wrapper is the HF plugin's torch.manual_seed(0) model = ToyMoE() yield model diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index ecbc648b235..1106d330e3b 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -22,7 +22,7 @@ import pytest import torch import torch.nn.functional as F -from _test_utils.torch.quantization.models import TOY_MOE_EXPERTS, TOY_MOE_HIDDEN +from _test_utils.torch.quantization.models import TOY_MOE_EXPERTS, TOY_MOE_HIDDEN, ToyMoE import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -35,6 +35,7 @@ compute_hessian_inverse, gptq_blockwise_update, ) +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector def _problem(rows=6, cols=16, seed=0): @@ -207,12 +208,29 @@ def test_a_module_calibration_never_reached_gets_the_plain_search(): torch.testing.assert_close(model["unused"].weight.detach(), plain, rtol=0, atol=0) -def test_gptq_quantizes_each_fused_expert_as_a_linear_fed_only_its_tokens(moe): +def _int4_config(algorithm): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 128}}, + "enable": True, + }, + ], + "algorithm": algorithm, + } + + +# iq2_xxs runs GGMLGPTQHelper and INT4 the base GPTQHelper. At this width a GGML weight is one +# 256-column group, which GGML GPTQ encodes without the Hessian, so only INT4 checks the routing. +@pytest.mark.parametrize("num_bits", ["iq2_xxs", 4]) +def test_gptq_quantizes_each_fused_expert_as_a_linear_fed_only_its_tokens(moe, num_bits): # The expert no token reaches is compared against a linear calibration never reaches. experts = moe.mlp.experts original = copy.deepcopy(experts) inputs = _inputs()[:96, :TOY_MOE_HIDDEN] - config = _iq_config(GPTQ, "iq2_xxs") + config = _iq_config(GPTQ, num_bits) if num_bits == "iq2_xxs" else _int4_config(GPTQ) mtq.quantize(moe, config, forward_loop=lambda m: m(inputs)) for idx in range(TOY_MOE_EXPERTS): @@ -224,11 +242,53 @@ def test_gptq_quantizes_each_fused_expert_as_a_linear_fed_only_its_tokens(moe): mtq.quantize(reference, config, forward_loop=lambda m, x=expert_inputs: m(x)) assert torch.equal(getattr(experts, name)[idx], reference.weight) - quantizer = getattr(experts, f"{name}_weight_quantizers")[idx] - assert torch.equal( - pinned_packed_weight(quantizer, "iq2_xxs"), - pinned_packed_weight(reference.weight_quantizer, "iq2_xxs"), - ) + if num_bits == "iq2_xxs": + quantizer = getattr(experts, f"{name}_weight_quantizers")[idx] + assert torch.equal( + pinned_packed_weight(quantizer, "iq2_xxs"), + pinned_packed_weight(reference.weight_quantizer, "iq2_xxs"), + ) + + +class _MoEStack(torch.nn.Module): + """Toy MoE decoder layers under ``layers``.""" + + def __init__(self, *layers): + super().__init__() + self.layers = torch.nn.ModuleList(layers) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +def test_layerwise_gptq_calibrates_fused_experts_on_the_earlier_layers_output(monkeypatch, moe): + monkeypatch.setattr( + LayerActivationCollector, + "_decoder_layer_support", + [(lambda m: isinstance(m, _MoEStack), lambda m: m.layers)], + ) + stack = _MoEStack(moe, ToyMoE()) + inputs = _inputs()[:96, :TOY_MOE_HIDDEN] + + def run(layerwise): + model = copy.deepcopy(stack) + algorithm = {**GPTQ, "layerwise": {"enable": layerwise}} + # INT4: one GGML group per weight here would ignore the Hessian, and with it the inputs. + mtq.quantize(model, _int4_config(algorithm), forward_loop=lambda m: m(inputs)) + return [layer.mlp.experts for layer in model.layers] + + def same(a, b): + return all(torch.equal(getattr(a, n), getattr(b, n)) for n in ("gate_up_proj", "down_proj")) + + single_pass, layerwise = run(False), run(True) + for before, after in zip(stack.layers, layerwise): + assert not same(before.mlp.experts, after) # GPTQ wrote every layer's experts + # Layer 0's experts see unquantized inputs either way. + assert same(single_pass[0], layerwise[0]) + # Layer 1's see layer 0's quantized experts only when layerwise. + assert not same(single_pass[1], layerwise[1]) def test_pin_outlives_the_weight_tensor(): From cedb19e1e9366a14de8a22e92d126b1b5340fa74 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 21 Oct 2026 07:34:55 +0000 Subject: [PATCH 3/3] Refuse GPTQ and imatrix on packed native FP8 experts Native Transformers FP8 experts store their weights in 8 bits and hand calibration a dequantized copy per expert, so the in-place update of GPTQ and imatrix landed in that copy and was silently lost. Both now refuse such experts before any calibration state changes and ask for unpack_weight(), as folding already does; after unpacking, the updates persist. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 23 ++++++++- .../plugins/test_fused_experts.py | 50 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3dd2bc96bea..5e7e62c4189 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -964,6 +964,22 @@ def _fused_expert_projections(module: nn.Module): yield weight_name, weight, quantizers, input_quantizer +def _refuse_packed_fused_expert_weights(name, module, method, updates): + """Refuse fused experts that store a projection ``method`` updates in 8 bits. + + Such experts (native FP8 experts) hand calibration dequantized copies of their weights, so an + in-place update would land in a temporary and be lost. ``updates(quantizer)`` says whether + ``method`` writes the weight that quantizer belongs to. + """ + for weight_name, weight, quantizers, _ in _fused_expert_projections(module): + if weight.element_size() == 1 and any(updates(q) for q in quantizers): + raise ValueError( + f"{method} updates {name}.{weight_name} in place, but it is stored in 8 bits " + "(native FP8 experts), so the update would be lost. Call unpack_weight() on the " + "module before calibrating." + ) + + class _WeightInputHook: """Identify active input collectors so auxiliary modules can request a calibration forward.""" @@ -2332,6 +2348,10 @@ def gptq( """ total_start = time.time() + fused_experts = [(n, m) for n, m in model.named_modules() if _is_quant_fused_experts(m)] + for name, module in fused_experts: + _refuse_packed_fused_expert_weights(name, module, "GPTQ", lambda q: q.is_enabled) + # TODO: Add support for other scale setting strateiges like weight-mse or local-hessian max_calibrate(model, forward_loop=forward_loop) @@ -2340,7 +2360,6 @@ def gptq( for n, m in model.named_modules() if is_quantized_linear(m) and m.weight_quantizer.is_enabled ] - fused_experts = [(n, m) for n, m in model.named_modules() if _is_quant_fused_experts(m)] if not quantized_layers and not fused_experts: print_rank_0("No quantized linear layers or fused experts found, skipping GPTQ") return @@ -2494,6 +2513,8 @@ def _accepted(quantizer): # 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: + if _is_quant_fused_experts(module): + _refuse_packed_fused_expert_weights(name, module, "imatrix", lambda q: id(q) in targets) for weight, quantizer in module.iter_weights_for_calibration(): if id(quantizer) in targets: try: diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 4f224e3b7c0..6eaf13f0bc1 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1983,3 +1983,53 @@ def test_native_fp8_experts_unpacked_temporary_fold(has_gate, exception, snapsho assert all(q.is_enabled for q in quantizers) torch.testing.assert_close(experts.state_dict(), state, rtol=0, atol=0) torch.testing.assert_close(experts(inputs, indices, routing), expected) + + +@pytest.mark.parametrize( + ("cfg", "algorithm"), + [ + ({"num_bits": 4, "block_sizes": {-1: 128}}, {"method": "gptq", "block_size": 128}), + ({"num_bits": "iq2_xxs", "block_sizes": {-1: 256}, "backend": "ggml"}, "imatrix"), + ], + ids=["gptq", "imatrix"], +) +def test_native_fp8_experts_reject_packed_weight_updates(cfg, algorithm): + # GPTQ and imatrix write the experts' weights in place; packed FP8 experts only hand them + # dequantized copies, so the update would be lost. + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts(config, block_size=None, has_gate=True) + for name in ("gate_up_proj", "down_proj"): + weight = getattr(experts, name) + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + getattr(experts, f"{name}_scale_inv").fill_(0.02) + quant_cfg = [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": cfg}, + ] + mtq.quantize(experts, {"quant_cfg": quant_cfg, "algorithm": "max"}, forward_loop=None) + inputs = torch.randn(64, 256) + indices = (torch.arange(64) % 2).unsqueeze(1) + routing = torch.ones(64, 1) + state = copy.deepcopy(experts.state_dict()) + + with pytest.raises(ValueError, match=r"Call unpack_weight\(\)"): + mtq.calibrate(experts, algorithm, forward_loop=lambda m: m(inputs, indices, routing)) + # Compare FP8 payloads in FP32 because CPU close/equality kernels do not support FP8. + torch.testing.assert_close( + {name: value.float() for name, value in experts.state_dict().items()}, + {name: value.float() for name, value in state.items()}, + rtol=0, + atol=0, + ) + + experts.unpack_weight() + unpacked = { + name: getattr(experts, name).detach().clone() for name in ("gate_up_proj", "down_proj") + } + mtq.calibrate(experts, algorithm, forward_loop=lambda m: m(inputs, indices, routing)) + for name, weight in unpacked.items(): + assert not torch.equal(getattr(experts, name), weight)