From 4a7dd87011a1f9dd3111bb17cc0c493659078a9f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 14 Aug 2026 14:42:52 -0700 Subject: [PATCH] [#17723][fix] Re-path dense Qwen quantization entries Signed-off-by: Mihai Chiorean --- .../_torch/models/modeling_qwen3_5.py | 43 +++- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../models/test_qwen3_5_dense_mlp_quant.py | 221 ++++++++++++++++++ 3 files changed, 253 insertions(+), 12 deletions(-) create mode 100644 tests/unittest/_torch/models/test_qwen3_5_dense_mlp_quant.py diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index b71fe8495928..60a22cb58627 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -516,12 +516,13 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): linear_attn.out_proj) silently fall back to the MIXED_PRECISION global config -> unquantized, and their quantized checkpoint weights fail to load. - On SM100/SM103, W4A16_NVFP4 routed experts AND dense MLP projections - (gate_proj/up_proj/down_proj) are promoted to NVFP4 so the CuteDSL/TRTLLM - GEMM path can consume the checkpoint's packed FP4 weights and static input - scales. Dense MLP keys are additionally re-pathed to the doubled - ``.mlp.mlp.`` form to match the ``_DenseMlpAdapter`` runtime module tree. - Other W4A16_NVFP4 modules retain their original algorithm. + Dense MLP projections (gate_proj/up_proj/down_proj directly under ``.mlp``) + are re-pathed to the doubled ``.mlp.mlp.`` form to match the + ``_DenseMlpAdapter`` runtime module tree, whatever their algorithm. On + SM100/SM103, W4A16_NVFP4 routed experts and dense MLP projections are + additionally promoted to NVFP4 so the CuteDSL/TRTLLM GEMM path can consume + the checkpoint's packed FP4 weights and static input scales. Other + W4A16_NVFP4 modules retain their original algorithm. Mutates ``quant_config_dict`` in place (model_config is frozen). @@ -531,7 +532,15 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): FP8 entry is synthesized under the fused module name so the Linear is built FP8; the weight mapper then requantizes the split weights onto one shared scale (_requantize_linear_attn_fp8_qkvz). Incomplete or non-FP8 - sets get no fused entry, and the mapper dequantizes them to bf16 instead. + sets get no fused entry, and the mapper dequantizes them to bf16 instead + (_dequantize_linear_attn_fp8_per_tensor). That includes rowwise + FP8_PER_CHANNEL_PER_TOKEN in_proj (the compressed-tensors Qwen3.8-27B + recipe): requantizing split projections onto one shared scale is only + lossless for a per-tensor scale, so the rowwise case takes the bf16 dequant, + which is exact. Their per-projection entries are left keyed on the split + checkpoint names, which no runtime module carries, so the fused Linear stays + unquantized -- correct, and a fused rowwise-FP8 path would be a performance + change, not a correctness fix. The ``lm_head`` entry is promoted W4A16_NVFP4 -> NVFP4 when ``keep_lm_head_quant`` (see _lm_head_nvfp4_enabled) and dropped otherwise: @@ -597,12 +606,22 @@ def _normalize_qwen35_quant_config_dict(model_config, keep_lm_head_quant=False): # Translate the per-layer key to that path so # ``apply_layerwise_quant_config`` matches it; otherwise the dense # MLP silently falls back to the global MIXED_PRECISION config and - # its quantized checkpoint weights fail to load. On SM100/SM103 also - # promote W4A16_NVFP4 -> NVFP4 so the CuteDSL/TRTLLM GEMM path can - # consume the checkpoint's packed FP4 weights and static input scales. + # its quantized checkpoint weights fail to load. The re-path is + # independent of the algorithm -- the weight mapper's + # ``_remap_dense_mlp_weights`` moves *every* dense MLP tensor to the + # doubled path, so any per-layer entry that keeps the checkpoint + # path is dead. Mixed compressed-tensors checkpoints exercise both + # sides of that: Qwen3.8-27B-NVFP4 has NVFP4 dense MLP in blocks + # 0-55 and FP8 dense MLP in blocks 56-63. dense_mlp_match = re.search(r"\.mlp\.(gate_proj|up_proj|down_proj)$", name) - if dense_mlp_match and cfg.quant_algo == QuantAlgo.W4A16_NVFP4: - if convert_to_nvfp4: + if dense_mlp_match: + # On SM100/SM103 promote W4A16_NVFP4 -> NVFP4 so the + # CuteDSL/TRTLLM GEMM path can consume the checkpoint's packed + # FP4 weights and static input scales. Algorithms the + # checkpoint states outright (NVFP4, FP8, FP8 rowwise) are + # never rewritten: their Linear methods load the stored tensors + # directly on every SM that has the kernels. + if convert_to_nvfp4 and cfg.quant_algo == QuantAlgo.W4A16_NVFP4: cfg = cfg.model_copy(update={"quant_algo": QuantAlgo.NVFP4}) proj = dense_mlp_match.group(1) name = name[: -len(dense_mlp_match.group(0))] + f".mlp.mlp.{proj}" diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 663835918d05..025ef40e0210 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_dense_mlp_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_dense_mlp_quant.py b/tests/unittest/_torch/models/test_qwen3_5_dense_mlp_quant.py new file mode 100644 index 000000000000..46d24846e28a --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_5_dense_mlp_quant.py @@ -0,0 +1,221 @@ +# 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-layer quant-config key translation for dense Qwen3.5/3.8 MLP blocks. + +``_DenseMlpAdapter`` wraps ``GatedMLP`` as ``self.mlp``, so a dense MLP +projection lives at ``model.layers.N.mlp.mlp.*`` at runtime while the checkpoint +stores it at ``model.layers.N.mlp.*``. The weight mapper moves every dense MLP +*tensor* onto the doubled path unconditionally, so any per-layer *quant* entry +left on the checkpoint path is dead: the module is built from the global +(unquantized) MIXED_PRECISION config and its quantized weights fail to load, or +load with their scales silently dropped. + +CPU only: no weights are loaded and no module is constructed. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm._torch.models.modeling_qwen3_5 import _normalize_qwen35_quant_config_dict +from tensorrt_llm._torch.modules.linear import ( + FP8RowwiseLinearMethod, + NVFP4LinearMethod, + W4A16NVFP4LinearMethod, + get_quant_method, +) +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +NUM_HIDDEN_LAYERS = 64 + +BLACKWELL_SMS = [100, 103] +# SM121 (DGX Spark) is where the Qwen3.8-27B-NVFP4 recipe is served; SM120/89/90 +# stand in for the other non-promoting architectures. +OTHER_SMS = [89, 90, 120, 121] +ALL_SMS = BLACKWELL_SMS + OTHER_SMS + +# What the compressed-tensors Qwen3.8-27B recipe resolves to per dense MLP +# block: NVFP4 (W4A4) for the early blocks, rowwise FP8 for the tail. +CHECKPOINT_ALGOS = [QuantAlgo.NVFP4, QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, QuantAlgo.FP8] + +DENSE_PROJECTIONS = ["gate_proj", "up_proj", "down_proj"] + + +def _model_config(quant_config_dict) -> SimpleNamespace: + return SimpleNamespace( + quant_config=QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, + kv_cache_quant_algo=QuantAlgo.FP8, + exclude_modules=[], + ), + quant_config_dict=quant_config_dict, + pretrained_config=SimpleNamespace(num_hidden_layers=NUM_HIDDEN_LAYERS), + mapping=SimpleNamespace(tp_size=1, enable_attention_dp=False), + ) + + +def _normalize(quant_config_dict, sm_version: int) -> dict: + model_config = _model_config(quant_config_dict) + with patch( + "tensorrt_llm._torch.models.modeling_qwen3_5.get_sm_version", return_value=sm_version + ): + _normalize_qwen35_quant_config_dict(model_config) + return model_config.quant_config_dict + + +@pytest.mark.parametrize("sm_version", ALL_SMS) +@pytest.mark.parametrize("algo", CHECKPOINT_ALGOS) +@pytest.mark.parametrize("proj", DENSE_PROJECTIONS) +def test_dense_mlp_repathed_for_every_algorithm(algo, proj, sm_version) -> None: + """The re-path is a module-tree fact, not an algorithm-specific one.""" + key = f"model.layers.7.mlp.{proj}" + normalized = _normalize({key: QuantConfig(quant_algo=algo)}, sm_version) + + assert key not in normalized + assert f"model.layers.7.mlp.mlp.{proj}" in normalized + # An algorithm the checkpoint states outright is never rewritten. + assert normalized[f"model.layers.7.mlp.mlp.{proj}"].quant_algo == algo + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS) +@pytest.mark.parametrize("proj", DENSE_PROJECTIONS) +def test_dense_mlp_w4a16_nvfp4_promoted_on_blackwell(proj, sm_version) -> None: + normalized = _normalize( + {f"model.layers.7.mlp.{proj}": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)}, + sm_version, + ) + + assert normalized[f"model.layers.7.mlp.mlp.{proj}"].quant_algo == QuantAlgo.NVFP4 + + +@pytest.mark.parametrize("sm_version", OTHER_SMS) +@pytest.mark.parametrize("proj", DENSE_PROJECTIONS) +def test_dense_mlp_w4a16_nvfp4_repathed_without_promotion(proj, sm_version) -> None: + """Off SM100/103 the entry is still re-pathed -- only the promotion is gated.""" + normalized = _normalize( + {f"model.layers.7.mlp.{proj}": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)}, + sm_version, + ) + + assert normalized[f"model.layers.7.mlp.mlp.{proj}"].quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_dense_mlp_repathed_from_vlm_namespace() -> None: + """Qwen3.8-27B-NVFP4 is a VLM checkpoint: keys arrive language_model-prefixed.""" + normalized = _normalize( + { + f"model.language_model.layers.0.mlp.{proj}": QuantConfig(quant_algo=QuantAlgo.NVFP4) + for proj in DENSE_PROJECTIONS + }, + 121, + ) + + assert set(normalized) == {f"model.layers.0.mlp.mlp.{proj}" for proj in DENSE_PROJECTIONS} + + +def test_dense_mlp_repathed_from_mtp_namespace() -> None: + normalized = _normalize( + {"mtp.layers.0.mlp.down_proj": QuantConfig(quant_algo=QuantAlgo.NVFP4)}, 121 + ) + + assert set(normalized) == {f"model.layers.{NUM_HIDDEN_LAYERS}.mlp.mlp.down_proj"} + + +@pytest.mark.parametrize("sm_version", ALL_SMS) +def test_non_dense_mlp_entries_are_untouched(sm_version) -> None: + """Attention, linear-attention and MoE-expert keys keep their paths.""" + entries = { + "model.layers.3.self_attn.q_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + "model.layers.3.self_attn.o_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + "model.layers.0.linear_attn.out_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + "model.layers.0.mlp.shared_expert.gate_proj": QuantAlgo.NVFP4, + } + normalized = _normalize({k: QuantConfig(quant_algo=v) for k, v in entries.items()}, sm_version) + + assert set(normalized) == set(entries) + for key, algo in entries.items(): + assert normalized[key].quant_algo == algo + + +@pytest.mark.parametrize("sm_version", BLACKWELL_SMS) +def test_moe_experts_promotion_is_unchanged(sm_version) -> None: + """Regression guard for the ModelOpt Qwen3.5/3.6 MoE path.""" + normalized = _normalize( + {"model.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)}, sm_version + ) + + assert normalized["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.NVFP4 + + +@pytest.mark.parametrize("sm_version", OTHER_SMS) +def test_moe_experts_not_promoted_off_blackwell(sm_version) -> None: + normalized = _normalize( + {"model.layers.0.mlp.experts": QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4)}, sm_version + ) + + assert normalized["model.layers.0.mlp.experts"].quant_algo == QuantAlgo.W4A16_NVFP4 + + +@pytest.mark.parametrize( + "algo, expected_method", + [ + # nvfp4-pack-quantized with FP4 input_activations parses to NVFP4 + # (W4A4). SM120/121 has CUTLASS FP4 GEMM tiles, so this is the intended + # path there -- Marlin (W4A16) is only substituted for W4A16_NVFP4. + (QuantAlgo.NVFP4, NVFP4LinearMethod), + (QuantAlgo.W4A16_NVFP4, W4A16NVFP4LinearMethod), + # float-quantized channel/token: e4m3 weight + per-channel [out, 1] + # weight_scale, flattened onto the 1-D buffer by load_weights_vanilla. + (QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, FP8RowwiseLinearMethod), + ], +) +def test_preserved_algorithm_resolves_to_the_loading_method(algo, expected_method) -> None: + """The algorithms kept above must select a method that reads the stored tensors.""" + normalized = _normalize({"model.layers.7.mlp.down_proj": QuantConfig(quant_algo=algo)}, 121) + cfg = normalized["model.layers.7.mlp.mlp.down_proj"] + + assert type(get_quant_method(cfg)) is expected_method + + +def test_split_linear_attn_fp8_fusion_is_unchanged() -> None: + """Per-tensor FP8 in_proj still fuses; rowwise FP8 still does not.""" + per_tensor = _normalize( + { + "model.layers.0.linear_attn.in_proj_qkv": QuantConfig(quant_algo=QuantAlgo.FP8), + "model.layers.0.linear_attn.in_proj_z": QuantConfig(quant_algo=QuantAlgo.FP8), + }, + 121, + ) + assert set(per_tensor) == {"model.layers.0.linear_attn.in_proj_qkvz"} + + rowwise = _normalize( + { + "model.layers.0.linear_attn.in_proj_qkv": QuantConfig( + quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN + ), + "model.layers.0.linear_attn.in_proj_z": QuantConfig( + quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN + ), + }, + 121, + ) + # No fused entry: the fused Linear stays unquantized and the weight mapper + # dequantizes the split projections to bf16 (exactly, via the [out, 1] + # per-channel scale). The split keys match no runtime module. + assert "model.layers.0.linear_attn.in_proj_qkvz" not in rowwise