From 586103d27718c07d58cdaaebf38ece763310fde6 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 25 Oct 2026 07:08:58 +0000 Subject: [PATCH 1/6] Skip Qwen3.8's QSA indexer loop during calibration when it keeps every token Transformers' Qwen4ExpTextQSAIndexer loops in Python over every (sample, query) pair: about 340 s per forward of 1,024 x 512 calibration tokens on each of Qwen3.8-Flash-Next's 12 full-attention layers, against 33 s with it skipped. A query keeps its block_topk best complete key blocks plus the incomplete tail, so when the keys hold no more blocks than that, the selection is the visibility mask the indexer was given. The calibration hook for model type qwen4_exp now gives each indexer instance a forward that returns that mask directly, bit-identical to the loop, and runs the loop otherwise. The hf.py dispatcher lets a model plugin define only the hooks it needs. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + modelopt/torch/models/__init__.py | 1 + modelopt/torch/models/hf.py | 18 ++-- modelopt/torch/models/qwen4_exp/__init__.py | 18 ++++ .../torch/models/qwen4_exp/qsa_indexer.py | 69 ++++++++++++ modelopt/torch/models/qwen4_exp/specs.py | 25 +++++ .../models/test_qwen4_exp_qsa_indexer.py | 100 ++++++++++++++++++ 7 files changed, 225 insertions(+), 7 deletions(-) create mode 100644 modelopt/torch/models/qwen4_exp/__init__.py create mode 100644 modelopt/torch/models/qwen4_exp/qsa_indexer.py create mode 100644 modelopt/torch/models/qwen4_exp/specs.py create mode 100644 tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c13b4b92956..283c221d8c8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog - Add ``hessian_input_dtype`` to the GPTQ and ``local_hessian`` algorithm configs. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, several times faster and without a float32 copy of each batch; the default ``float32`` keeps the existing float32 GEMM. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2. +- ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration. *Speculative Decoding* diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index 98aaf9e57da..a855418664e 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -93,6 +93,7 @@ qwen3_moe, qwen3_next, qwen3_vl_moe, + qwen4_exp, step3p5, step3p7, ) diff --git a/modelopt/torch/models/hf.py b/modelopt/torch/models/hf.py index 3b765ddd2ba..22472911a9f 100644 --- a/modelopt/torch/models/hf.py +++ b/modelopt/torch/models/hf.py @@ -26,9 +26,12 @@ ] # Keep modeling dependencies lazy: importing model specs must not import calibration code. +# A plugin defines any of prepare_for_loading, has_mtp_weights and prepare_for_calibration. _PLUGINS = { "nemotron_h": "nemotron_h.mtp", "nemotron_h_omni": "nemotron_h.mtp", + "qwen4_exp": "qwen4_exp.qsa_indexer", + "qwen4_exp_text": "qwen4_exp.qsa_indexer", } @@ -42,16 +45,16 @@ def prepare_model_for_loading( model_type, checkpoint_path: str, trust_remote_code: bool, *, model_class=None ): """Prepare checkpoint-only modules, optionally on the loader's selected concrete class.""" - plugin = _get_plugin(model_type) - if plugin is None: + hook = getattr(_get_plugin(model_type), "prepare_for_loading", None) + if hook is None: return nullcontext() - return plugin.prepare_for_loading(checkpoint_path, trust_remote_code, model_class=model_class) + return hook(checkpoint_path, trust_remote_code, model_class=model_class) def checkpoint_has_mtp(model_type, checkpoint_path: str) -> bool: """Use the model plugin to detect MTP tensors before selecting a supported loading path.""" - plugin = _get_plugin(model_type) - return plugin is not None and plugin.has_mtp_weights(checkpoint_path) + hook = getattr(_get_plugin(model_type), "has_mtp_weights", None) + return hook is not None and hook(checkpoint_path) def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: @@ -65,5 +68,6 @@ def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwa def prepare_model_for_calibration(model) -> None: """Run the optional model-specific hook that augments the default calibration forward.""" plugin = _get_plugin(getattr(getattr(model, "config", None), "model_type", None)) - if plugin is not None: - plugin.prepare_for_calibration(model) + hook = getattr(plugin, "prepare_for_calibration", None) + if hook is not None: + hook(model) diff --git a/modelopt/torch/models/qwen4_exp/__init__.py b/modelopt/torch/models/qwen4_exp/__init__.py new file mode 100644 index 00000000000..79c462197c2 --- /dev/null +++ b/modelopt/torch/models/qwen4_exp/__init__.py @@ -0,0 +1,18 @@ +# 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. + +"""Per-model specs for HF model type ``qwen4_exp``.""" + +from . import specs diff --git a/modelopt/torch/models/qwen4_exp/qsa_indexer.py b/modelopt/torch/models/qwen4_exp/qsa_indexer.py new file mode 100644 index 00000000000..db9c53c476c --- /dev/null +++ b/modelopt/torch/models/qwen4_exp/qsa_indexer.py @@ -0,0 +1,69 @@ +# 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. + +"""Calibration support for Qwen3.8's QSA indexer.""" + +import types + +import torch + +__all__ = ["prepare_for_calibration"] + + +def _forward(self, hidden_states, position_embeddings, attention_mask, past_key_values): + """Transformers' indexer, short-circuited when every visible key fits the token budget. + + The reference loops in Python over every (sample, query) pair: about 340 s per forward of + 1,024 x 512 calibration tokens, on each full-attention layer. A query keeps its + ``block_topk`` best complete key blocks plus the incomplete tail, so when the keys hold no + more blocks than that, it keeps every visible token: the selection is the visibility mask the + indexer was given, which this returns directly. + """ + if attention_mask.shape[-1] // self.compress_ratio > self.block_topk: + return self._modelopt_reference_forward( + hidden_states, position_embeddings, attention_mask, past_key_values + ) + if past_key_values is not None: + # The reference appends this step's indexer keys to the cache before selecting. + batch_size, seq_length, _ = hidden_states.shape + token_k = self.index_qk_proj(hidden_states)[..., self.index_n_heads * self.index_head_dim :] + past_key_values.update_indexer( + token_k.reshape(batch_size, seq_length, -1, self.index_head_dim).squeeze(2), + self.layer_idx, + ) + # The reference reads head slot 0 of the 4-D mask and returns one head slot. + mask = attention_mask[:, :1] + if attention_mask.is_floating_point(): + min_dtype = torch.finfo(attention_mask.dtype).min + return torch.where(mask == 0, mask.new_zeros(()), min_dtype) + return mask.clone() + + +def prepare_for_calibration(model) -> bool: + """Give each QSA indexer of ``model`` the fast path; return whether ``model`` has one. + + Calibration sequences rarely exceed the indexer's token budget, so on Qwen3.8-Flash-Next + (2,048 tokens) this skips the loop for every ``--calib_seq`` up to 2,048, with bit-identical + results, and runs the reference otherwise. + """ + found = False + for module in model.modules(): + if type(module).__name__ != "Qwen4ExpTextQSAIndexer": + continue + found = True + if not hasattr(module, "_modelopt_reference_forward"): + module._modelopt_reference_forward = module.forward + module.forward = types.MethodType(_forward, module) + return found diff --git a/modelopt/torch/models/qwen4_exp/specs.py b/modelopt/torch/models/qwen4_exp/specs.py new file mode 100644 index 00000000000..a32675ac351 --- /dev/null +++ b/modelopt/torch/models/qwen4_exp/specs.py @@ -0,0 +1,25 @@ +# 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. + +"""Qwen3.8 specs (HF model type ``qwen4_exp``).""" + +from ..specs import ModelSpec, register + +__all__: list[str] = [] + +# Only the model type is recorded, so the dispatcher in ``hf.py`` can find this package's +# calibration helper. The fused MoE layout is left undeclared: export and quantization handle +# these experts generically, and declaring it would change their behavior for this model. +register(ModelSpec(model_type="qwen4_exp", min_transformers_version="5.16")) diff --git a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py new file mode 100644 index 00000000000..38a31f23356 --- /dev/null +++ b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py @@ -0,0 +1,100 @@ +# 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. + +"""Qwen3.8's QSA indexer fast path against Transformers' reference loop.""" + +import types + +import pytest +import torch +import torch.nn as nn + +modeling = pytest.importorskip("transformers.models.qwen4_exp.modeling_qwen4_exp") +from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig + +from modelopt.torch.models.hf import prepare_model_for_calibration + +SEQ = 64 + + +class _IndexerCache: + """Records the indexer keys a forward appends, as the model's cache would.""" + + def __init__(self): + self.keys = [] + + def update_indexer(self, keys, layer_idx): + self.keys.append(keys) + return keys + + +def _mask(pads, floating): + """Causal mask over SEQ tokens, with padded key positions hidden per row.""" + mask = torch.tril(torch.ones(SEQ, SEQ, dtype=torch.bool)).expand(len(pads), 1, SEQ, SEQ).clone() + for row, (side, n) in enumerate(pads): + mask[row, :, :, slice(0, n) if side == "left" else slice(SEQ - n, SEQ)] = False + if floating: + return torch.where(mask, 0.0, torch.finfo(torch.float32).min) + return mask + + +def _reference_must_not_run(*args, **kwargs): + raise AssertionError("the reference indexer loop ran although every visible key fits") + + +@pytest.mark.parametrize( + ("budget", "pads", "floating", "cache"), + [ + (64, [("right", 0), ("right", 9)], False, False), + (64, [("left", 13), ("right", 0)], True, False), + (64, [("left", 5), ("right", 3)], False, True), + (32, [("right", 0)], False, False), # 16 key blocks, 8 kept: the loop must run + ], + ids=["bool_mask", "float_mask", "cache", "over_budget"], +) +def test_qsa_indexer_fast_path_matches_the_reference(monkeypatch, budget, pads, floating, cache): + torch.manual_seed(0) + config = Qwen4ExpTextConfig( + hidden_size=64, + head_dim=16, + num_attention_heads=4, + num_key_value_heads=2, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=16, + indexer_budget=budget, + indexer_compress_ratio=4, + ) + model = nn.Module() + model.config = types.SimpleNamespace(model_type="qwen4_exp") + model.indexer = modeling.Qwen4ExpTextQSAIndexer(config, layer_idx=0) + hidden = torch.randn(len(pads), SEQ, config.hidden_size) + rope = modeling.Qwen4ExpTextRotaryEmbedding(config)( + hidden, torch.arange(SEQ).expand(3, len(pads), SEQ) + ) + mask = _mask(pads, floating) + expected_cache, cache_seen = (_IndexerCache(), _IndexerCache()) if cache else (None, None) + + with torch.no_grad(): + expected = type(model.indexer).forward(model.indexer, hidden, rope, mask, expected_cache) + prepare_model_for_calibration(model) + if SEQ // config.indexer_compress_ratio <= model.indexer.block_topk: + monkeypatch.setattr(type(model.indexer), "forward", _reference_must_not_run) + actual = model.indexer(hidden, rope, mask, cache_seen) + + assert actual.dtype == expected.dtype + assert torch.equal(actual, expected) + if cache: + assert torch.equal(cache_seen.keys[0], expected_cache.keys[0]) From 89cd188bdd2bf4aeabf253c5c5a0d4839ceb0dea Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 21:41:48 +0000 Subject: [PATCH 2/6] Run the QSA indexer's projection on the fast path too The fast path returned the visibility mask without calling index_qk_proj when no cache was passed. The reference calls it on every forward, so a recipe that enables input quantizers on every linear (FP8, for example) would leave that projection's input quantizer without calibration data and export it uncalibrated (review). The fast path now runs the projection once, as the reference does, and reuses its keys for the cache. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/models/qwen4_exp/qsa_indexer.py | 4 +++- tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/qwen4_exp/qsa_indexer.py b/modelopt/torch/models/qwen4_exp/qsa_indexer.py index db9c53c476c..ea882004025 100644 --- a/modelopt/torch/models/qwen4_exp/qsa_indexer.py +++ b/modelopt/torch/models/qwen4_exp/qsa_indexer.py @@ -35,10 +35,12 @@ def _forward(self, hidden_states, position_embeddings, attention_mask, past_key_ return self._modelopt_reference_forward( hidden_states, position_embeddings, attention_mask, past_key_values ) + # The reference runs the projection on every forward, so a quantizer on it gets calibrated. + qk = self.index_qk_proj(hidden_states) if past_key_values is not None: # The reference appends this step's indexer keys to the cache before selecting. batch_size, seq_length, _ = hidden_states.shape - token_k = self.index_qk_proj(hidden_states)[..., self.index_n_heads * self.index_head_dim :] + token_k = qk[..., self.index_n_heads * self.index_head_dim :] past_key_values.update_indexer( token_k.reshape(batch_size, seq_length, -1, self.index_head_dim).squeeze(2), self.layer_idx, diff --git a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py index 38a31f23356..1343e93dd4e 100644 --- a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py +++ b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py @@ -92,9 +92,13 @@ def test_qsa_indexer_fast_path_matches_the_reference(monkeypatch, budget, pads, prepare_model_for_calibration(model) if SEQ // config.indexer_compress_ratio <= model.indexer.block_topk: monkeypatch.setattr(type(model.indexer), "forward", _reference_must_not_run) + projections = [] + model.indexer.index_qk_proj.register_forward_hook(lambda *_: projections.append(None)) actual = model.indexer(hidden, rope, mask, cache_seen) assert actual.dtype == expected.dtype assert torch.equal(actual, expected) + # The projection still runs once, as in the reference, so a quantizer on it gets calibrated. + assert len(projections) == 1 if cache: assert torch.equal(cache_seen.keys[0], expected_cache.keys[0]) From 7ce378b698e7f6796b11d1f446f4249e807973b7 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 21:44:43 +0000 Subject: [PATCH 3/6] Run imatrix encode workers in the caller's inference mode Inference mode, like grad mode and the current device, is per thread, so a worker encoding weights that are inference tensors failed with "Inplace update to inference tensor outside InferenceMode" when the caller ran imatrix under torch.inference_mode() (review). Each worker now enters inference mode when the caller is in it, and no_grad otherwise. The 2-GPU test runs both ways. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 6 ++-- .../quantization/test_ggml_imatrix_gpu.py | 29 +++++++++++-------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index a784e63f57f..0a13fcfff88 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2765,10 +2765,12 @@ def _encode_on_each_gpu(by_device, encode): # CUDA extension, happens once before the workers start. first, *by_device[gpus[0]] = by_device[gpus[0]] encode(*first) + inference = torch.is_inference_mode_enabled() def worker(device): - torch.cuda.set_device(device) # per thread, like grad mode - with torch.no_grad(): + # The device, grad mode and inference mode are per thread: take the caller's. + torch.cuda.set_device(device) + with torch.inference_mode() if inference else torch.no_grad(): for item in by_device[device]: encode(*item) torch.cuda.synchronize(device) diff --git a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py index 31590a0e063..1ab9812762f 100644 --- a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py +++ b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py @@ -15,9 +15,9 @@ """GPU tests for imatrix calibration of the GGML formats.""" -import copy import time +import pytest import torch from torch import nn @@ -49,8 +49,10 @@ def forward(self, x): return x -def test_imatrix_encodes_each_gpus_weights_at_once(need_2_gpus, monkeypatch): - """Weights on two GPUs are encoded concurrently, to the same codes as on one GPU.""" +@pytest.mark.parametrize("inference_mode", [False, True]) +def test_imatrix_encodes_each_gpus_weights_at_once(need_2_gpus, monkeypatch, inference_mode): + """Weights on two GPUs are encoded concurrently, to the same codes as on one GPU, also as + inference tensors under the caller's inference mode, which worker threads do not inherit.""" spans = [] encode = GGMLImatrixEncoder.encode @@ -58,20 +60,23 @@ def timed_encode(self, weight, quantizer, second_moments): start = time.perf_counter() decoded = encode(self, weight, quantizer, second_moments) torch.cuda.synchronize(weight.device) - time.sleep(0.2) # long enough that encodes on different GPUs must overlap to save time + # Stands in for a long encode: encodes on two GPUs overlap only if they run at once. + time.sleep(0.2) spans.append((weight.device.index, start, time.perf_counter())) return decoded monkeypatch.setattr(GGMLImatrixEncoder, "encode", timed_encode) - torch.manual_seed(0) - reference = _Layers().cuda() - model = copy.deepcopy(reference) - model.layers[2:].to("cuda:1") - inputs = torch.randn(8, 128, 512, device="cuda") + with torch.inference_mode(inference_mode): + torch.manual_seed(0) + reference = _Layers().cuda() + torch.manual_seed(0) + model = _Layers().cuda() + model.layers[2:].to("cuda:1") + inputs = torch.randn(8, 128, 512, device="cuda") - mtq.quantize(reference, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs)) - spans.clear() - mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs)) + mtq.quantize(reference, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs)) + spans.clear() + mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs)) for expected, layer in zip(reference.layers, model.layers): assert torch.equal(layer.weight.cuda(0), expected.weight) From ba0a5ee4feb18d9127cbe2c48d1a3e389e484d20 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 22:02:36 +0000 Subject: [PATCH 4/6] Treat a plugin's prefers_bundled_code hook as optional prefers_bundled_code, added on main after this PR opened, called the hook unconditionally, so hf_ptq.py with trust_remote_code on a qwen4_exp checkpoint raised AttributeError before loading: the Qwen3.8 plugin defines only prepare_for_calibration (review). It now falls back to False when the plugin lacks the hook, like the other plugin hooks. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/models/hf.py | 9 ++++----- .../unit/torch/models/test_qwen4_exp_qsa_indexer.py | 13 ++++++++++++- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/hf.py b/modelopt/torch/models/hf.py index 22472911a9f..baea9c3c325 100644 --- a/modelopt/torch/models/hf.py +++ b/modelopt/torch/models/hf.py @@ -26,7 +26,8 @@ ] # Keep modeling dependencies lazy: importing model specs must not import calibration code. -# A plugin defines any of prepare_for_loading, has_mtp_weights and prepare_for_calibration. +# A plugin defines any of prepare_for_loading, has_mtp_weights, prefers_bundled_code and +# prepare_for_calibration. _PLUGINS = { "nemotron_h": "nemotron_h.mtp", "nemotron_h_omni": "nemotron_h.mtp", @@ -59,10 +60,8 @@ def checkpoint_has_mtp(model_type, checkpoint_path: str) -> bool: def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: """Use the model plugin to tell whether the checkpoint's bundled code must replace the built-in class.""" - plugin = _get_plugin(model_type) - return plugin is not None and plugin.prefers_bundled_code( - hf_config, checkpoint_path, config_kwargs - ) + hook = getattr(_get_plugin(model_type), "prefers_bundled_code", None) + return hook is not None and hook(hf_config, checkpoint_path, config_kwargs) def prepare_model_for_calibration(model) -> None: diff --git a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py index 1343e93dd4e..131fdcf6d1b 100644 --- a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py +++ b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py @@ -24,7 +24,11 @@ modeling = pytest.importorskip("transformers.models.qwen4_exp.modeling_qwen4_exp") from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig -from modelopt.torch.models.hf import prepare_model_for_calibration +from modelopt.torch.models.hf import ( + checkpoint_has_mtp, + prefers_bundled_code, + prepare_model_for_calibration, +) SEQ = 64 @@ -102,3 +106,10 @@ def test_qsa_indexer_fast_path_matches_the_reference(monkeypatch, budget, pads, assert len(projections) == 1 if cache: assert torch.equal(cache_seen.keys[0], expected_cache.keys[0]) + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text"]) +def test_loading_hooks_the_plugin_lacks_keep_their_defaults(model_type): + """The Qwen3.8 plugin only has a calibration hook; the loading hooks fall back to the defaults.""" + assert not prefers_bundled_code(model_type, None, "ckpt", {}) + assert not checkpoint_has_mtp(model_type, "ckpt") From 18e5989c0864b67b5eca5f086e77ae1dfe84a0bb Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 23:12:41 +0000 Subject: [PATCH 5/6] Keep Accelerate hooks around the QSA fast path; test it harder Review (meenchen): - Bind the fast path with bind_forward_method instead of assigning forward, so an Accelerate hook a device_map load already put on the indexer still wraps it and moves its inputs. - The equivalence test now also forbids the forward the plugin saved, not only the class forward, so a fast path that always delegated to the loop would fail. - Cover an incomplete key tail, a real accumulating indexer cache, and decoding that crosses from the fast path into the loop, comparing every step and the cached keys with an unpatched indexer. - Drop the timing from the docstring; it belongs in the PR description. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- .../torch/models/qwen4_exp/qsa_indexer.py | 11 +- .../models/test_qwen4_exp_qsa_indexer.py | 149 +++++++++++++----- 2 files changed, 112 insertions(+), 48 deletions(-) diff --git a/modelopt/torch/models/qwen4_exp/qsa_indexer.py b/modelopt/torch/models/qwen4_exp/qsa_indexer.py index ea882004025..8d10bf4bc8b 100644 --- a/modelopt/torch/models/qwen4_exp/qsa_indexer.py +++ b/modelopt/torch/models/qwen4_exp/qsa_indexer.py @@ -15,18 +15,17 @@ """Calibration support for Qwen3.8's QSA indexer.""" -import types - import torch +from modelopt.torch.utils.network import bind_forward_method + __all__ = ["prepare_for_calibration"] def _forward(self, hidden_states, position_embeddings, attention_mask, past_key_values): """Transformers' indexer, short-circuited when every visible key fits the token budget. - The reference loops in Python over every (sample, query) pair: about 340 s per forward of - 1,024 x 512 calibration tokens, on each full-attention layer. A query keeps its + The reference loops in Python over every (sample, query) pair. A query keeps its ``block_topk`` best complete key blocks plus the incomplete tail, so when the keys hold no more blocks than that, it keeps every visible token: the selection is the visibility mask the indexer was given, which this returns directly. @@ -66,6 +65,6 @@ def prepare_for_calibration(model) -> bool: continue found = True if not hasattr(module, "_modelopt_reference_forward"): - module._modelopt_reference_forward = module.forward - module.forward = types.MethodType(_forward, module) + # Keeps an Accelerate hook already on the indexer, such as device_map's, around the fast path. + bind_forward_method(module, _forward, "_modelopt_reference_forward") return found diff --git a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py index 131fdcf6d1b..8420e37cc69 100644 --- a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py +++ b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py @@ -15,6 +15,7 @@ """Qwen3.8's QSA indexer fast path against Transformers' reference loop.""" +import copy import types import pytest @@ -22,6 +23,7 @@ import torch.nn as nn modeling = pytest.importorskip("transformers.models.qwen4_exp.modeling_qwen4_exp") +from transformers.cache_utils import Cache, DynamicIndexedLayer from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig from modelopt.torch.models.hf import ( @@ -30,25 +32,45 @@ prepare_model_for_calibration, ) -SEQ = 64 + +def _model(budget): + """A model holding one QSA indexer with 4-token key blocks, ``budget // 4`` of them kept.""" + config = Qwen4ExpTextConfig( + hidden_size=64, + head_dim=16, + num_attention_heads=4, + num_key_value_heads=2, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=16, + indexer_budget=budget, + indexer_compress_ratio=4, + ) + model = nn.Module() + model.config = types.SimpleNamespace(model_type="qwen4_exp") + model.indexer = modeling.Qwen4ExpTextQSAIndexer(config, layer_idx=0) + return model, config -class _IndexerCache: - """Records the indexer keys a forward appends, as the model's cache would.""" +def _inputs(config, batch, seq): + hidden = torch.randn(batch, seq, config.hidden_size) + rope = modeling.Qwen4ExpTextRotaryEmbedding(config)( + hidden, torch.arange(seq).expand(3, batch, seq) + ) + return hidden, rope - def __init__(self): - self.keys = [] - def update_indexer(self, keys, layer_idx): - self.keys.append(keys) - return keys +def _cache(): + return Cache(layers=[DynamicIndexedLayer()]) -def _mask(pads, floating): - """Causal mask over SEQ tokens, with padded key positions hidden per row.""" - mask = torch.tril(torch.ones(SEQ, SEQ, dtype=torch.bool)).expand(len(pads), 1, SEQ, SEQ).clone() +def _mask(pads, floating, seq, queries=None): + """Causal mask of the last ``queries`` tokens over ``seq`` keys, padded key positions hidden.""" + queries = seq if queries is None else queries + causal = torch.tril(torch.ones(seq, seq, dtype=torch.bool))[seq - queries :] + mask = causal.expand(len(pads), 1, queries, seq).clone() for row, (side, n) in enumerate(pads): - mask[row, :, :, slice(0, n) if side == "left" else slice(SEQ - n, SEQ)] = False + mask[row, :, :, slice(0, n) if side == "left" else slice(seq - n, seq)] = False if floating: return torch.where(mask, 0.0, torch.finfo(torch.float32).min) return mask @@ -58,44 +80,37 @@ def _reference_must_not_run(*args, **kwargs): raise AssertionError("the reference indexer loop ran although every visible key fits") +def _forbid_the_loop(monkeypatch, indexer): + """Fail if the loop runs: through the class, or through the forward the plugin saved.""" + monkeypatch.setattr(type(indexer), "forward", _reference_must_not_run) + monkeypatch.setattr(indexer, "_modelopt_reference_forward", _reference_must_not_run) + + @pytest.mark.parametrize( - ("budget", "pads", "floating", "cache"), + ("budget", "seq", "pads", "floating", "cache"), [ - (64, [("right", 0), ("right", 9)], False, False), - (64, [("left", 13), ("right", 0)], True, False), - (64, [("left", 5), ("right", 3)], False, True), - (32, [("right", 0)], False, False), # 16 key blocks, 8 kept: the loop must run + (64, 64, [("right", 0), ("right", 9)], False, False), + (64, 64, [("left", 13), ("right", 0)], True, False), + (64, 64, [("left", 5), ("right", 3)], False, True), + (64, 62, [("left", 7), ("right", 0)], False, True), # 15 key blocks and a 2-token tail + (32, 64, [("right", 0)], False, False), # 16 key blocks, 8 kept: the loop must run ], - ids=["bool_mask", "float_mask", "cache", "over_budget"], + ids=["bool_mask", "float_mask", "cache", "incomplete_tail", "over_budget"], ) -def test_qsa_indexer_fast_path_matches_the_reference(monkeypatch, budget, pads, floating, cache): +def test_qsa_indexer_fast_path_matches_the_reference( + monkeypatch, budget, seq, pads, floating, cache +): torch.manual_seed(0) - config = Qwen4ExpTextConfig( - hidden_size=64, - head_dim=16, - num_attention_heads=4, - num_key_value_heads=2, - indexer_n_heads=2, - indexer_kv_heads=1, - indexer_head_dim=16, - indexer_budget=budget, - indexer_compress_ratio=4, - ) - model = nn.Module() - model.config = types.SimpleNamespace(model_type="qwen4_exp") - model.indexer = modeling.Qwen4ExpTextQSAIndexer(config, layer_idx=0) - hidden = torch.randn(len(pads), SEQ, config.hidden_size) - rope = modeling.Qwen4ExpTextRotaryEmbedding(config)( - hidden, torch.arange(SEQ).expand(3, len(pads), SEQ) - ) - mask = _mask(pads, floating) - expected_cache, cache_seen = (_IndexerCache(), _IndexerCache()) if cache else (None, None) + model, config = _model(budget) + hidden, rope = _inputs(config, len(pads), seq) + mask = _mask(pads, floating, seq) + expected_cache, cache_seen = (_cache(), _cache()) if cache else (None, None) with torch.no_grad(): expected = type(model.indexer).forward(model.indexer, hidden, rope, mask, expected_cache) prepare_model_for_calibration(model) - if SEQ // config.indexer_compress_ratio <= model.indexer.block_topk: - monkeypatch.setattr(type(model.indexer), "forward", _reference_must_not_run) + if seq // config.indexer_compress_ratio <= model.indexer.block_topk: + _forbid_the_loop(monkeypatch, model.indexer) projections = [] model.indexer.index_qk_proj.register_forward_hook(lambda *_: projections.append(None)) actual = model.indexer(hidden, rope, mask, cache_seen) @@ -105,7 +120,57 @@ def test_qsa_indexer_fast_path_matches_the_reference(monkeypatch, budget, pads, # The projection still runs once, as in the reference, so a quantizer on it gets calibrated. assert len(projections) == 1 if cache: - assert torch.equal(cache_seen.keys[0], expected_cache.keys[0]) + assert torch.equal(cache_seen.layers[0].indexer_keys, expected_cache.layers[0].indexer_keys) + + +def test_qsa_indexer_decoding_past_the_budget_matches_the_reference(): + """Prefill 30 tokens, then decode one token at a time past the 32-token budget: every step + matches the reference, as the indexer moves from the fast path to the loop, and so do the + indexer keys the cache accumulates.""" + torch.manual_seed(0) + model, config = _model(32) + reference = copy.deepcopy(model.indexer) + prepare_model_for_calibration(model) + hidden, rope = _inputs(config, 1, 42) + caches = (_cache(), _cache()) + + with torch.no_grad(): + for start, end in [(0, 30), *((step, step + 1) for step in range(30, 42))]: + args = (hidden[:, start:end], (rope[0][:, :end], rope[1][:, :end])) + mask = _mask([("right", 0)], False, end, queries=end - start) + expected = reference(*args, mask, caches[0]) + assert torch.equal(model.indexer(*args, mask, caches[1]), expected), (start, end) + + assert torch.equal(caches[1].layers[0].indexer_keys, caches[0].layers[0].indexer_keys) + + +def test_qsa_indexer_fast_path_runs_inside_an_accelerate_hook(monkeypatch): + """``device_map`` loading puts an Accelerate hook on the indexer before calibration support is + added; the fast path must run inside it, so the hook still moves the inputs.""" + hooks = pytest.importorskip("accelerate.hooks") + + class _CountingHook(hooks.ModelHook): + calls = 0 + + def pre_forward(self, module, *args, **kwargs): + self.calls += 1 + return args, kwargs + + torch.manual_seed(0) + model, config = _model(64) + hidden, rope = _inputs(config, 1, 64) + mask = _mask([("right", 0)], False, 64) + hook = _CountingHook() + hooks.add_hook_to_module(model.indexer, hook) + + with torch.no_grad(): + expected = model.indexer(hidden, rope, mask, None) + prepare_model_for_calibration(model) + _forbid_the_loop(monkeypatch, model.indexer) + actual = model.indexer(hidden, rope, mask, None) + + assert hook.calls == 2 + assert torch.equal(actual, expected) @pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text"]) From 08f82275a31fcf3984bae2c8f4631370de28a7bd Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 30 Oct 2026 04:16:52 +0000 Subject: [PATCH 6/6] Warn on a missing QSA indexer; refuse misspelled plugin hooks Review (Fridah-nv, non-blocking): - prepare_for_calibration warns when a Qwen3.8 model has no indexer it recognizes, such as after a Transformers class rename, instead of silently falling back to the slow loop. - The fast path hands a call with arguments it does not know, from a newer Transformers, to the reference forward instead of raising TypeError. - Optional hooks made a misspelled hook a silent no-op. A plugin now lists exactly its hooks in __all__, and the dispatcher refuses any other name there. models/README.md documents the four hooks. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/models/README.md | 12 ++++++ modelopt/torch/models/hf.py | 17 ++++++-- .../torch/models/qwen4_exp/qsa_indexer.py | 18 +++++++-- .../models/test_qwen4_exp_qsa_indexer.py | 40 +++++++++++++++++++ 4 files changed, 80 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/models/README.md b/modelopt/torch/models/README.md index c0caf7f6e01..1d54ffa47dd 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -105,5 +105,17 @@ only; subsystem logic stays in the subsystem. Keep architecture-specific loading and calibration helpers alongside the model's `specs.py` (for example, `nemotron_h/mtp.py`). The `hf.py` dispatcher loads these helpers on demand; examples call its loading context and calibration hook. + +A helper module registered in `hf.py`'s `_PLUGINS` defines any of these hooks, all +optional, and lists exactly the ones it defines in `__all__`. The dispatcher rejects any +other name there, so a misspelled hook fails instead of being skipped. + +| Hook | Called by | Default without it | +|---|---|---| +| `prepare_for_loading(checkpoint_path, trust_remote_code, *, model_class)` | `prepare_model_for_loading` | no context | +| `has_mtp_weights(checkpoint_path)` | `checkpoint_has_mtp` | `False` | +| `prefers_bundled_code(hf_config, checkpoint_path, config_kwargs)` | `prefers_bundled_code` | `False` | +| `prepare_for_calibration(model)` | `prepare_model_for_calibration` | nothing | + Do not import runtime helpers from a model package's `__init__.py`: registering specs must remain independent of optional Transformers and calibration imports. diff --git a/modelopt/torch/models/hf.py b/modelopt/torch/models/hf.py index baea9c3c325..236f5386821 100644 --- a/modelopt/torch/models/hf.py +++ b/modelopt/torch/models/hf.py @@ -26,8 +26,10 @@ ] # Keep modeling dependencies lazy: importing model specs must not import calibration code. -# A plugin defines any of prepare_for_loading, has_mtp_weights, prefers_bundled_code and -# prepare_for_calibration. +# A plugin defines any of these hooks and lists exactly those in its __all__ (see README.md). +_HOOKS = frozenset( + {"prepare_for_loading", "has_mtp_weights", "prefers_bundled_code", "prepare_for_calibration"} +) _PLUGINS = { "nemotron_h": "nemotron_h.mtp", "nemotron_h_omni": "nemotron_h.mtp", @@ -39,7 +41,16 @@ def _get_plugin(model_type): """Return the optional ModelOpt support plugin for a Hugging Face model type.""" name = _PLUGINS.get(model_type) - return None if name is None else importlib.import_module(f".{name}", __package__) + if name is None: + return None + plugin = importlib.import_module(f".{name}", __package__) + # Hooks are optional, so a misspelled one would be skipped silently: refuse it instead. + unknown = set(getattr(plugin, "__all__", ())) - _HOOKS + if unknown: + raise ValueError( + f"{plugin.__name__} exports {sorted(unknown)}, not hooks of {sorted(_HOOKS)}" + ) + return plugin def prepare_model_for_loading( diff --git a/modelopt/torch/models/qwen4_exp/qsa_indexer.py b/modelopt/torch/models/qwen4_exp/qsa_indexer.py index 8d10bf4bc8b..be9080b88d4 100644 --- a/modelopt/torch/models/qwen4_exp/qsa_indexer.py +++ b/modelopt/torch/models/qwen4_exp/qsa_indexer.py @@ -17,22 +17,26 @@ import torch +from modelopt.torch.utils import warn_rank_0 from modelopt.torch.utils.network import bind_forward_method __all__ = ["prepare_for_calibration"] -def _forward(self, hidden_states, position_embeddings, attention_mask, past_key_values): +def _forward( + self, hidden_states, position_embeddings, attention_mask, past_key_values, *args, **kwargs +): """Transformers' indexer, short-circuited when every visible key fits the token budget. The reference loops in Python over every (sample, query) pair. A query keeps its ``block_topk`` best complete key blocks plus the incomplete tail, so when the keys hold no more blocks than that, it keeps every visible token: the selection is the visibility mask the - indexer was given, which this returns directly. + indexer was given, which this returns directly. A call with arguments this does not know, + from a newer Transformers, goes to the reference too. """ - if attention_mask.shape[-1] // self.compress_ratio > self.block_topk: + if args or kwargs or attention_mask.shape[-1] // self.compress_ratio > self.block_topk: return self._modelopt_reference_forward( - hidden_states, position_embeddings, attention_mask, past_key_values + hidden_states, position_embeddings, attention_mask, past_key_values, *args, **kwargs ) # The reference runs the projection on every forward, so a quantizer on it gets calibrated. qk = self.index_qk_proj(hidden_states) @@ -67,4 +71,10 @@ def prepare_for_calibration(model) -> bool: if not hasattr(module, "_modelopt_reference_forward"): # Keeps an Accelerate hook already on the indexer, such as device_map's, around the fast path. bind_forward_method(module, _forward, "_modelopt_reference_forward") + if not found: + warn_rank_0( + "No Qwen4ExpTextQSAIndexer found in this Qwen3.8 model, so calibration runs " + "Transformers' QSA indexer loop, which can take hours. If Transformers renamed the " + "class, update modelopt/torch/models/qwen4_exp/qsa_indexer.py." + ) return found diff --git a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py index 8420e37cc69..24389b5fcc0 100644 --- a/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py +++ b/tests/unit/torch/models/test_qwen4_exp_qsa_indexer.py @@ -16,6 +16,7 @@ """Qwen3.8's QSA indexer fast path against Transformers' reference loop.""" import copy +import sys import types import pytest @@ -26,6 +27,7 @@ from transformers.cache_utils import Cache, DynamicIndexedLayer from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig +from modelopt.torch.models import hf from modelopt.torch.models.hf import ( checkpoint_has_mtp, prefers_bundled_code, @@ -178,3 +180,41 @@ def test_loading_hooks_the_plugin_lacks_keep_their_defaults(model_type): """The Qwen3.8 plugin only has a calibration hook; the loading hooks fall back to the defaults.""" assert not prefers_bundled_code(model_type, None, "ckpt", {}) assert not checkpoint_has_mtp(model_type, "ckpt") + + +def test_qsa_indexer_hands_unknown_arguments_to_the_reference(): + """A newer Transformers may pass the indexer an argument the fast path does not know; the call + goes to the reference forward with it instead of raising ``TypeError``.""" + model, config = _model(64) + hidden, rope = _inputs(config, 1, 64) + prepare_model_for_calibration(model) + calls = [] + model.indexer._modelopt_reference_forward = lambda *a, **kw: calls.append(kw) or "reference" + + mask = _mask([("right", 0)], False, 64) + assert model.indexer(hidden, rope, mask, None, new_argument=1) == "reference" + assert calls == [{"new_argument": 1}] + + +def test_a_qwen3_8_model_without_a_known_indexer_warns(): + """If Transformers renames the indexer class, calibration warns rather than silently running + the loop.""" + model = nn.Module() + model.config = types.SimpleNamespace(model_type="qwen4_exp") + model.layer = nn.Linear(4, 4) + with pytest.warns(UserWarning, match="Qwen4ExpTextQSAIndexer"): + prepare_model_for_calibration(model) + + +def test_a_plugin_exporting_an_unknown_hook_is_refused(monkeypatch): + """Hooks are optional, so a misspelled one in a plugin's ``__all__`` raises instead of being + skipped.""" + plugin = types.ModuleType("modelopt.torch.models.misspelled_plugin") + plugin.__all__ = ["prepare_for_calibraton"] + plugin.prepare_for_calibraton = lambda model: None + monkeypatch.setitem(sys.modules, plugin.__name__, plugin) + monkeypatch.setitem(hf._PLUGINS, "misspelled", "misspelled_plugin") + model = nn.Module() + model.config = types.SimpleNamespace(model_type="misspelled") + with pytest.raises(ValueError, match="prepare_for_calibraton"): + prepare_model_for_calibration(model)