diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 2efc1e775..7a482f999 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -943,6 +943,23 @@ class CommunicateWithAllReduceAndLayerNormFn: residual_input_mode=residual_input_mode, ) + if ( + (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) + and ( + residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL] + ) + and (hidden_states_output_mode == ScatterMode.TP_ATTN_FULL) + and (residual_output_mode == ScatterMode.TP_ATTN_FULL) + and context.attn_tp_size > 1 + ): + # Used when the dense MLP is tensor-parallelized along the + # attention TP group (``moe_dense_tp_size > 1``): hidden states + # need an all-reduce inside the attention TP group before the + # next layernorm, while staying in TP_ATTN_FULL on both sides. + return ( + CommunicateWithAllReduceAndLayerNormFn._tp_attn_all_reduce_and_layernorm + ) + raise NotImplementedError( f"{hidden_states_input_mode=} {residual_input_mode=} {hidden_states_output_mode=} {residual_output_mode=}" ) @@ -960,6 +977,25 @@ class CommunicateWithAllReduceAndLayerNormFn: hidden_states, residual = layernorm(hidden_states, residual) return hidden_states, residual + @staticmethod + def _tp_attn_all_reduce_and_layernorm( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + ): + """All-reduce hidden states inside the attention TP group, then layernorm. + + Used when the dense MLP shares the attention TP group + (``moe_dense_tp_size > 1``): both hidden states and residual stay in + ``TP_ATTN_FULL`` across the boundary. + """ + hidden_states = get_attention_tp_group().all_reduce(hidden_states) + if hidden_states.shape[0] != 0: + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + @staticmethod def _gather_hidden_states_and_residual( hidden_states: torch.Tensor, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 247d113d8..76d9b440f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7199,10 +7199,11 @@ class ServerArgs: assert self.base_gpu_id >= 0, "base_gpu_id must be non-negative" assert self.gpu_id_step >= 1, "gpu_id_step must be positive" - assert self.moe_dense_tp_size in { - 1, + assert self.moe_dense_tp_size in ( None, - }, "moe_dense_tp_size only support 1 and None currently" + 1, + self.tp_size, + ), "moe_dense_tp_size only supports None, 1, or tp_size currently" # Check served model name to not have colon as it is reserved for LoRA adapter syntax if not is_runai_obj_uri(self.served_model_name): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index feb505d5d..f677b3ea5 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3119,8 +3119,7 @@ def require_attn_tp_gather(server_args: ServerArgs): from sglang.srt.layers.moe.utils import get_moe_a2a_backend - assert server_args.moe_dense_tp_size in [1, None] - if not get_moe_a2a_backend().is_none() or server_args.moe_dense_tp_size == 1: + if not get_moe_a2a_backend().is_none() or server_args.moe_dense_tp_size is not None: if server_args.enable_dp_attention: return server_args.dp_size < server_args.tp_size else: