From 1f87d8f512e09a6fe230da1e6d4ab228f6bbb91c Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 09:04:35 +0800 Subject: [PATCH] [diffusion] fix: stop reserving nccl device buffers for single-rank groups (#35538) --- .../runtime/cache/cache_dit_integration.py | 4 ++- .../runtime/distributed/group_coordinator.py | 24 ++++++++----- .../runtime/distributed/parallel_groups.py | 12 +++---- .../runtime/distributed/parallel_state.py | 7 ++-- .../unit/test_single_rank_device_group.py | 35 +++++++++++++++++++ 5 files changed, 62 insertions(+), 20 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_single_rank_device_group.py diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py index 6b6bf2db4..afbbf995e 100644 --- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py +++ b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py @@ -77,7 +77,9 @@ def _patch_cache_dit_similarity(): tp_sp_group = getattr(self, "_sglang_tp_sp_group", None) target_group = tp_sp_group or sp_group or tp_group - if target_group is None: + # Averaging over a one-rank group returns the input, so skip the + # collective rather than pay for a round trip that cannot change it. + if target_group is None or dist.get_world_size(target_group) == 1: return _original_similarity( self, t1, diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index bc564f82e..62b1ff27b 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -123,6 +123,18 @@ class GraphCaptureContext: stream: torch.cuda.Stream | None +def new_device_group(ranks, backend=None): + """Create a process group for device collectives. + + A single-rank group never runs one: every collective short-circuits on + world_size == 1. NCCL would still allocate its per-channel device buffers + for it, which costs ~390 MiB a group. + """ + return torch.distributed.new_group( + ranks, backend="gloo" if len(ranks) == 1 else backend + ) + + class GroupCoordinator: """ PyTorch ProcessGroup wrapper for a group of processes. @@ -169,9 +181,7 @@ class GroupCoordinator: self.cpu_group = None for ranks in group_ranks: - device_group = torch.distributed.new_group( - ranks, backend=torch_distributed_backend - ) + device_group = new_device_group(ranks, torch_distributed_backend) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): @@ -863,9 +873,7 @@ class PipelineGroupCoordinator(GroupCoordinator): self.device_groups = [] if len(group_ranks[0]) > 2 or len(group_ranks[0]) == 1: for ranks in group_ranks: - device_group = torch.distributed.new_group( - ranks, backend=torch_distributed_backend - ) + device_group = new_device_group(ranks, torch_distributed_backend) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): @@ -927,9 +935,7 @@ class PipelineGroupCoordinator(GroupCoordinator): ] = None self.skip_device_group = None for ranks in group_ranks: - skip_device_group = torch.distributed.new_group( - ranks, backend=torch_distributed_backend - ) + skip_device_group = new_device_group(ranks, torch_distributed_backend) if self.rank in ranks: self.skip_device_group = skip_device_group assert self.skip_device_group is not None diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py index 567c9b30a..d1d3bc2f6 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py @@ -1,7 +1,7 @@ # Reference: https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py -import torch +from .group_coordinator import new_device_group class Singleton: @@ -56,34 +56,34 @@ def set_seq_parallel_pg_by_sp_groups( def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]: return [ranks[i] for i in indices] - # Important: call torch.distributed.new_group in the same order on all ranks. + # Important: create the groups in the same order on all ranks. for sp_ranks in sp_groups: if use_ulysses_low: for i in range(num_ulysses_pgs): idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree)) ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) - group = torch.distributed.new_group(ulysses_ranks) + group = new_device_group(ulysses_ranks) if rank in ulysses_ranks: ulyssess_pg = group for i in range(num_ring_pgs): idx = list(range(i, sp_degree, num_ring_pgs)) ring_ranks = _map_indices_to_ranks(sp_ranks, idx) - group = torch.distributed.new_group(ring_ranks) + group = new_device_group(ring_ranks) if rank in ring_ranks: ring_pg = group else: for i in range(num_ring_pgs): idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree)) ring_ranks = _map_indices_to_ranks(sp_ranks, idx) - group = torch.distributed.new_group(ring_ranks) + group = new_device_group(ring_ranks) if rank in ring_ranks: ring_pg = group for i in range(num_ulysses_pgs): idx = list(range(i, sp_degree, num_ulysses_pgs)) ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) - group = torch.distributed.new_group(ulysses_ranks) + group = new_device_group(ulysses_ranks) if rank in ulysses_ranks: ulyssess_pg = group diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 70e0a4cb2..cbce8735e 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -56,6 +56,7 @@ from .group_coordinator import ( PipelineGroupCoordinator, SequenceParallelGroupCoordinator, get_local_torch_device, + new_device_group, ) logger = init_logger(__name__) @@ -999,9 +1000,7 @@ def init_dit_group( ) -> None: global _DIT assert _DIT is None, "DIT group is already initialized" - _DIT = torch.distributed.new_group( - ranks=list(range(dit_parallel_size)), backend=backend - ) + _DIT = new_device_group(list(range(dit_parallel_size)), backend) def get_dit_group() -> ProcessGroup: @@ -1018,7 +1017,7 @@ def init_vae_group( global _VAE assert _VAE is None, "VAE parallel group is already initialized" vae_ranks = list(range(dit_parallel_size, dit_parallel_size + vae_parallel_size)) - _VAE = torch.distributed.new_group(ranks=vae_ranks, backend=backend) + _VAE = new_device_group(vae_ranks, backend) def destroy_model_parallel() -> None: diff --git a/python/sglang/multimodal_gen/test/unit/test_single_rank_device_group.py b/python/sglang/multimodal_gen/test/unit/test_single_rank_device_group.py new file mode 100644 index 000000000..f10da5283 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_single_rank_device_group.py @@ -0,0 +1,35 @@ +"""Single-rank groups get gloo, so NCCL does not reserve device buffers for them.""" + +import unittest +from unittest.mock import patch + +from sglang.multimodal_gen.runtime.distributed.group_coordinator import ( + new_device_group, +) + +NEW_GROUP_PATH = "torch.distributed.new_group" + + +class TestSingleRankDeviceGroup(unittest.TestCase): + def test_single_rank_group_avoids_the_device_backend(self): + for ranks, requested in [([0], "nccl"), ([3], "hccl"), ([0], None)]: + with self.subTest(ranks=ranks, requested=requested): + with patch(NEW_GROUP_PATH) as new_group: + new_device_group(ranks, requested) + new_group.assert_called_once_with(ranks, backend="gloo") + + def test_multi_rank_group_keeps_the_requested_backend(self): + for ranks, requested in [([0, 1], "nccl"), ([0, 1, 2, 3], None)]: + with self.subTest(ranks=ranks, requested=requested): + with patch(NEW_GROUP_PATH) as new_group: + new_device_group(ranks, requested) + new_group.assert_called_once_with(ranks, backend=requested) + + def test_backend_defaults_to_none_for_multi_rank(self): + with patch(NEW_GROUP_PATH) as new_group: + new_device_group([0, 1]) + new_group.assert_called_once_with([0, 1], backend=None) + + +if __name__ == "__main__": + unittest.main()