From 75b969ec660819b6f57b9ff7002f4af9c3af0ff8 Mon Sep 17 00:00:00 2001 From: bhsueh <11360707+byshiue@users.noreply.github.com> Date: Thu, 22 May 2025 11:45:19 +0000 Subject: [PATCH 1/2] fix bug of qwen3 fp4 workflow with EP Signed-off-by: bhsueh <11360707+byshiue@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_qwen3_moe.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index b0c24936688d..add40c8c8f4f 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -68,12 +68,17 @@ def forward( hidden_states = hidden_states.view(-1, self.hidden_dim) all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if self.enable_attention_dp and len(all_rank_num_tokens) > 1: + max_num_token = max(all_rank_num_tokens) + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, router_logits, all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) + use_dp_padding=True) if not self.enable_attention_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce( From c660c44d4ae0b16b98c1fb1464f0de48ceb81f94 Mon Sep 17 00:00:00 2001 From: bhsueh <11360707+byshiue@users.noreply.github.com> Date: Fri, 23 May 2025 01:55:33 +0000 Subject: [PATCH 2/2] fix bug of qwen3_moe with ep Signed-off-by: bhsueh <11360707+byshiue@users.noreply.github.com> --- .../_torch/models/modeling_qwen3_moe.py | 57 ++++++++++++++++--- 1 file changed, 48 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index add40c8c8f4f..3aaa1495283f 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -6,14 +6,18 @@ from tqdm import tqdm from transformers import Qwen3MoeConfig +from tensorrt_llm._mnnvl_utils import MnnvlMemory + from ..attention_backend import AttentionMetadata -from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, + allgather) from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import FusedMoE, RenormalizeMoeRoutingMethod from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm +from ..utils import disable_fp4_allgather from .modeling_qwen3 import Qwen3Attention from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, EagerFusionConfig, duplicate_kv_weight, @@ -32,12 +36,15 @@ def __init__( self.hidden_dim = config.hidden_size self.ffn_dim = config.intermediate_size self.moe_intermediate_size = config.moe_intermediate_size - # self.shared_expert_intermediate_size = config.shared_expert_intermediate_size # not used in qwen3 self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.enable_attention_dp = model_config.mapping.enable_attention_dp self.mapping = model_config.mapping self.allreduce = AllReduce(self.mapping) + self.enable_alltoall = Qwen3MoE.should_enable_alltoall( + model_config, self.top_k) + if self.enable_alltoall: + MnnvlMemory.initialize() # moe gate (linear layer) only runs in half/full precision for now self.gate = Linear(self.hidden_dim, @@ -57,6 +64,25 @@ def __init__( model_config=model_config, ) + @staticmethod + def should_enable_alltoall(model_config: ModelConfig, top_k: int) -> bool: + if not model_config.mapping.enable_attention_dp: + return False + + if model_config.mapping.tp_size == 1: + return False + + if not MnnvlMemory.supports_mnnvl(): + return False + + if os.environ.get("TRTLLM_MOE_DISABLE_ALLTOALLV", "0") == "1": + return False + + if model_config.mapping.moe_ep_size <= top_k: + return False + + return True + def forward( self, hidden_states: torch.Tensor, @@ -66,19 +92,32 @@ def forward( assert hidden_states.shape[-1] == self.hidden_dim orig_shape = hidden_states.shape hidden_states = hidden_states.view(-1, self.hidden_dim) - + use_dp_padding = False all_rank_num_tokens = attn_metadata.all_rank_num_tokens - if self.enable_attention_dp and len(all_rank_num_tokens) > 1: - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) + + if self.enable_attention_dp and self.mapping.tp_size > 1: + # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization + # to reduce allreduce BW + if disable_fp4_allgather() and not self.enable_alltoall: + hidden_states = allgather(hidden_states, + self.mapping, + dim=0, + sizes=all_rank_num_tokens) + elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq + and self.experts.has_nvfp4): + # Use padding when not using the cutlass path or when x_sf in self.experts is not None + use_dp_padding = True + max_num_token = max(all_rank_num_tokens) + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max_num_token - hidden_states.shape[0])) + router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, router_logits, all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=True) + use_dp_padding=use_dp_padding) if not self.enable_attention_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce(