From 6d95602ea3b1923baaa3a99545fed191d95e14ae Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 10 Apr 2026 13:23:23 -0700 Subject: [PATCH] Reduce GPU memory for MoE parallel groups (#22515) Co-authored-by: Claude Opus 4.6 --- python/sglang/srt/distributed/parallel_state.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 447715230..92e52b10e 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -609,7 +609,7 @@ class GroupCoordinator: and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_) ): outplace_all_reduce_method = "torch_symm_mem" - elif is_in_piecewise_cuda_graph(): + elif is_in_piecewise_cuda_graph() and self.pynccl_comm is not None: # For piecewise cuda graph, we use pynccl outplace allreduce outplace_all_reduce_method = "pynccl" if outplace_all_reduce_method is not None: @@ -1923,7 +1923,6 @@ def initialize_model_parallel( if moe_ep_size == tensor_model_parallel_size: _MOE_EP = _TP else: - # TODO(ch-wan): use split_group to save memory group_ranks = [] for tp_group_idx in range(num_tensor_model_parallel_groups): for moe_dp_idx in range(moe_dp_size): @@ -1940,6 +1939,8 @@ def initialize_model_parallel( group_ranks, get_world_group().local_rank, backend, + use_pynccl=False, + use_custom_allreduce=False, group_name="moe_ep", ) @@ -1948,7 +1949,6 @@ def initialize_model_parallel( if moe_tp_size == tensor_model_parallel_size: _MOE_TP = _TP else: - # TODO(ch-wan): use split_group to save memory group_ranks = [] for tp_group_idx in range(num_tensor_model_parallel_groups): for ep_dp_combined_idx in range(moe_ep_size * moe_dp_size): @@ -1966,6 +1966,8 @@ def initialize_model_parallel( group_ranks, get_world_group().local_rank, backend, + use_pynccl=False, + use_custom_allreduce=False, group_name="moe_tp", )