From 5404a8599eb06caa595063a12bf43971de6540ed Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Wed, 12 Aug 2026 18:27:28 +0200 Subject: [PATCH 1/3] [None][fix] Honour producer ignored layers for fused MoE and any quant format Signed-off-by: Joe Rowell --- tensorrt_llm/_torch/model_config.py | 21 +++- tensorrt_llm/_torch/models/modeling_laguna.py | 32 ++++++- tensorrt_llm/_torch/models/modeling_utils.py | 21 ++++ .../_torch/models/test_laguna_gproj_quant.py | 89 +++++++++++++++++ tests/unittest/_torch/test_hf_quant_config.py | 96 +++++++++++++++++++ 5 files changed, 252 insertions(+), 7 deletions(-) create mode 100644 tests/unittest/_torch/models/test_laguna_gproj_quant.py diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b3ee33c3cca1..d3e9e7a0defb 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -654,11 +654,8 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): default_exclude = ["*kv_b_proj*", "*k_b_proj*", "*eh_proj"] # Merge HF config's modules_to_not_convert with default exclude_modules - if hf_exclude_modules is not None: - quant_config.exclude_modules = list( - set(hf_exclude_modules + default_exclude)) - else: - quant_config.exclude_modules = default_exclude + merged = list(hf_exclude_modules or []) + default_exclude + quant_config.exclude_modules = list(dict.fromkeys(merged)) # MXFP4 checkpoints. elif hf_quant_config.get("quant_method") == "mxfp4": quant_config.quant_algo = ModelConfig.get_mxfp4_quant_algo( @@ -713,6 +710,20 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): dict.fromkeys(hf_exclude_modules + default_exclude)) else: quant_config.exclude_modules = default_exclude + + # Honour the producer's "leave these layers unquantized" list, whatever + # the quant format. Compressed-tensors spells it "ignore"; modelopt, + # HF fp8 and mxfp8 spell it "ignored_layers". These layers carry no + # quant scales in the checkpoint and must be built as bf16 (per-head + # g_proj whose out dim is not a multiple of the block alignment, and + # MoE experts the producer kept in bf16 for quality). Merged on top of + # whatever per-format defaults were set above. + producer_ignored = list(hf_quant_config.get("ignored_layers", []) or []) + producer_ignored += list(hf_quant_config.get("ignore", []) or []) + if producer_ignored: + existing = quant_config.exclude_modules or [] + quant_config.exclude_modules = list( + dict.fromkeys(list(existing) + producer_ignored)) return quant_config, layer_quant_config @staticmethod diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 03257d22cedf..987325002287 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -23,6 +23,7 @@ from transformers import PretrainedConfig from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata from ..attention_backend.interface import ( @@ -226,6 +227,29 @@ def forward( # --------------------------------------------------------------------------- +def g_proj_quant_config(model_config, layer_idx: Optional[int]) -> Optional[QuantConfig]: + """Resolve g_proj's quant config, dropping it if the checkpoint excludes it. + + g_proj's output dim is the head count, not a multiple of the FP8 block + alignment, so a quantised g_proj trips Linear.__init__'s assert under + tp_size > 1. __post_init__'s generic exclusion pass runs too late. + """ + quant_config = model_config.get_quant_config() + if quant_config is None or layer_idx is None: + return quant_config + + name = f"model.layers.{layer_idx}.self_attn.g_proj" + if not quant_config.is_module_excluded_from_quantization(name): + return quant_config + + # Keep kv_cache_quant_algo: QuantMode derives the KV-cache mode from it, so + # dropping it would silently disable FP8 KV cache for this module. + return QuantConfig( + quant_algo=None, + kv_cache_quant_algo=quant_config.kv_cache_quant_algo, + ) + + class LagunaAttention(QKNormRoPEAttention): """Laguna attention with per-head softplus gating, per-layer heads, dual RoPE, and sliding window.""" @@ -298,7 +322,7 @@ def __init__( dtype=config.torch_dtype, mapping=model_config.mapping, tensor_parallel_mode=g_tp_mode, - quant_config=model_config.get_quant_config(), + quant_config=g_proj_quant_config(model_config, layer_idx), ) @staticmethod @@ -690,9 +714,13 @@ def handle_special_instance_module( .replace("up_proj", "w3") .replace("down_proj", "w2") ) - if use_fp8_block_rename: + if use_fp8_block_rename and "weight_scale_inv" not in new_wn: # ModelOpt uses "weight_scale"; TRT-LLM DeepSeekFP8 # quantization expects "weight_scale_inv" (same values). + # Skip names already using "weight_scale_inv" + # (compressed-tensors / DeepSeek style): a blanket replace + # turns it into "weight_scale_inv_inv", which silently + # drops the block scales and leaves them uninitialised. new_wn = new_wn.replace("weight_scale", "weight_scale_inv") # filter_weights strips module prefix leaving "experts.X.w1.weight"; # VANILLA mode expects "X.w1.weight" (integer prefix only). diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1a140ea15e03..9ab8f14d4396 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -622,6 +622,27 @@ def apply_quant_config_exclude_modules(self): name.replace('qkv_proj', 'k_proj'), name.replace('qkv_proj', 'v_proj') ] + elif isinstance(module, (MoE, VanillaMoE)): + # Fused MoE: a checkpoint's ignore list names the + # per-expert weights (re:...experts.[0-9]+.gate_proj), + # but the fused module has no per-expert child, so such + # a rule never matches and those experts get quantized + # despite being bf16 in the checkpoint. Offer + # representative per-expert names, mirroring the Linear + # expansion above. + # + # Strip a trailing ".backend": ConfigurableMoE wraps the + # module that actually owns the weights, and excluding + # only the wrapper does nothing because create_weights() + # delegates to the backend. Without this the backend + # stays quantized and the output is still degenerate. + base = (name[:-len('.backend')] + if name.endswith('.backend') else name) + candidates += [ + f'{base}.0.gate_proj', + f'{base}.0.up_proj', + f'{base}.0.down_proj', + ] is_excluded = any( quant_config.is_module_excluded_from_quantization(n) for n in candidates) diff --git a/tests/unittest/_torch/models/test_laguna_gproj_quant.py b/tests/unittest/_torch/models/test_laguna_gproj_quant.py new file mode 100644 index 000000000000..ed3d014f8958 --- /dev/null +++ b/tests/unittest/_torch/models/test_laguna_gproj_quant.py @@ -0,0 +1,89 @@ +# 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. +"""g_proj must not be built quantised when the checkpoint excludes it. + +Only tp_size > 1 takes the sharding path, so this presents as a sharding +failure rather than a quantisation one. The existing Laguna accuracy tests are +all TP=1 and would not catch a regression. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.models.modeling_laguna import g_proj_quant_config +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +LAYER_IDX = 3 +G_PROJ = f"model.layers.{LAYER_IDX}.self_attn.g_proj" + + +def _model_config(exclude_modules): + quant_config = QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + kv_cache_quant_algo=QuantAlgo.FP8, + exclude_modules=exclude_modules, + ) + return SimpleNamespace(get_quant_config=lambda name=None: quant_config) + + +# Exact name, ancestor walk, ".*" subtree, glob and "re:" regex. The fleet uses +# literal names for g_proj but regexes for the fused-MoE exclusions. +@pytest.mark.parametrize( + "pattern", + [ + G_PROJ, + f"model.layers.{LAYER_IDX}.self_attn", + f"model.layers.{LAYER_IDX}", + f"model.layers.{LAYER_IDX}.*", + "*g_proj*", + r"re:.*\.self_attn\.g_proj$", + ], +) +def test_excluded_g_proj_is_not_quantised(pattern): + resolved = g_proj_quant_config(_model_config([pattern, "lm_head"]), LAYER_IDX) + + assert resolved.quant_algo is None, f"{pattern!r} did not exclude g_proj" + # KV cache mode is derived from this, so it must survive the override. + assert resolved.kv_cache_quant_algo == QuantAlgo.FP8 + + +def test_unexcluded_g_proj_keeps_the_model_quant_config(): + """Control: a checkpoint that does quantise g_proj is left alone.""" + model_config = _model_config(["lm_head", "model.layers.9.self_attn.g_proj"]) + + resolved = g_proj_quant_config(model_config, LAYER_IDX) + + assert resolved is model_config.get_quant_config() + assert resolved.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + + +def test_unquantised_model_is_passed_through(): + model_config = SimpleNamespace(get_quant_config=lambda name=None: None) + + assert g_proj_quant_config(model_config, LAYER_IDX) is None + + +def test_missing_layer_idx_does_not_match_a_literal_none(): + """layer_idx is Optional; without the guard the name would render as + "model.layers.None.self_attn.g_proj" and silently match nothing.""" + model_config = _model_config(["model.layers.None.self_attn.g_proj"]) + + resolved = g_proj_quant_config(model_config, None) + + assert resolved.quant_algo == QuantAlgo.FP8_BLOCK_SCALES diff --git a/tests/unittest/_torch/test_hf_quant_config.py b/tests/unittest/_torch/test_hf_quant_config.py index 0c51180d2cb3..d1fed5139e78 100644 --- a/tests/unittest/_torch/test_hf_quant_config.py +++ b/tests/unittest/_torch/test_hf_quant_config.py @@ -62,3 +62,99 @@ def test_load_hf_quant_config_parses_nvfp4_with_kv_cache_scheme(): assert quant_config.group_size == 16 assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 assert set(quant_config.exclude_modules) == {gate_exclude, "lm_head"} + + +def _fp8_block_scales_config(**overrides): + config = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + config.update(overrides) + return config + + +def test_fp8_block_scales_honours_ignored_layers(): + """vLLM/AutoFP8-style "fp8" configs record BF16 layers as "ignored_layers". + + These checkpoints carry no "modules_to_not_convert" at all, so before this + was read every exclusion was dropped and the layers were built as FP8. + """ + ignored = ["lm_head", "model.layers.0.self_attn.g_proj"] + hf_quant_config = _fp8_block_scales_config(ignored_layers=ignored) + + quant_config, _ = ModelConfig.load_hf_quant_config(hf_quant_config, moe_backend="CUTLASS") + + assert quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + for name in ignored: + assert name in quant_config.exclude_modules + assert quant_config.is_module_excluded_from_quantization(name) + # the built-in defaults are still applied + assert "*kv_b_proj*" in quant_config.exclude_modules + + +def test_fp8_block_scales_merges_both_exclusion_keys(): + """Both keys are honoured, and an entry named twice appears once. + + Producers commonly write the same module under both keys, and one of the + entries here also collides with a built-in default, so this fails if the + merge stops de-duplicating. + """ + # "lm_head" appears only under ignored_layers, so this still fails if that + # key is not read; "*kv_b_proj*" duplicates a built-in default. + hf_quant_config = _fp8_block_scales_config( + ignored_layers=["lm_head", "*kv_b_proj*"], + modules_to_not_convert=["model.layers.1.mlp.down_proj"], + ) + + quant_config, _ = ModelConfig.load_hf_quant_config(hf_quant_config, moe_backend="CUTLASS") + + excluded = quant_config.exclude_modules + assert "lm_head" in excluded + assert "model.layers.1.mlp.down_proj" in excluded + assert excluded.count("lm_head") == 1 + assert excluded.count("*kv_b_proj*") == 1 + assert len(excluded) == len(set(excluded)) + + +def test_fp8_block_scales_without_exclusions_keeps_defaults(): + quant_config, _ = ModelConfig.load_hf_quant_config( + _fp8_block_scales_config(), moe_backend="CUTLASS" + ) + + assert quant_config.exclude_modules == ["*kv_b_proj*", "*k_b_proj*", "*eh_proj"] + + +def test_compressed_tensors_honours_ignore_key(): + """Characterization: compressed-tensors keeps its "ignore" list. + + update_quant_config_from_compressed_tensors already merged "ignore", so + this passes with or without the general merge below. Kept to pin that the + hoisted merge does not REGRESS the path that already worked, including the + per-expert regex form a hybrid checkpoint uses. + """ + ignore = [ + "lm_head", + r"re:^model\.layers\.4[0-7]\.mlp\.experts\.[0-9]+\.(gate_proj|up_proj|down_proj)$", + ] + hf_quant_config = _compressed_tensors_nvfp4_config(ignore=ignore) + + quant_config, _ = ModelConfig.load_hf_quant_config(hf_quant_config, moe_backend="CUTLASS") + + for entry in ignore: + assert entry in quant_config.exclude_modules, f"{entry!r} was dropped" + + +def test_ignore_and_ignored_layers_are_both_merged(): + """A config carrying both spellings keeps both, de-duplicated.""" + hf_quant_config = _fp8_block_scales_config( + ignored_layers=["model.layers.3.self_attn.g_proj", "lm_head"], + ignore=["model.layers.9.mlp.experts", "lm_head"], + ) + + quant_config, _ = ModelConfig.load_hf_quant_config(hf_quant_config, moe_backend="CUTLASS") + + ex = quant_config.exclude_modules + for entry in ("model.layers.3.self_attn.g_proj", "model.layers.9.mlp.experts", "lm_head"): + assert entry in ex, f"{entry!r} missing" + assert ex.count("lm_head") == 1, "duplicate not collapsed" From 863f3551535aeb4684cf9a1e8ca5caf448edefb5 Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Thu, 13 Aug 2026 12:15:02 +0200 Subject: [PATCH 2/3] [None][chore] Annotate new functions and register the CPU tests Signed-off-by: Joe Rowell --- tensorrt_llm/_torch/models/modeling_laguna.py | 5 ++++- tests/integration/test_lists/test-db/l0_cpu.yml | 2 ++ .../_torch/models/test_laguna_gproj_quant.py | 10 +++++----- tests/unittest/_torch/test_hf_quant_config.py | 14 ++++++++------ 4 files changed, 19 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 987325002287..2183dd18cfa8 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -32,6 +32,7 @@ RopeParams, ) from ..distributed import AllReduce, AllReduceParams +from ..model_config import ModelConfig from ..modules.attention import _helix_cp_allgather_input, _helix_cp_output_projection from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding @@ -227,7 +228,9 @@ def forward( # --------------------------------------------------------------------------- -def g_proj_quant_config(model_config, layer_idx: Optional[int]) -> Optional[QuantConfig]: +def g_proj_quant_config( + model_config: ModelConfig[PretrainedConfig], layer_idx: Optional[int] +) -> Optional[QuantConfig]: """Resolve g_proj's quant config, dropping it if the checkpoint excludes it. g_proj's output dim is the head count, not a multiple of the FP8 block diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c605b8e36d6e..2d3222b7d8fb 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -33,10 +33,12 @@ l0_cpu: - unittest/_torch/modeling - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/checkpoints + - unittest/_torch/models/test_laguna_gproj_quant.py - unittest/_torch/modules - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/test_hf_quant_config.py - unittest/_torch/test_mmap_utils.py - unittest/_torch/test_model_config.py - unittest/_torch/thop/parallel_hw_agnostic diff --git a/tests/unittest/_torch/models/test_laguna_gproj_quant.py b/tests/unittest/_torch/models/test_laguna_gproj_quant.py index ed3d014f8958..a0b8068521a9 100644 --- a/tests/unittest/_torch/models/test_laguna_gproj_quant.py +++ b/tests/unittest/_torch/models/test_laguna_gproj_quant.py @@ -33,7 +33,7 @@ G_PROJ = f"model.layers.{LAYER_IDX}.self_attn.g_proj" -def _model_config(exclude_modules): +def _model_config(exclude_modules: list[str]) -> SimpleNamespace: quant_config = QuantConfig( quant_algo=QuantAlgo.FP8_BLOCK_SCALES, kv_cache_quant_algo=QuantAlgo.FP8, @@ -55,7 +55,7 @@ def _model_config(exclude_modules): r"re:.*\.self_attn\.g_proj$", ], ) -def test_excluded_g_proj_is_not_quantised(pattern): +def test_excluded_g_proj_is_not_quantised(pattern: str) -> None: resolved = g_proj_quant_config(_model_config([pattern, "lm_head"]), LAYER_IDX) assert resolved.quant_algo is None, f"{pattern!r} did not exclude g_proj" @@ -63,7 +63,7 @@ def test_excluded_g_proj_is_not_quantised(pattern): assert resolved.kv_cache_quant_algo == QuantAlgo.FP8 -def test_unexcluded_g_proj_keeps_the_model_quant_config(): +def test_unexcluded_g_proj_keeps_the_model_quant_config() -> None: """Control: a checkpoint that does quantise g_proj is left alone.""" model_config = _model_config(["lm_head", "model.layers.9.self_attn.g_proj"]) @@ -73,13 +73,13 @@ def test_unexcluded_g_proj_keeps_the_model_quant_config(): assert resolved.quant_algo == QuantAlgo.FP8_BLOCK_SCALES -def test_unquantised_model_is_passed_through(): +def test_unquantised_model_is_passed_through() -> None: model_config = SimpleNamespace(get_quant_config=lambda name=None: None) assert g_proj_quant_config(model_config, LAYER_IDX) is None -def test_missing_layer_idx_does_not_match_a_literal_none(): +def test_missing_layer_idx_does_not_match_a_literal_none() -> None: """layer_idx is Optional; without the guard the name would render as "model.layers.None.self_attn.g_proj" and silently match nothing.""" model_config = _model_config(["model.layers.None.self_attn.g_proj"]) diff --git a/tests/unittest/_torch/test_hf_quant_config.py b/tests/unittest/_torch/test_hf_quant_config.py index d1fed5139e78..58166f55aaec 100644 --- a/tests/unittest/_torch/test_hf_quant_config.py +++ b/tests/unittest/_torch/test_hf_quant_config.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + import pytest from tensorrt_llm._torch.model_config import ModelConfig @@ -64,7 +66,7 @@ def test_load_hf_quant_config_parses_nvfp4_with_kv_cache_scheme(): assert set(quant_config.exclude_modules) == {gate_exclude, "lm_head"} -def _fp8_block_scales_config(**overrides): +def _fp8_block_scales_config(**overrides: Any) -> dict[str, Any]: config = { "quant_method": "fp8", "activation_scheme": "dynamic", @@ -74,7 +76,7 @@ def _fp8_block_scales_config(**overrides): return config -def test_fp8_block_scales_honours_ignored_layers(): +def test_fp8_block_scales_honours_ignored_layers() -> None: """vLLM/AutoFP8-style "fp8" configs record BF16 layers as "ignored_layers". These checkpoints carry no "modules_to_not_convert" at all, so before this @@ -93,7 +95,7 @@ def test_fp8_block_scales_honours_ignored_layers(): assert "*kv_b_proj*" in quant_config.exclude_modules -def test_fp8_block_scales_merges_both_exclusion_keys(): +def test_fp8_block_scales_merges_both_exclusion_keys() -> None: """Both keys are honoured, and an entry named twice appears once. Producers commonly write the same module under both keys, and one of the @@ -117,7 +119,7 @@ def test_fp8_block_scales_merges_both_exclusion_keys(): assert len(excluded) == len(set(excluded)) -def test_fp8_block_scales_without_exclusions_keeps_defaults(): +def test_fp8_block_scales_without_exclusions_keeps_defaults() -> None: quant_config, _ = ModelConfig.load_hf_quant_config( _fp8_block_scales_config(), moe_backend="CUTLASS" ) @@ -125,7 +127,7 @@ def test_fp8_block_scales_without_exclusions_keeps_defaults(): assert quant_config.exclude_modules == ["*kv_b_proj*", "*k_b_proj*", "*eh_proj"] -def test_compressed_tensors_honours_ignore_key(): +def test_compressed_tensors_honours_ignore_key() -> None: """Characterization: compressed-tensors keeps its "ignore" list. update_quant_config_from_compressed_tensors already merged "ignore", so @@ -145,7 +147,7 @@ def test_compressed_tensors_honours_ignore_key(): assert entry in quant_config.exclude_modules, f"{entry!r} was dropped" -def test_ignore_and_ignored_layers_are_both_merged(): +def test_ignore_and_ignored_layers_are_both_merged() -> None: """A config carrying both spellings keeps both, de-duplicated.""" hf_quant_config = _fp8_block_scales_config( ignored_layers=["model.layers.3.self_attn.g_proj", "lm_head"], From 5a217878002a77833d84c911f48596012602ddf0 Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Wed, 19 Aug 2026 11:42:54 +0100 Subject: [PATCH 3/3] fused moe: honour ignore rules for every expert, warn on partial list Signed-off-by: Joe Rowell --- tensorrt_llm/_torch/model_config.py | 13 ++- tensorrt_llm/_torch/models/modeling_laguna.py | 14 ++- tensorrt_llm/_torch/models/modeling_utils.py | 49 ++++++-- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../models/test_fused_moe_ignored_layers.py | 105 ++++++++++++++++++ 5 files changed, 161 insertions(+), 21 deletions(-) create mode 100644 tests/unittest/_torch/models/test_fused_moe_ignored_layers.py diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index d3e9e7a0defb..5bdf1d3ef0f3 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -712,12 +712,13 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): quant_config.exclude_modules = default_exclude # Honour the producer's "leave these layers unquantized" list, whatever - # the quant format. Compressed-tensors spells it "ignore"; modelopt, - # HF fp8 and mxfp8 spell it "ignored_layers". These layers carry no - # quant scales in the checkpoint and must be built as bf16 (per-head - # g_proj whose out dim is not a multiple of the block alignment, and - # MoE experts the producer kept in bf16 for quality). Merged on top of - # whatever per-format defaults were set above. + # the quant format. Compressed-tensors spells it "ignore"; HF fp8 and + # mxfp8 spell it "ignored_layers" (modelopt returns early above, so it + # never reaches here). These layers carry no quant scales in the + # checkpoint and must be built as bf16 (per-head g_proj whose out dim is + # not a multiple of the block alignment, and MoE experts the producer + # kept in bf16 for quality). Merged on top of whatever per-format + # defaults were set above. producer_ignored = list(hf_quant_config.get("ignored_layers", []) or []) producer_ignored += list(hf_quant_config.get("ignore", []) or []) if producer_ignored: diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 2183dd18cfa8..72fc1e024c5a 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -237,11 +237,19 @@ def g_proj_quant_config( alignment, so a quantised g_proj trips Linear.__init__'s assert under tp_size > 1. __post_init__'s generic exclusion pass runs too late. """ - quant_config = model_config.get_quant_config() - if quant_config is None or layer_idx is None: - return quant_config + if layer_idx is None: + return model_config.get_quant_config() name = f"model.layers.{layer_idx}.self_attn.g_proj" + # Resolve this layer's own config for a MIXED_PRECISION checkpoint; fall + # back to the global config when there is no per-layer entry for g_proj. + try: + quant_config = model_config.get_quant_config(name) + except ValueError: + quant_config = model_config.get_quant_config() + if quant_config is None: + return quant_config + if not quant_config.is_module_excluded_from_quantization(name): return quant_config diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 9ab8f14d4396..9713bfd8b9d2 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -605,6 +605,7 @@ def apply_quant_config_exclude_modules(self): if quant_config.exclude_modules is not None: for name, module in self.named_modules(): candidates = [name] + moe_expert_exclusions = None if isinstance(module, Linear): weight_mode = module.weights_loading_config.weight_mode if weight_mode == WeightMode.FUSED_GATE_UP_LINEAR: @@ -627,25 +628,49 @@ def apply_quant_config_exclude_modules(self): # per-expert weights (re:...experts.[0-9]+.gate_proj), # but the fused module has no per-expert child, so such # a rule never matches and those experts get quantized - # despite being bf16 in the checkpoint. Offer - # representative per-expert names, mirroring the Linear - # expansion above. + # despite being bf16 in the checkpoint. Probe every + # expert index so a rule targeting any single expert + # (not only expert 0) is honoured. # - # Strip a trailing ".backend": ConfigurableMoE wraps the - # module that actually owns the weights, and excluding - # only the wrapper does nothing because create_weights() - # delegates to the backend. Without this the backend - # stays quantized and the output is still degenerate. + # Strip a trailing ".backend" for these synthesized + # candidates: ConfigurableMoE wraps the module that owns + # the weights, so a synthesized name must point at the + # backend (create_weights() delegates to it). A plain + # rule naming the wrapper already reaches the backend via + # is_module_excluded_from_quantization's ancestor walk, + # so the strip is only needed here. base = (name[:-len('.backend')] if name.endswith('.backend') else name) - candidates += [ - f'{base}.0.gate_proj', - f'{base}.0.up_proj', - f'{base}.0.down_proj', + num_experts = getattr(module, 'num_experts', None) + expert_ids = (range(num_experts) + if isinstance(num_experts, int) + and num_experts > 0 else [0]) + moe_expert_exclusions = [ + any( + quant_config. + is_module_excluded_from_quantization( + f'{base}.{e}.{proj}') + for proj in ('gate_proj', 'up_proj', + 'down_proj')) for e in expert_ids ] is_excluded = any( quant_config.is_module_excluded_from_quantization(n) for n in candidates) + if moe_expert_exclusions is not None: + n_excluded = sum(moe_expert_exclusions) + if 0 < n_excluded < len(moe_expert_exclusions): + # A fused MoE shares one quant config across all + # experts, so a partial ignore list cannot be + # honoured per-expert. Treat the module as excluded + # (never leave a bf16 expert quantized) but warn so + # the partial list is visible rather than silently + # rounded to all-or-nothing. + logger.warning( + f"Quantization ignore list excludes " + f"{n_excluded}/{len(moe_expert_exclusions)} " + f"experts of fused module '{base}'; treating " + f"the whole module as excluded.") + is_excluded = is_excluded or n_excluded > 0 if is_excluded and getattr(module, "quant_config", None) is not None: module.quant_config = new_config diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 2d3222b7d8fb..dae115c2e7f7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -33,6 +33,7 @@ l0_cpu: - unittest/_torch/modeling - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/checkpoints + - unittest/_torch/models/test_fused_moe_ignored_layers.py - unittest/_torch/models/test_laguna_gproj_quant.py - unittest/_torch/modules - unittest/_torch/multimodal diff --git a/tests/unittest/_torch/models/test_fused_moe_ignored_layers.py b/tests/unittest/_torch/models/test_fused_moe_ignored_layers.py new file mode 100644 index 000000000000..b034338fbd16 --- /dev/null +++ b/tests/unittest/_torch/models/test_fused_moe_ignored_layers.py @@ -0,0 +1,105 @@ +# 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. +"""A producer's per-expert ignore rule must de-quantize the fused MoE. + +The fused module has no per-expert child, so a rule naming an expert's weights +(``...experts..up_proj``) only matches synthesized candidates. The pass must +probe every expert index (not just expert 0), strip a ``.backend`` wrapper, and +warn rather than silently round a partial list to all-or-nothing. +""" + +from types import SimpleNamespace + +import pytest +import torch.nn as nn + +from tensorrt_llm._torch.models import modeling_utils +from tensorrt_llm._torch.modules.fused_moe import MoE +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +NUM_EXPERTS = 8 +BASE = "model.layers.3.mlp.experts" + + +class _FakeMoE(MoE): + """Minimal MoE stand-in exposing only what the exclusion pass reads.""" + + def __init__(self, num_experts: int) -> None: + nn.Module.__init__(self) + self.num_experts = num_experts + self.quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + self._weights_created = True + + +def _attach(root: nn.Module, path: str, leaf: nn.Module) -> None: + """Register ``leaf`` under ``root`` at a dotted path, so ``named_modules`` + yields ``path`` (plain nn.Modules stand in for the intermediate wrappers).""" + parent = root + parts = path.split(".") + for part in parts[:-1]: + child = parent._modules.get(part) + if child is None: + child = nn.Module() + parent.add_module(part, child) + parent = child + parent.add_module(parts[-1], leaf) + + +def _apply_exclusions(moe_path: str, exclude_modules: list[str]) -> _FakeMoE: + moe = _FakeMoE(NUM_EXPERTS) + model = nn.Module() + _attach(model, moe_path, moe) + model.model_config = SimpleNamespace( + quant_config=QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, exclude_modules=list(exclude_modules) + ) + ) + modeling_utils.DecoderModelForCausalLM.apply_quant_config_exclude_modules(model) + return moe + + +def test_exact_nonzero_expert_index_dequantizes() -> None: + """Regression: an exact rule for a nonzero expert must de-quantize the + module. The old code only probed expert 0, so such rules never matched.""" + last = NUM_EXPERTS - 1 + moe = _apply_exclusions(BASE, [f"{BASE}.{last}.up_proj", "lm_head"]) + + assert moe.quant_config.quant_algo is None + assert moe._weights_created is False + + +def test_backend_wrapped_moe_is_dequantized() -> None: + """ConfigurableMoE wraps the weight-owning backend; the ``.backend`` suffix + must be stripped so the producer's expert rule reaches the backend.""" + moe = _apply_exclusions(f"{BASE}.backend", [f"{BASE}.5.down_proj", "lm_head"]) + + assert moe.quant_config.quant_algo is None + assert moe._weights_created is False + + +def test_partial_expert_list_warns_and_dequantizes(monkeypatch: pytest.MonkeyPatch) -> None: + """A fused module shares one quant config, so a partial list cannot be + honoured per-expert: warn, and treat the whole module as excluded.""" + warnings: list[str] = [] + monkeypatch.setattr(modeling_utils.logger, "warning", lambda msg, *a, **k: warnings.append(msg)) + + exclude = [f"{BASE}.{e}.up_proj" for e in range(3)] + ["lm_head"] + moe = _apply_exclusions(BASE, exclude) + + assert moe.quant_config.quant_algo is None + assert any("fused module" in w for w in warnings)