Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 3 additions & 9 deletions tensorrt_llm/_torch/models/modeling_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
27 changes: 25 additions & 2 deletions tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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.
Expand Down
3 changes: 3 additions & 0 deletions tests/integration/defs/accuracy/references/mmlu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 30 additions & 0 deletions tests/integration/defs/accuracy/test_llm_api_pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
1 change: 1 addition & 0 deletions tests/integration/test_lists/qa/llm_function_core.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 2 additions & 0 deletions tests/integration/test_lists/test-db/l0_dgx_h100.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
122 changes: 122 additions & 0 deletions tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py
Original file line number Diff line number Diff line change
@@ -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()
Comment on lines +71 to +88

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Add coverage for the other quantization lifecycle branches.

This test covers quant_config_dict deferral. It does not cover a non-empty exclude_modules value at Line 354. It also does not verify that _override_quant_config remains authoritative at Lines 664-668.

Add one focused case for exclusions and one for explicit override precedence.

Test coverage summary: insufficient. The added unit test covers layerwise configuration propagation. It does not cover all changed allocation branches. Integration test-list registration does not apply to this unit-test module.

As per path instructions, changed test code requires a coverage verdict.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/fused_moe/test_configurable_moe.py` around
lines 70 - 87, Add two focused tests alongside
test_layerwise_quant_config_is_applied_before_weight_creation: cover the
non-empty exclude_modules branch in the relevant backend allocation flow, and
verify that an explicit _override_quant_config remains authoritative over other
quantization configuration sources. Keep each case isolated and assert the
resulting quant_config and weight-creation behavior, including the expected
allocation/deferment outcome for exclusions.

Source: Path instructions



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()
Loading