diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b3ee33c3cca1..5bdf1d3ef0f3 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,21 @@ 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"; 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: + 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..72fc1e024c5a 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 ( @@ -31,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 @@ -226,6 +228,39 @@ def forward( # --------------------------------------------------------------------------- +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 + alignment, so a quantised g_proj trips Linear.__init__'s assert under + tp_size > 1. __post_init__'s generic exclusion pass runs too late. + """ + 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 + + # 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 +333,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 +725,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..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: @@ -622,9 +623,54 @@ 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. Probe every + # expert index so a rule targeting any single expert + # (not only expert 0) is honoured. + # + # 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) + 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 c605b8e36d6e..dae115c2e7f7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -33,10 +33,13 @@ 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 - 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_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) 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..a0b8068521a9 --- /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: list[str]) -> SimpleNamespace: + 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: 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" + # 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() -> 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"]) + + 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() -> 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() -> 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..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 @@ -62,3 +64,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: Any) -> dict[str, Any]: + 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() -> 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 + 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() -> 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 + 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() -> None: + 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() -> None: + """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() -> 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"], + 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"