diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index b71fe8495928..5795074ac7e8 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -413,39 +413,55 @@ def _normalize_qwen35_moe_vl_config(model_config) -> None: _normalize_qwen35_vl_config(model_config, inner_arch="Qwen3_5MoeForCausalLM") -def _lm_head_nvfp4_enabled(model_config): +_LM_HEAD_FP8_ALGOS = (QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN,) + + +def _lm_head_quant_enabled(model_config: ModelConfig) -> bool: """Whether the checkpoint's quantized lm_head should stay quantized. - ModelOpt MIXED_PRECISION exports for Qwen3.5/3.6 quantize lm_head to - W4A16_NVFP4 (packed FP4 weight + per-group FP8 scales). On SM100/103 the - NVFP4 (W4A4) Linear path can consume it directly, cutting the lm_head - GEMM's weight traffic 4x vs the bf16 dequant fallback -- the decode - lm_head is purely weight-bandwidth-bound. Conditions mirror what the - quantized LMHead supports (see LMHead.__init__ guards) plus the paths - that bypass the Linear machinery entirely: - - - tie_word_embeddings shares the weight with the embedding lookup, which - needs a dense bf16 weight; - - ADP builds a TP-less LMHead (or slices the raw weight for the - spec-decoding head), assuming bf16; - - COLUMN TP pads the vocab shard when it doesn't divide evenly, which the - quantized path does not support (vocab 248320 divides all common tp). - - Must be evaluated BEFORE _normalize_qwen35_quant_config_dict runs (it - promotes the entry to NVFP4 or drops it). + ModelOpt W4A16_NVFP4 heads stay quantized only on SM100/103. Mixed-group + FP8 heads use the checkpoint's explicit per-layer entry on every SM. + Unsupported FP8 configurations fail here because no loader path currently + dequantizes their weight and per-channel scale to a dense weight. """ qcd = getattr(model_config, "quant_config_dict", None) or {} cfg = qcd.get("lm_head") + if cfg is None: + return False + if cfg.quant_algo == QuantAlgo.W4A16_NVFP4: + if get_sm_version() not in (100, 103): + return False + is_fp8 = False + elif cfg.quant_algo in _LM_HEAD_FP8_ALGOS: + is_fp8 = True + elif cfg.quant_algo in (QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"FP8 lm_head algorithm {cfg.quant_algo} is unsupported; " + "falling back to bf16 would discard its quantization scales." + ) + else: + return False + pretrained = model_config.pretrained_config mapping = model_config.mapping - return ( - cfg is not None - and cfg.quant_algo == QuantAlgo.W4A16_NVFP4 - and get_sm_version() in (100, 103) - and not getattr(pretrained, "tie_word_embeddings", False) - and not mapping.enable_attention_dp - and getattr(pretrained, "vocab_size", 0) % mapping.tp_size == 0 - ) + unsupported_reason = None + if getattr(pretrained, "tie_word_embeddings", False): + unsupported_reason = "tied embeddings require the dense embedding weight" + elif mapping.enable_attention_dp and getattr(mapping, "enable_lm_head_tp_in_adp", False): + unsupported_reason = "lm_head TP in attention-DP slices the raw weight" + elif not mapping.enable_attention_dp: + vocab_size = getattr(pretrained, "vocab_size", None) + if vocab_size is None: + unsupported_reason = "vocab_size is required for quantized LMHead" + elif vocab_size % mapping.tp_size != 0: + unsupported_reason = "vocab padding is unsupported for quantized LMHead" + + if unsupported_reason is not None and is_fp8: + raise ValueError( + f"FP8 lm_head cannot fall back to bf16: {unsupported_reason}; " + "dequantizing FP8 lm_head weights is not implemented." + ) + return unsupported_reason is None def _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=False): @@ -458,7 +474,7 @@ def _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=False): function translates the patterns so that ``apply_quant_config_exclude_modules`` can match them. - ``keep_lm_head_quant`` (see _lm_head_nvfp4_enabled) skips the lm_head + ``keep_lm_head_quant`` (see _lm_head_quant_enabled) skips the lm_head force-exclusion so the quantized LMHead path can engage. """ qc = model_config.quant_config @@ -497,7 +513,7 @@ def _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=False): # By default LMHead allocates an unquantized (bf16) weight, so a quantized # lm_head (e.g. NVFP4 in some ModelOpt MIXED_PRECISION exports) must be # excluded from quant and the weight mapper dequantizes it to bf16. When - # the quantized LMHead path is enabled (_lm_head_nvfp4_enabled), lm_head + # the quantized LMHead path is enabled (_lm_head_quant_enabled), lm_head # must NOT be excluded so DecoderModelForCausalLM builds it quantized. if not keep_lm_head_quant: normalized.add("lm_head") @@ -533,10 +549,13 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): shared scale (_requantize_linear_attn_fp8_qkvz). Incomplete or non-FP8 sets get no fused entry, and the mapper dequantizes them to bf16 instead. - The ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 when - ``keep_lm_head_quant`` (see _lm_head_nvfp4_enabled) and dropped otherwise: - a leftover entry would make DecoderModelForCausalLM build a quantized - LMHead whose weights the mapper had already dequantized to bf16. + The ``lm_head`` entry is kept when ``keep_lm_head_quant`` (see + _lm_head_quant_enabled) and dropped otherwise: a leftover entry would make + DecoderModelForCausalLM build a quantized LMHead whose weights the mapper + had already dequantized to bf16. A kept W4A16_NVFP4 entry is promoted to + NVFP4 on SM100/103; kept FP8 entries (compressed-tensors mixed-group + checkpoints) stay on their checkpoint algorithm, which the FP8 Linear + methods load directly. """ qcd = getattr(model_config, "quant_config_dict", None) if not qcd: @@ -554,16 +573,17 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): continue if name == "lm_head": if keep_lm_head_quant: - normalized[name] = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) + # Only W4A16_NVFP4 is promoted -- to the SM100/103 NVFP4 (W4A4) + # Linear path _lm_head_quant_enabled gated on. FP8 entries keep + # the checkpoint's algorithm: their Linear methods consume + # lm_head.weight + lm_head.weight_scale as stored. + if convert_to_nvfp4 and cfg.quant_algo == QuantAlgo.W4A16_NVFP4: + cfg = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) + normalized[name] = cfg else: - # Make the fallback visible: the checkpoint quantizes lm_head - # but this configuration can't keep it quantized (see - # _lm_head_nvfp4_enabled), so it dequantizes to bf16 at load. logger.info( f"lm_head quant entry ({cfg.quant_algo}) dropped: " - "unsupported configuration for quantized LMHead " - "(requires SM100/103, untied embeddings, no attention-DP, " - "vocab divisible by tp_size); lm_head runs bf16" + "no compatible quantized LMHead path" ) continue from_mtp = name.startswith("mtp.") @@ -645,7 +665,7 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: return _get_qwen35_moe_model_defaults(llm_args) def __init__(self, model_config): - keep_lm_head_quant = _lm_head_nvfp4_enabled(model_config) + keep_lm_head_quant = _lm_head_quant_enabled(model_config) _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=keep_lm_head_quant) super().__init__(model_config) @@ -662,7 +682,7 @@ class Qwen3_5ForCausalLM(Qwen3NextForCausalLM): """ def __init__(self, model_config): - keep_lm_head_quant = _lm_head_nvfp4_enabled(model_config) + keep_lm_head_quant = _lm_head_quant_enabled(model_config) _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=keep_lm_head_quant) super().__init__(model_config) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 663835918d05..0eb712720234 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -32,6 +32,7 @@ l0_cpu: - unittest/_torch/memory - unittest/_torch/modeling - unittest/_torch/models/checkpoints + - unittest/_torch/models/test_qwen3_5_lm_head_quant.py - unittest/_torch/modules - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py diff --git a/tests/unittest/_torch/models/test_qwen3_5_lm_head_quant.py b/tests/unittest/_torch/models/test_qwen3_5_lm_head_quant.py new file mode 100644 index 000000000000..723e896b54b7 --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_5_lm_head_quant.py @@ -0,0 +1,278 @@ +# 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.5/3.8 lm_head quantization decisions.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm._torch.models.modeling_qwen3_5 import ( + _lm_head_quant_enabled, + _normalize_qwen35_exclude_modules, + _normalize_qwen35_quant_config_dict, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +NUM_HIDDEN_LAYERS = 64 +# Qwen3.8-27B; divisible by every common tp_size, so COLUMN TP never pads. +VOCAB_SIZE = 248320 + +# A peer entry that must survive every lm_head decision unchanged. +PEER_KEY = "model.language_model.layers.0.self_attn.q_proj" +PEER_KEY_NORMALIZED = "model.layers.0.self_attn.q_proj" + +# What llm-compressor's ``ignore`` list looks like for these checkpoints: never +# mentions lm_head (the checkpoint quantizes it), so the exclusion below is the +# one the normalizer adds. +CHECKPOINT_IGNORE = ["model.language_model.layers.0.linear_attn.in_proj_a"] + +# The rowwise FP8 spelling used by compressed-tensors mixed-group checkpoints. +FP8_ALGOS = [QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN] + +BLACKWELL_SMS = [100, 103] +NON_BLACKWELL_SMS = [89, 90, 120] + + +def _model_config( + lm_head_algo, + *, + tp_size: int = 1, + vocab_size: int = VOCAB_SIZE, + tie_word_embeddings: bool = False, + enable_attention_dp: bool = False, + enable_lm_head_tp_in_adp: bool = False, +) -> SimpleNamespace: + """Build the minimal model_config the lm_head normalization pass reads.""" + quant_config_dict = {PEER_KEY: QuantConfig(quant_algo=QuantAlgo.FP8)} + if lm_head_algo is not None: + quant_config_dict["lm_head"] = QuantConfig(quant_algo=lm_head_algo) + return SimpleNamespace( + quant_config=QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, + kv_cache_quant_algo=QuantAlgo.FP8, + exclude_modules=list(CHECKPOINT_IGNORE), + ), + quant_config_dict=quant_config_dict, + pretrained_config=SimpleNamespace( + num_hidden_layers=NUM_HIDDEN_LAYERS, + vocab_size=vocab_size, + tie_word_embeddings=tie_word_embeddings, + ), + mapping=SimpleNamespace( + tp_size=tp_size, + enable_attention_dp=enable_attention_dp, + enable_lm_head_tp_in_adp=enable_lm_head_tp_in_adp, + ), + ) + + +def _normalize(model_config) -> bool: + """Run the pass both Qwen3.5 entry points run, and report the decision.""" + keep_lm_head_quant = _lm_head_quant_enabled(model_config) + _normalize_qwen35_exclude_modules(model_config, keep_lm_head_quant=keep_lm_head_quant) + _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=keep_lm_head_quant) + return keep_lm_head_quant + + +def _assert_lm_head_quantized(model_config, expected_algo: QuantAlgo) -> None: + """The normalized config must make the generic path build a quantized head.""" + assert model_config.quant_config_dict["lm_head"].quant_algo == expected_algo + assert "lm_head" not in model_config.quant_config.exclude_modules + resolved = DecoderModelForCausalLM._resolve_lm_head_quant_config(model_config) + assert resolved is not None + assert resolved.quant_algo == expected_algo + + +def _assert_lm_head_bf16(model_config) -> None: + """The normalized config must leave LMHead unquantized.""" + assert "lm_head" not in model_config.quant_config_dict + assert "lm_head" in model_config.quant_config.exclude_modules + assert DecoderModelForCausalLM._resolve_lm_head_quant_config(model_config) is None + + +def _assert_peer_entry_intact(model_config) -> None: + """lm_head handling must not disturb the rest of quant_config_dict.""" + peer = model_config.quant_config_dict[PEER_KEY_NORMALIZED] + assert peer.quant_algo == QuantAlgo.FP8 + + +# --- FP8 lm_head (compressed-tensors mixed-group checkpoints) ---------------- +# +# The checkpoint stores lm_head.weight as e4m3 plus a per-channel +# lm_head.weight_scale and no packed-FP4 tensors, so nothing can dequantize it +# to bf16: the entry has to be kept, on any SM, with its algorithm untouched. + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS + NON_BLACKWELL_SMS) +@pytest.mark.parametrize("algo", FP8_ALGOS) +def test_fp8_lm_head_is_retained_unchanged(algo, sm_version): + model_config = _model_config(algo) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + assert _normalize(model_config) + + # Kept on the checkpoint's own algorithm -- no NVFP4 promotion. + _assert_lm_head_quantized(model_config, algo) + _assert_peer_entry_intact(model_config) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("algo", FP8_ALGOS) +def test_fp8_lm_head_retained_across_tp_sizes(algo, tp_size): + # vocab_size divides every tp_size here, so COLUMN TP never pads the vocab + # shard -- the case the quantized LMHead rejects at construction. + model_config = _model_config(algo, tp_size=tp_size) + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=100): + assert _normalize(model_config) + + _assert_lm_head_quantized(model_config, algo) + + +# --- Unsupported FP8 configurations fail instead of dropping the scale ------ + + +@pytest.mark.parametrize( + "case,kwargs", + [ + ("tied embeddings share the dense weight", {"tie_word_embeddings": True}), + ( + "lm_head TP in attention-DP slices the raw weight", + {"enable_attention_dp": True, "enable_lm_head_tp_in_adp": True}, + ), + ("vocab does not divide tp_size (COLUMN TP pads)", {"vocab_size": 250, "tp_size": 4}), + ], +) +@pytest.mark.parametrize("algo", FP8_ALGOS) +def test_unsupported_fp8_configuration_raises(algo, case, kwargs): + model_config = _model_config(algo, **kwargs) + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=100): + with pytest.raises(ValueError, match="FP8 lm_head cannot fall back to bf16"): + _normalize(model_config) + + +def test_fp8_lm_head_requires_vocab_size(): + model_config = _model_config(QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN) + del model_config.pretrained_config.vocab_size + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=100): + with pytest.raises(ValueError, match="vocab_size is required"): + _normalize(model_config) + + +@pytest.mark.parametrize("algo", FP8_ALGOS + [QuantAlgo.W4A16_NVFP4]) +def test_attention_dp_without_lm_head_tp_keeps_quantized_head(algo): + model_config = _model_config(algo, enable_attention_dp=True) + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=100): + assert _normalize(model_config) + + expected_algo = QuantAlgo.NVFP4 if algo == QuantAlgo.W4A16_NVFP4 else algo + _assert_lm_head_quantized(model_config, expected_algo) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"tie_word_embeddings": True}, + {"enable_attention_dp": True, "enable_lm_head_tp_in_adp": True}, + {"vocab_size": 250, "tp_size": 4}, + ], +) +def test_unsupported_w4a16_configuration_falls_back_to_bf16(kwargs): + model_config = _model_config(QuantAlgo.W4A16_NVFP4, **kwargs) + + with patch("tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=100): + assert not _normalize(model_config) + + _assert_lm_head_bf16(model_config) + _assert_peer_entry_intact(model_config) + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS + NON_BLACKWELL_SMS) +@pytest.mark.parametrize("algo", [QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES]) +def test_unsupported_fp8_lm_head_algo_raises(algo, sm_version): + model_config = _model_config(algo) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + with pytest.raises(ValueError, match="FP8 lm_head algorithm"): + _normalize(model_config) + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS + NON_BLACKWELL_SMS) +@pytest.mark.parametrize("algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_AWQ, QuantAlgo.INT8]) +def test_other_lm_head_algorithms_keep_existing_fallback(algo, sm_version): + model_config = _model_config(algo) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + assert not _normalize(model_config) + + _assert_lm_head_bf16(model_config) + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS + NON_BLACKWELL_SMS) +def test_bf16_lm_head_checkpoint_is_untouched(sm_version): + # No lm_head entry at all (checkpoint left it bf16): lm_head stays excluded + # from quantization exactly as before. + model_config = _model_config(None) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + assert not _normalize(model_config) + + _assert_lm_head_bf16(model_config) + _assert_peer_entry_intact(model_config) + + +# --- W4A16_NVFP4 lm_head (ModelOpt MIXED_PRECISION) must not regress -------- + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS) +def test_w4a16_nvfp4_lm_head_promoted_on_blackwell(sm_version): + model_config = _model_config(QuantAlgo.W4A16_NVFP4) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + assert _normalize(model_config) + + # Promoted to the NVFP4 (W4A4) Linear path, as before. + _assert_lm_head_quantized(model_config, QuantAlgo.NVFP4) + _assert_peer_entry_intact(model_config) + + +@pytest.mark.parametrize("sm_version", NON_BLACKWELL_SMS) +def test_w4a16_nvfp4_lm_head_falls_back_off_blackwell(sm_version): + model_config = _model_config(QuantAlgo.W4A16_NVFP4) + + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + assert not _normalize(model_config) + + # No W4A4 path here: the entry is dropped and the weight mapper + # dequantizes the packed FP4 weight to bf16. + _assert_lm_head_bf16(model_config)