From 464e9b0fcf8a1c766d83cbf0ea14295ba1c48862 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:26:58 -0700 Subject: [PATCH 1/2] [https://nvbugs/6891163][fix] Replicate Qwen3 embedding under attention DP Under attention DP without lm_head TP, DecoderModelForCausalLM builds a replicated lm_head, while Qwen3Model always built a TP-sharded embed_tokens. Tied-embedding checkpoints such as Qwen3-0.6B then failed the tie assertion. Build embed_tokens with the same parallelism as the lm_head, and remove the waiver for test_llm_api_attention_dp_kv_events. Co-Authored-By: Claude Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_qwen3.py | 27 +++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3.py b/tensorrt_llm/_torch/models/modeling_qwen3.py index 24cd6d8d418d..880aee242cd6 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3.py @@ -208,15 +208,24 @@ def __init__(self, config = self.model_config self.mapping_with_cp = mapping_with_cp - self.embed_tokens = Embedding( - config.pretrained_config.vocab_size, - config.pretrained_config.hidden_size, - dtype=config.pretrained_config.torch_dtype, - mapping=config.mapping, - tensor_parallel_mode=TensorParallelMode.COLUMN, - gather_output=True, - allreduce_strategy=config.allreduce_strategy, - ) + # Under attention DP each rank embeds its own tokens, so the embedding + # is replicated like the attention-DP lm_head it may be tied to. + if config.mapping.enable_attention_dp and not config.mapping.enable_lm_head_tp_in_adp: + self.embed_tokens = Embedding( + config.pretrained_config.vocab_size, + config.pretrained_config.hidden_size, + dtype=config.pretrained_config.torch_dtype, + ) + else: + self.embed_tokens = Embedding( + config.pretrained_config.vocab_size, + config.pretrained_config.hidden_size, + dtype=config.pretrained_config.torch_dtype, + mapping=config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + allreduce_strategy=config.allreduce_strategy, + ) self.layers = nn.ModuleList([ Qwen3DecoderLayer( model_config, From fc65496d9f0b4fa0c4ed4aea8470ee93485ddd91 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:38:13 -0700 Subject: [PATCH 2/2] [nvbugs/6891163][chore] Remove stale waiver after fix Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1c14893fa48f..3e8375eddc7c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -192,7 +192,6 @@ unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) unittest/llmapi/test_llm_args.py::TestKvCacheManagerV2AutoResolution::test_required_v2_refuses_incompatible_disagg[nixl-cpp] SKIP (https://nvbugs/6917693) unittest/llmapi/test_llm_args.py::TestKvCacheManagerV2AutoResolution::test_required_v2_refuses_incompatible_disagg[ucx-python] SKIP (https://nvbugs/6917693) -unittest/llmapi/test_llm_kv_cache_events.py::test_llm_api_attention_dp_kv_events SKIP (https://nvbugs/6891163) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_llm_disagg_streaming_gen_cancelled[PYTHON] SKIP (https://nvbugs/6903683)