From f45603739658c780b71c311fc1e79b770b842f67 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 25 May 2025 08:34:35 +0800 Subject: [PATCH] Utilize static dispatching for communicator (#6577) --- python/sglang/srt/layers/communicator.py | 376 ++++++++++++----------- 1 file changed, 196 insertions(+), 180 deletions(-) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 6a8610cb2..362fe3ba3 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -14,7 +14,8 @@ from dataclasses import dataclass from enum import Enum, auto -from typing import Dict, Optional, Tuple +from functools import partial +from typing import Dict, Optional import torch.distributed @@ -145,6 +146,36 @@ class LayerCommunicator: ScatterMode.FULL: self.tp_size, } + self._context = _Context( + process_group_sizes=self.process_group_sizes, + attn_tp_rank=self.attn_tp_rank, + attn_tp_size=self.attn_tp_size, + local_attn_dp_size=self.local_attn_dp_size, + tp_size=self.tp_size, + ) + self._communicate_simple_fn = _CommunicateSimpleFn.get_fn( + input_mode=self.layer_scatter_modes.layer_input_mode, + output_mode=self.layer_scatter_modes.attn_mode, + context=self._context, + ) + self._communicate_with_all_reduce_and_layer_norm_fn = ( + _CommunicateWithAllReduceAndLayerNormFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.attn_mode, + residual_input_mode=self.layer_scatter_modes.layer_input_mode, + hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, + residual_output_mode=self.layer_scatter_modes.middle_residual_mode, + context=self._context, + ) + ) + self._communicate_summable_tensor_pair_fn = ( + _CommunicateSummableTensorPairFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, + residual_input_mode=self.layer_scatter_modes.middle_residual_mode, + output_mode=self.layer_scatter_modes.layer_output_mode, + context=self._context, + ) + ) + def prepare_attn( self, hidden_states: torch.Tensor, @@ -160,12 +191,10 @@ class LayerCommunicator: else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = _communicate_simple( + hidden_states = self._communicate_simple_fn( hidden_states=hidden_states, forward_batch=forward_batch, - input_mode=self.layer_scatter_modes.layer_input_mode, - output_mode=self.layer_scatter_modes.attn_mode, - context=self._compute_context(forward_batch), + context=self._context, ) return hidden_states, residual @@ -176,16 +205,12 @@ class LayerCommunicator: residual: torch.Tensor, forward_batch: ForwardBatch, ): - return _communicate_with_all_reduce_and_layer_norm( + return self._communicate_with_all_reduce_and_layer_norm_fn( hidden_states=hidden_states, residual=residual, forward_batch=forward_batch, - hidden_states_input_mode=self.layer_scatter_modes.attn_mode, - residual_input_mode=self.layer_scatter_modes.layer_input_mode, - hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, - residual_output_mode=self.layer_scatter_modes.middle_residual_mode, layernorm=self.post_attention_layernorm, - context=self._compute_context(forward_batch), + context=self._context, ) def postprocess_layer( @@ -194,58 +219,16 @@ class LayerCommunicator: residual: torch.Tensor, forward_batch: ForwardBatch, ): - return _communicate_summable_tensor_pair( + return self._communicate_summable_tensor_pair_fn( hidden_states=hidden_states, residual=residual, forward_batch=forward_batch, - hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, - residual_input_mode=self.layer_scatter_modes.middle_residual_mode, - output_mode=self.layer_scatter_modes.layer_output_mode, - context=self._compute_context(forward_batch), + context=self._context, ) - def _compute_context(self, forward_batch: ForwardBatch): - return _Context( - num_tokens_of_mode=_compute_num_tokens_of_mode( - forward_batch, - attn_tp_rank=self.attn_tp_rank, - attn_tp_size=self.attn_tp_size, - ), - process_group_sizes=self.process_group_sizes, - attn_tp_rank=self.attn_tp_rank, - attn_tp_size=self.attn_tp_size, - local_attn_dp_size=self.local_attn_dp_size, - tp_size=self.tp_size, - ) - - -def _compute_num_tokens_of_mode( - forward_batch: ForwardBatch, attn_tp_rank: int, attn_tp_size: int -): - tp_attn_full_num_tokens = forward_batch.input_ids.shape[0] - return { - ScatterMode.SCATTERED: _torch_tensor_split_len( - tp_attn_full_num_tokens, attn_tp_size, attn_tp_rank - ), - ScatterMode.TP_ATTN_FULL: tp_attn_full_num_tokens, - ScatterMode.FULL: ( - forward_batch.gathered_buffer.shape[0] - if global_server_args_dict["enable_dp_attention"] - else forward_batch.input_ids.shape[0] - ), - } - - -def _torch_tensor_split_len(tensor_len: int, n: int, output_index: int): - if output_index < int(tensor_len % n): - return int(tensor_len / n) + 1 - else: - return int(tensor_len / n) - @dataclass class _Context: - num_tokens_of_mode: Dict["ScatterMode", int] process_group_sizes: Dict["ScatterMode", int] attn_tp_rank: int attn_tp_size: int @@ -255,75 +238,63 @@ class _Context: def is_same_group_size(self, a: "ScatterMode", b: "ScatterMode"): return self.process_group_sizes[a] == self.process_group_sizes[b] - def check_shape(self, x: torch.Tensor, mode: ScatterMode): - if x is None: - return - - actual_num_tokens = x.shape[0] - expect_num_tokens = self.num_tokens_of_mode[mode] - assert ( - actual_num_tokens == expect_num_tokens - ), f"{actual_num_tokens=} {expect_num_tokens=} {mode=} {x.shape=} {self.num_tokens_of_mode=} {self.process_group_sizes=}" - return x - - def check_shapes( - self, xs: Tuple[torch.Tensor, ...], modes: Tuple[ScatterMode, ...] - ) -> Tuple[torch.Tensor, ...]: - return tuple( - [self.check_shape(x, mode) for x, mode in zip(xs, modes, strict=True)] - ) - - -def _communicate_simple( - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - input_mode: ScatterMode, - output_mode: ScatterMode, - context: _Context, -) -> torch.Tensor: - def _inner(): - nonlocal hidden_states +class _CommunicateSimpleFn: + @staticmethod + def get_fn( + input_mode: ScatterMode, + output_mode: ScatterMode, + context: _Context, + ): if context.is_same_group_size(input_mode, output_mode): - return hidden_states + return _CommunicateSimpleFn._trivial if (input_mode == ScatterMode.SCATTERED) and ( output_mode == ScatterMode.TP_ATTN_FULL ): - hidden_states, local_hidden_states = ( - forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], - hidden_states, - ) - attn_tp_all_gather( - list(hidden_states.tensor_split(context.attn_tp_size)), - local_hidden_states, - ) - return hidden_states + return _CommunicateSimpleFn._scattered_to_tp_attn_full raise NotImplementedError(f"{input_mode=} {output_mode=}") - context.check_shape(hidden_states, input_mode) - return context.check_shape(_inner(), output_mode) + @staticmethod + def _trivial( + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + context: _Context, + ) -> torch.Tensor: + return hidden_states + + @staticmethod + def _scattered_to_tp_attn_full( + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + context: _Context, + ) -> torch.Tensor: + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + attn_tp_all_gather( + list(hidden_states.tensor_split(context.attn_tp_size)), + local_hidden_states, + ) + return hidden_states -def _communicate_with_all_reduce_and_layer_norm( - hidden_states: torch.Tensor, - residual: torch.Tensor, - hidden_states_input_mode: ScatterMode, - residual_input_mode: ScatterMode, - hidden_states_output_mode: ScatterMode, - residual_output_mode: ScatterMode, - forward_batch: ForwardBatch, - layernorm: torch.nn.Module, - context: _Context, -): +class _CommunicateWithAllReduceAndLayerNormFn: """Besides communication, needs to 1. All reduce in tp_attn_group on hidden_states 2. Apply layer norm """ - def _inner(): - nonlocal hidden_states, residual + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + hidden_states_output_mode: ScatterMode, + residual_output_mode: ScatterMode, + context: _Context, + ): if ( context.is_same_group_size( @@ -332,10 +303,7 @@ def _communicate_with_all_reduce_and_layer_norm( and context.is_same_group_size(residual_input_mode, residual_output_mode) and context.attn_tp_size == 1 ): - # TODO move these `if shape != 0` into LayerNorm itself - if hidden_states.shape[0] != 0: - hidden_states, residual = layernorm(hidden_states, residual) - return hidden_states, residual + return _CommunicateWithAllReduceAndLayerNormFn._simple if ( (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) @@ -343,21 +311,7 @@ def _communicate_with_all_reduce_and_layer_norm( and (hidden_states_output_mode == ScatterMode.FULL) and (residual_output_mode == ScatterMode.TP_ATTN_FULL) ): - if context.local_attn_dp_size != 1: - if context.attn_tp_rank == 0: - hidden_states += residual - hidden_states, local_hidden_states = ( - forward_batch.gathered_buffer, - hidden_states, - ) - dp_gather_partial(hidden_states, local_hidden_states, forward_batch) - dp_scatter(residual, hidden_states, forward_batch) - if hidden_states.shape[0] != 0: - hidden_states = layernorm(hidden_states) - else: - hidden_states = tensor_model_parallel_all_reduce(hidden_states) - hidden_states, residual = layernorm(hidden_states, residual) - return hidden_states, residual + return _CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states if ( (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) @@ -367,85 +321,147 @@ def _communicate_with_all_reduce_and_layer_norm( and (hidden_states_output_mode == ScatterMode.SCATTERED) and (residual_output_mode == ScatterMode.SCATTERED) ): - tensor_list = list(hidden_states.tensor_split(context.attn_tp_size)) - hidden_states = tensor_list[context.attn_tp_rank] - attn_tp_reduce_scatter(hidden_states, tensor_list) - if residual_input_mode == ScatterMode.TP_ATTN_FULL: - residual = residual.tensor_split(context.attn_tp_size)[ - context.attn_tp_rank - ] - if hidden_states.shape[0] != 0: - hidden_states, residual = layernorm(hidden_states, residual) - return hidden_states, residual + return partial( + _CommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual, + residual_input_mode=residual_input_mode, + ) raise NotImplementedError( f"{hidden_states_input_mode=} {residual_input_mode=} {residual_output_mode=} {residual_output_mode=}" ) - context.check_shapes( - (hidden_states, residual), (hidden_states_input_mode, residual_input_mode) - ) - return context.check_shapes( - _inner(), (hidden_states_output_mode, residual_output_mode) - ) + @staticmethod + def _simple( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: _Context, + ): + # TODO move these `if shape != 0` into LayerNorm itself + if hidden_states.shape[0] != 0: + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + + @staticmethod + def _gather_hidden_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: _Context, + ): + if context.local_attn_dp_size != 1: + if context.attn_tp_rank == 0: + hidden_states += residual + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer, + hidden_states, + ) + dp_gather_partial(hidden_states, local_hidden_states, forward_batch) + dp_scatter(residual, hidden_states, forward_batch) + if hidden_states.shape[0] != 0: + hidden_states = layernorm(hidden_states) + else: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + + @staticmethod + def _scatter_hidden_states_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: _Context, + *, + residual_input_mode, + ): + tensor_list = list(hidden_states.tensor_split(context.attn_tp_size)) + hidden_states = tensor_list[context.attn_tp_rank] + attn_tp_reduce_scatter(hidden_states, tensor_list) + if residual_input_mode == ScatterMode.TP_ATTN_FULL: + residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank] + if hidden_states.shape[0] != 0: + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual -def _communicate_summable_tensor_pair( - hidden_states: torch.Tensor, - residual: torch.Tensor, - forward_batch: ForwardBatch, - hidden_states_input_mode: ScatterMode, - residual_input_mode: ScatterMode, - output_mode: ScatterMode, - context: _Context, -): - """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed.""" +class _CommunicateSummableTensorPairFn: - def _inner(): - nonlocal hidden_states, residual + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + output_mode: ScatterMode, + context: _Context, + ): + """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed.""" if context.is_same_group_size( hidden_states_input_mode, output_mode ) and context.is_same_group_size(residual_input_mode, output_mode): - return hidden_states, residual + return _CommunicateSummableTensorPairFn._trivial if ( (hidden_states_input_mode == ScatterMode.FULL) and (residual_input_mode == ScatterMode.TP_ATTN_FULL) and (output_mode == ScatterMode.TP_ATTN_FULL) ): - # TODO(ch-wan): use reduce-scatter in MLP to avoid this scatter - # important: forward batch.gathered_buffer is used both after scatter and after gather. - # be careful about this! - hidden_states, global_hidden_states = ( - forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], - hidden_states, - ) - dp_scatter(hidden_states, global_hidden_states, forward_batch) - return hidden_states, residual + return _CommunicateSummableTensorPairFn._scatter_hidden_states if ( (hidden_states_input_mode == ScatterMode.SCATTERED) and (residual_input_mode == ScatterMode.SCATTERED) and (output_mode == ScatterMode.TP_ATTN_FULL) ): - hidden_states += residual - residual = None - hidden_states, local_hidden_states = ( - forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], - hidden_states, - ) - attn_tp_all_gather( - list(hidden_states.tensor_split(context.attn_tp_size)), - local_hidden_states, - ) - return hidden_states, residual + return _CommunicateSummableTensorPairFn._gather raise NotImplementedError( f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}" ) - context.check_shapes( - (hidden_states, residual), (hidden_states_input_mode, residual_input_mode) - ) - return context.check_shapes(_inner(), (output_mode, output_mode)) + @staticmethod + def _trivial( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: _Context, + ): + return hidden_states, residual + + @staticmethod + def _scatter_hidden_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: _Context, + ): + # TODO(ch-wan): use reduce-scatter in MLP to avoid this scatter + # important: forward batch.gathered_buffer is used both after scatter and after gather. + # be careful about this! + hidden_states, global_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + dp_scatter(hidden_states, global_hidden_states, forward_batch) + return hidden_states, residual + + @staticmethod + def _gather( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: _Context, + ): + hidden_states += residual + residual = None + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + attn_tp_all_gather( + list(hidden_states.tensor_split(context.attn_tp_size)), + local_hidden_states, + ) + return hidden_states, residual