diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index eb91a6456209..59a9e1d23067 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -622,16 +622,10 @@ def apply_quant_config_exclude_modules(self): # Reset _weights_created so create_weights() in # __post_init__ will re-create this module's weights # with the updated (non-quantized) config. Some - # wrappers (e.g. ConfigurableMoE) expose - # _weights_created as a read-only property that - # delegates to a child backend module — that backend - # is itself an nn.Module child and will be visited - # separately, so swallow the resulting AttributeError. + # Wrappers such as ConfigurableMoE delegate this state + # update to their child backend. if hasattr(module, '_weights_created'): - try: - module._weights_created = False - except AttributeError: - pass + module._weights_created = False def __post_init__(self): self.apply_layerwise_quant_config() diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index b2681db5a2a0..fcf4937ee530 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -172,6 +172,7 @@ def __init__( layer_idx=layer_idx, # ConfigurableMoE needs correct layer_idx for EPLB initialization **kwargs, ) + self._override_quant_config = override_quant_config if override_quant_config is not None: self.quant_config = override_quant_config @@ -343,8 +344,14 @@ def _create_and_sync_backend( # Sync done -- now the backend has enough info to allocate weight # tensors with the right shard / slot count. - if not backend_model_config.skip_create_weights_in_init: - self.backend.create_weights() + # Layerwise quantization is applied after the model is constructed. + # Defer allocation so the backend is built from the final wrapper + # quant_config rather than an earlier global value. Module exclusions + # reset _weights_created on matching modules during __post_init__, so + # unrelated exclusions retain the historical eager allocation path. + has_post_init_quant_config = model_config.quant_config_dict is not None + if not backend_model_config.skip_create_weights_in_init and not has_post_init_quant_config: + self.create_weights() def _supports_load_balancer(self) -> bool: """Check if this MoE implementation supports load balancer. @@ -645,6 +652,14 @@ def create_weights(self): assert hasattr(self.backend, "create_weights"), ( f"Backend {self.backend.__class__.__name__} must implement create_weights()" ) + # An explicit override is authoritative. Otherwise use the wrapper's + # final value, after model __post_init__ has applied layerwise and + # exclusion-based quantization settings. + self.backend.quant_config = ( + self._override_quant_config + if self._override_quant_config is not None + else self.quant_config + ) return self.backend.create_weights() def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): @@ -725,6 +740,14 @@ def _weights_created(self): ) return self.backend._weights_created + @_weights_created.setter + def _weights_created(self, value: bool) -> None: + """Update backend weight state during post-init quantization changes.""" + assert hasattr(self.backend, "_weights_created"), ( + f"Backend {self.backend.__class__.__name__} must have _weights_created attribute" + ) + self.backend._weights_created = value + # ========== Explicit Backend Attribute Proxies ========== # These properties delegate to backend for commonly accessed attributes # TODO: Unify the property access to backend in ConfigurableMoE. diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index a870fbba6402..375e8a0ef3bd 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -247,6 +247,9 @@ Qwen3/Qwen3-30B-A3B: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 accuracy: 79.53 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + accuracy: 79.53 - quant_algo: W4A8_MXFP4_FP8 accuracy: 79.78 - quant_algo: W4A8_MXFP4_MXFP8 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index df42b0615e65..ab9cbb31ba38 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4763,6 +4763,36 @@ def test_nvfp4( task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_hopper + @skip_post_blackwell + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", + [(2, 1, 1, False, False, True)], + ids=["tp2_ep1"]) + def test_w4a8( + self, + tp_size: int, + pp_size: int, + ep_size: int, + attention_dp: bool, + cuda_graph: bool, + overlap_scheduler: bool, + ) -> None: + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + cuda_graph_config=CudaGraphConfig() if cuda_graph else None) + + llm = LLM( + f"{llm_models_root()}/Qwen3/saved_models_Qwen3-30B-A3B_w4a8_hf", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + **pytorch_config, + enable_attention_dp=attention_dp) + with llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.parametrize("moe_backend", ["CUTLASS", "TRTLLM"]) @pytest.mark.parametrize("tp_size,pp_size,ep_size", [ (1, 1, 1), diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 470177b82816..a1b34e1fa11a 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -729,6 +729,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_trtllm-torch_compile=True] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRITON] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRTLLM] +accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8[tp2_ep1] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-CUTLASS] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-TRTLLM] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[mxfp8-latency-CUTLASS] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 3cda7d22ca40..76c207222cc2 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -65,6 +65,8 @@ l0_dgx_h100: # llmapi - unittest/llmapi/test_mpi_session.py::test_llmapi_launch_multiple_tasks - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_gen_only_spec_dec + # ------------- Model specific tests --------------- + - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8[tp2_ep1] - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py b/tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py new file mode 100644 index 000000000000..7c451fad4d11 --- /dev/null +++ b/tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py @@ -0,0 +1,122 @@ +# 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. + +from unittest.mock import Mock, patch + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.modules.fused_moe.configurable_moe import ( + _BACKEND_SYNC_ATTRS, + ConfigurableMoE, +) +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _wrapper() -> ConfigurableMoE: + wrapper = ConfigurableMoE.__new__(ConfigurableMoE) + torch.nn.Module.__init__(wrapper) + wrapper.num_experts = 8 + wrapper.hidden_size = 16 + wrapper.intermediate_size = 32 + wrapper.dtype = torch.bfloat16 + wrapper.reduce_results = False + wrapper.aux_stream_dict = None + wrapper.weight_loading_mode = None + wrapper.apply_router_weight_on_input = False + wrapper.activation_type = None + wrapper._override_quant_config = None + for attr in _BACKEND_SYNC_ATTRS: + setattr(wrapper, attr, None) + return wrapper + + +def _create_backend( + wrapper: ConfigurableMoE, + model_config: ModelConfig, + override_quant_config: QuantConfig | None = None, +) -> Mock: + backend = Mock() + with ( + patch( + "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", + return_value=Mock(), + ), + patch( + "tensorrt_llm._torch.modules.fused_moe.create_moe.create_moe_backend", + return_value=backend, + ), + ): + wrapper._create_and_sync_backend( + model_config=model_config, + routing_method=Mock(), + override_quant_config=override_quant_config, + ) + return backend + + +def test_layerwise_quant_config_is_applied_before_weight_creation() -> None: + global_config = QuantConfig() + layer_config = QuantConfig() + model_config = ModelConfig( + quant_config=global_config, + quant_config_dict={"model.layers.0.mlp.experts": layer_config}, + ) + wrapper = _wrapper() + wrapper.quant_config = global_config + + backend = _create_backend(wrapper, model_config) + + backend.create_weights.assert_not_called() + wrapper.quant_config = layer_config + wrapper.create_weights() + + assert backend.quant_config is layer_config + backend.create_weights.assert_called_once_with() + + +def test_exclusions_only_recreate_matching_moe_weights() -> None: + quant_config = QuantConfig( + quant_algo=QuantAlgo.FP8, + exclude_modules=["*kv_b_proj*", "*k_b_proj*", "*eh_proj"], + ) + model_config = ModelConfig(quant_config=quant_config) + wrapper = _wrapper() + wrapper.quant_config = quant_config + + backend = _create_backend(wrapper, model_config) + + assert backend.quant_config is quant_config + backend.create_weights.assert_called_once_with() + backend.create_weights.reset_mock() + backend._weights_created = True + + root = torch.nn.Module() + root.model_config = ModelConfig( + quant_config=QuantConfig( + quant_algo=QuantAlgo.FP8, + exclude_modules=["experts"], + ) + ) + root.experts = wrapper + + DecoderModelForCausalLM.apply_quant_config_exclude_modules(root) + + assert not backend._weights_created + wrapper.create_weights() + assert wrapper.quant_config.quant_algo is None + assert backend.quant_config.quant_algo is None + backend.create_weights.assert_called_once_with()