From e18606d1a36cf5d83567a9f8aa7acd281e30f523 Mon Sep 17 00:00:00 2001 From: ylichen Date: Sun, 17 May 2026 18:05:59 -0700 Subject: [PATCH 1/5] Fix the runtime error of w4a8 model with ConfiguragleMoE Signed-off-by: ylichen --- tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index b2681db5a2a0..45323bb1a30e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -645,6 +645,12 @@ def create_weights(self): assert hasattr(self.backend, "create_weights"), ( f"Backend {self.backend.__class__.__name__} must implement create_weights()" ) + # Sync quant_config to backend before creating weights. + # apply_layerwise_quant_config() matches modules by name (e.g. "model.layers.X.mlp.experts"), + # but the backend's module name has a ".backend" suffix that prevents matching. + # Propagate the (possibly updated) quant_config from ConfigurableMoE to backend here + # so that backend._get_quant_method() uses the correct per-layer quantization config. + self.backend.quant_config = self.quant_config return self.backend.create_weights() def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): From cf6765f83e68b78c633925962989f28aec27ba7c Mon Sep 17 00:00:00 2001 From: ylichen Date: Mon, 25 May 2026 01:54:43 -0700 Subject: [PATCH 2/5] Add unit test for w4a8 model Signed-off-by: ylichen --- .../defs/accuracy/test_llm_api_pytorch.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index df42b0615e65..d3785c55cd0f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4763,6 +4763,29 @@ 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, pp_size, ep_size, attention_dp, cuda_graph, + overlap_scheduler): + 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), From bbb4e6d790b2ef7722e65bc7fcc446e9839af796 Mon Sep 17 00:00:00 2001 From: ylichen Date: Sun, 28 Jun 2026 06:54:34 -0700 Subject: [PATCH 3/5] Add QA test for Qwen3-30B-A3B w4a8 model - Add test_w4a8[tp2_ep1] to llm_function_core.txt QA test list - Add test to l0_dgx_h100.yml 2-GPU pre_merge section - Add W4A8_AWQ accuracy reference (79.53) to mmlu.yaml Signed-off-by: ylichen --- tests/integration/defs/accuracy/references/mmlu.yaml | 2 ++ tests/integration/test_lists/qa/llm_function_core.txt | 1 + tests/integration/test_lists/test-db/l0_dgx_h100.yml | 2 ++ 3 files changed, 5 insertions(+) diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index a870fbba6402..96a1f143ca7b 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -247,6 +247,8 @@ Qwen3/Qwen3-30B-A3B: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 accuracy: 79.53 + - quant_algo: W4A8_AWQ + accuracy: 79.53 - quant_algo: W4A8_MXFP4_FP8 accuracy: 79.78 - quant_algo: W4A8_MXFP4_MXFP8 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: From 04c958612c16816bb01456316cb5da154fb08ef4 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Wed, 29 Jul 2026 08:26:23 +0800 Subject: [PATCH 4/5] Fix mixed W4A8 MMLU accuracy reference Register the Qwen3 W4A8 checkpoint using its global mixed-precision quantization metadata so accuracy reference lookup matches the runtime configuration. Changes - replace the W4A8_AWQ-only MMLU key with MIXED_PRECISION - include the checkpoint's FP8 KV-cache quantization key Validation - git diff --check - parse mmlu.yaml and assert the exact reference entry Result - the reported accuracy lookup resolves MIXED_PRECISION plus FP8 to the 79.53 reference Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> --- tests/integration/defs/accuracy/references/mmlu.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index 96a1f143ca7b..375e8a0ef3bd 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -247,7 +247,8 @@ Qwen3/Qwen3-30B-A3B: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 accuracy: 79.53 - - quant_algo: W4A8_AWQ + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 accuracy: 79.53 - quant_algo: W4A8_MXFP4_FP8 accuracy: 79.78 From 342b7a02011aef74217680de812588318d0ce4af Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Mon, 3 Aug 2026 18:14:17 +0800 Subject: [PATCH 5/5] [None][fix] preserve final quant config in ConfigurableMoE Defer ConfigurableMoE weight creation until layerwise quantization has selected the final configuration while preserving eager allocation for unrelated module exclusions. Preserve explicit backend overrides as authoritative. Changes - defer backend weight allocation only for post-init layerwise quantization - delegate weight-state resets so matching exclusions recreate backend weights - preserve eager allocation for unrelated FP8-block-scale exclusions - add focused lifecycle tests and annotate the W4A8 acceptance test Validation - pre-commit run --files tensorrt_llm/_torch/models/modeling_utils.py tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py - python -m pytest -q tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py - LLM_MODELS_ROOT=/home/scratch.trt_llm_data/llm-models python -m pytest -q -s tests/integration/defs/accuracy/test_llm_api_pytorch.py -k 'TestQwen3_30B_A3B and test_w4a8' Result - focused layerwise and exclusion lifecycle suite passes 2 tests - prior 2xH100 W4A8 acceptance produced MMLU accuracy 80.263 above the 77.713 threshold Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_utils.py | 12 +- .../modules/fused_moe/configurable_moe.py | 33 +++-- .../defs/accuracy/test_llm_api_pytorch.py | 11 +- .../fused_moe/test_configurable_moe.py | 122 ++++++++++++++++++ 4 files changed, 159 insertions(+), 19 deletions(-) create mode 100644 tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py 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 45323bb1a30e..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,12 +652,14 @@ def create_weights(self): assert hasattr(self.backend, "create_weights"), ( f"Backend {self.backend.__class__.__name__} must implement create_weights()" ) - # Sync quant_config to backend before creating weights. - # apply_layerwise_quant_config() matches modules by name (e.g. "model.layers.X.mlp.experts"), - # but the backend's module name has a ".backend" suffix that prevents matching. - # Propagate the (possibly updated) quant_config from ConfigurableMoE to backend here - # so that backend._get_quant_method() uses the correct per-layer quantization config. - self.backend.quant_config = self.quant_config + # 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): @@ -731,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/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index d3785c55cd0f..ab9cbb31ba38 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4769,8 +4769,15 @@ def test_nvfp4( "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, pp_size, ep_size, attention_dp, cuda_graph, - overlap_scheduler): + 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) 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()