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, 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)