@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 38fd2405f3
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 1, 22:30 UTC

use split_group for pytorch process group creation (#41980)

Signed-off-by: Tushar Jain <tushar00jain@users.noreply.github.com> Co-authored-by: Tushar Jain <tushar00jain@users.noreply.github.com>

Tushar Jain committed Jun 4, 2026 at 14:36 UTC 38fd2405f35e9a38347f34deec1ad6d1bce0c2fb
11 files changed +527 -50
tests/distributed/test_dcp_a2a.py
+8 -1
@@ -15,6 +15,7 @@ import pytest
15 import torch
16 import torch.distributed as dist
17
18 +import vllm.envs as envs
19 from vllm.config.parallel import ParallelConfig
20 from vllm.utils.network_utils import get_open_port
21 from vllm.utils.system_utils import update_environment_variables
@@ -379,7 +380,13 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
380 update_environment_variables(env)
381 local_rank = int(env["LOCAL_RANK"])
382 torch.accelerator.set_device_index(local_rank)
382 - dist.init_process_group(backend="nccl")
383 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
384 + dist.init_process_group(
385 + backend="cpu:gloo,cuda:nccl",
386 + device_id=torch.device(f"cuda:{local_rank}"),
387 + )
388 + else:
389 + dist.init_process_group(backend="nccl")
390 use_workspace = env.get("USE_WORKSPACE") == "1"
391 if use_workspace:
392 from vllm.v1.worker.workspace import init_workspace_manager
tests/distributed/test_pynccl.py
+23 -10
@@ -9,6 +9,7 @@ import pytest
9 import torch
10 import torch.distributed
11
12 +import vllm.envs as envs
13 from tests.utils import ensure_current_vllm_config
14 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
15 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
@@ -82,11 +83,18 @@ def test_pynccl():
83 @worker_fn_wrapper
84 def multiple_allreduce_worker_fn():
85 device = torch.device(f"cuda:{torch.distributed.get_rank()}")
85 - groups = [
86 - torch.distributed.new_group(ranks=[0, 1], backend="gloo"),
87 - torch.distributed.new_group(ranks=[2, 3], backend="gloo"),
88 - ]
89 - group = groups[0] if torch.distributed.get_rank() in [0, 1] else groups[1]
86 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
87 + # Eager-init path: parent PG has bound_device_id + a CPU backend,
88 + # so split_group is supported.
89 + group = torch.distributed.split_group(
90 + split_ranks=[[0, 1], [2, 3]], backend="cpu:gloo,cuda:nccl"
91 + )
92 + else:
93 + groups = [
94 + torch.distributed.new_group(ranks=[0, 1], backend="gloo"),
95 + torch.distributed.new_group(ranks=[2, 3], backend="gloo"),
96 + ]
97 + group = groups[0] if torch.distributed.get_rank() in [0, 1] else groups[1]
98 pynccl_comm = PyNcclCommunicator(group=group, device=device)
99 tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
100 # two groups can communicate independently
@@ -339,11 +347,16 @@ def test_pynccl_send_recv():
347 @worker_fn_wrapper
348 def multiple_send_recv_worker_fn():
349 device = torch.device(f"cuda:{torch.distributed.get_rank()}")
342 - groups = [
343 - torch.distributed.new_group(ranks=[0, 2], backend="gloo"),
344 - torch.distributed.new_group(ranks=[1, 3], backend="gloo"),
345 - ]
346 - group = groups[0] if torch.distributed.get_rank() in [0, 2] else groups[1]
350 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
351 + group = torch.distributed.split_group(
352 + split_ranks=[[0, 2], [1, 3]], backend="cpu:gloo,cuda:nccl"
353 + )
354 + else:
355 + groups = [
356 + torch.distributed.new_group(ranks=[0, 2], backend="gloo"),
357 + torch.distributed.new_group(ranks=[1, 3], backend="gloo"),
358 + ]
359 + group = groups[0] if torch.distributed.get_rank() in [0, 2] else groups[1]
360 pynccl_comm = PyNcclCommunicator(group=group, device=device)
361 if torch.distributed.get_rank() == 0:
362 tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
tests/distributed/test_quick_all_reduce.py
+22 -7
@@ -9,6 +9,7 @@ import ray
9 import torch
10 import torch.distributed as dist
11
12 +import vllm.envs as envs
13 from vllm import _custom_ops as ops
14 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
15 from vllm.distributed.device_communicators.quick_all_reduce import (
@@ -397,13 +398,27 @@ def qr_variable_input(rank, world_size):
398 ranks = []
399 for i in range(world_size):
400 ranks.append(i)
400 - dist.init_process_group(
401 - backend="nccl",
402 - init_method="tcp://127.0.0.1:29500",
403 - rank=rank,
404 - world_size=world_size,
405 - )
406 - cpu_group = torch.distributed.new_group(ranks, backend="nccl")
401 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
402 + dist.init_process_group(
403 + backend="cpu:gloo,cuda:nccl",
404 + init_method="tcp://127.0.0.1:29500",
405 + rank=rank,
406 + world_size=world_size,
407 + device_id=device,
408 + )
409 + else:
410 + dist.init_process_group(
411 + backend="nccl",
412 + init_method="tcp://127.0.0.1:29500",
413 + rank=rank,
414 + world_size=world_size,
415 + )
416 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
417 + cpu_group = torch.distributed.split_group(
418 + split_ranks=[ranks], backend="cpu:gloo,cuda:nccl"
419 + )
420 + else:
421 + cpu_group = torch.distributed.new_group(ranks, backend="nccl")
422
423 handle = ops.qr_get_handle(_ptr)
424 world_size = dist.get_world_size(group=cpu_group)
tests/distributed/test_split_group.py new
+233
@@ -0,0 +1,233 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for split_group in GroupCoordinator.
4 +
5 +These tests verify that:
6 +1. split_group is used for both device and CPU group creation.
7 +2. Multiple subgroups work correctly with split_group.
8 +3. Both GPU and CPU all-reduce work on split groups.
9 +"""
10 +
11 +import os
12 +from typing import Any
13 +
14 +import multiprocess as mp
15 +import pytest
16 +import torch
17 +import torch.distributed
18 +
19 +import vllm.envs as envs
20 +from vllm.distributed.parallel_state import (
21 + GroupCoordinator,
22 + init_distributed_environment,
23 +)
24 +from vllm.utils.system_utils import update_environment_variables
25 +
26 +# The whole module exercises the split_group code path, which is opt-in
27 +# behind VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1.
28 +pytestmark = pytest.mark.skipif(
29 + not envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP,
30 + reason=("VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1 not set; split_group path is opt-in."),
31 +)
32 +
33 +mp.set_start_method("spawn", force=True)
34 +
35 +
36 +def distributed_run(fn, world_size):
37 + number_of_processes = world_size
38 + processes: list[mp.Process] = []
39 + for i in range(number_of_processes):
40 + env: dict[str, str] = {}
41 + env["RANK"] = str(i)
42 + env["LOCAL_RANK"] = str(i)
43 + env["WORLD_SIZE"] = str(number_of_processes)
44 + env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
45 + env["MASTER_ADDR"] = "localhost"
46 + env["MASTER_PORT"] = "12346"
47 + # propagate the opt-in flag to the spawned child workers
48 + env["VLLM_DISTRIBUTED_USE_SPLIT_GROUP"] = "1"
49 + p = mp.Process(target=fn, args=(env,))
50 + processes.append(p)
51 + p.start()
52 +
53 + for p in processes:
54 + p.join()
55 +
56 + for p in processes:
57 + assert p.exitcode == 0
58 +
59 +
60 +def worker_fn_wrapper(fn):
61 + def wrapped_fn(env):
62 + update_environment_variables(env)
63 + local_rank = os.environ["LOCAL_RANK"]
64 + device = torch.device(f"cuda:{local_rank}")
65 + torch.accelerator.set_device_index(device)
66 + init_distributed_environment()
67 + fn()
68 +
69 + return wrapped_fn
70 +
71 +
72 +def _verify_device_group(coordinator: GroupCoordinator):
73 + """Verify device group works via all-reduce."""
74 + local_rank = torch.distributed.get_rank()
75 + device = torch.device(f"cuda:{local_rank}")
76 + tensor = torch.ones(16, 16, dtype=torch.float32, device=device)
77 + torch.distributed.all_reduce(tensor, group=coordinator.device_group)
78 + torch.accelerator.synchronize()
79 + expected = coordinator.world_size
80 + assert torch.all(tensor == expected).cpu().item(), (
81 + f"Device group all-reduce failed: expected {expected}, "
82 + f"got {tensor.flatten()[0].item()}"
83 + )
84 +
85 +
86 +def _verify_cpu_group(coordinator: GroupCoordinator):
87 + """Verify CPU group works via all-reduce."""
88 + tensor = torch.ones(16, dtype=torch.float32)
89 + torch.distributed.all_reduce(tensor, group=coordinator.cpu_group)
90 + expected = coordinator.world_size
91 + assert torch.all(tensor == expected).cpu().item(), (
92 + f"CPU group all-reduce failed: expected {expected}, "
93 + f"got {tensor.flatten()[0].item()}"
94 + )
95 +
96 +
97 +# ---------------------------------------------------------------------------
98 +# Test 1: Basic split_group path with 2 GPUs
99 +# ---------------------------------------------------------------------------
100 +@worker_fn_wrapper
101 +def split_group_basic_worker():
102 + rank = torch.distributed.get_rank()
103 + world_size = torch.distributed.get_world_size()
104 + group_ranks = [list(range(world_size))]
105 +
106 + coordinator = GroupCoordinator(
107 + group_ranks=group_ranks,
108 + local_rank=rank,
109 + torch_distributed_backend="nccl",
110 + use_device_communicator=False,
111 + group_name="test_split_basic",
112 + )
113 +
114 + _verify_device_group(coordinator)
115 + _verify_cpu_group(coordinator)
116 +
117 +
118 +@pytest.mark.skipif(
119 + torch.accelerator.device_count() < 2,
120 + reason="Need at least 2 GPUs to run the test.",
121 +)
122 +def test_split_group_basic():
123 + """Test basic GroupCoordinator creation with split_group."""
124 + distributed_run(split_group_basic_worker, 2)
125 +
126 +
127 +# ---------------------------------------------------------------------------
128 +# Test 2: Multiple subgroups with split_group (4 GPUs)
129 +# ---------------------------------------------------------------------------
130 +@worker_fn_wrapper
131 +def split_group_multiple_subgroups_worker():
132 + rank = torch.distributed.get_rank()
133 + group_ranks = [[0, 1], [2, 3]]
134 +
135 + coordinator = GroupCoordinator(
136 + group_ranks=group_ranks,
137 + local_rank=rank,
138 + torch_distributed_backend="nccl",
139 + use_device_communicator=False,
140 + group_name="test_split_multi",
141 + )
142 +
143 + assert coordinator.world_size == 2
144 +
145 + _verify_device_group(coordinator)
146 + _verify_cpu_group(coordinator)
147 +
148 + if rank in [0, 1]:
149 + assert coordinator.ranks == [0, 1]
150 + else:
151 + assert coordinator.ranks == [2, 3]
152 +
153 +
154 +@pytest.mark.skipif(
155 + torch.accelerator.device_count() < 4,
156 + reason="Need at least 4 GPUs to run the test.",
157 +)
158 +def test_split_group_multiple_subgroups():
159 + """Test GroupCoordinator with multiple independent subgroups."""
160 + distributed_run(split_group_multiple_subgroups_worker, 4)
161 +
162 +
163 +# ---------------------------------------------------------------------------
164 +# Test 3: split_group contract — every parent rank must enter with the same
165 +# ``split_ranks``. NCCL happens to produce
166 +# correct subgroups for disjoint partitions because the wrapper hashes
167 +# ``my_group`` to derive the comm-split color, but the contract violation is
168 +# real and would break under non-partition / non-NCCL backends. This test
169 +# captures the actual ``split_ranks`` argument passed on every rank and
170 +# asserts they match.
171 +# ---------------------------------------------------------------------------
172 +@worker_fn_wrapper
173 +def split_group_contract_worker():
174 + rank = torch.distributed.get_rank()
175 + group_ranks = [[0, 1], [2, 3]]
176 +
177 + captured: list[list[list[int]]] = []
178 + original_split_group = torch.distributed.split_group
179 +
180 + def capturing_split_group(*args, split_ranks=None, **kwargs):
181 + captured.append([list(g) for g in split_ranks])
182 + return original_split_group(*args, split_ranks=split_ranks, **kwargs)
183 +
184 + torch.distributed.split_group = capturing_split_group
185 + try:
186 + GroupCoordinator(
187 + group_ranks=group_ranks,
188 + local_rank=rank,
189 + torch_distributed_backend="nccl",
190 + use_device_communicator=False,
191 + group_name="test_split_contract",
192 + )
193 + finally:
194 + torch.distributed.split_group = original_split_group
195 +
196 + # GroupCoordinator builds two subgroups (device + cpu) per coordinator,
197 + # so every rank must have made exactly two split_group calls.
198 + if len(captured) != 2:
199 + raise AssertionError(
200 + f"rank {rank} expected 2 split_group calls (device + cpu), "
201 + f"got {len(captured)}: {captured}"
202 + )
203 +
204 + world_size = torch.distributed.get_world_size()
205 + for call_idx in range(2):
206 + gathered: list[Any] = [None] * world_size
207 + torch.distributed.all_gather_object(gathered, captured[call_idx])
208 + # Normalize for stable comparison (sort each subgroup and the outer list).
209 + norm = [
210 + sorted([sorted(sg) for sg in per_rank_args]) for per_rank_args in gathered
211 + ]
212 + reference = norm[0]
213 + for r, args in enumerate(norm):
214 + if args != reference:
215 + raise AssertionError(
216 + f"split_group contract violation on call #{call_idx}: "
217 + f"rank {r} passed split_ranks={gathered[r]}, but rank 0 "
218 + f"passed split_ranks={gathered[0]}. PyTorch requires every "
219 + "parent rank to enter split_group with the same split_ranks."
220 + )
221 +
222 +
223 +@pytest.mark.skipif(
224 + torch.accelerator.device_count() < 4,
225 + reason="Need at least 4 GPUs to run the test.",
226 +)
227 +def test_split_group_contract_same_split_ranks_on_all_ranks():
228 + """All parent ranks must call torch.distributed.split_group with the same
229 + ``split_ranks`` argument. This catches the bug where each rank passed
230 + only its own subgroup (``split_ranks=[ranks]``), which NCCL forgives for
231 + disjoint partitions but is a documented contract violation.
232 + """
233 + distributed_run(split_group_contract_worker, 4)
tests/distributed/test_torchrun_example.py
+15 -2
@@ -5,13 +5,26 @@
5 import os
6 import random
7
8 +import torch
9 import torch.distributed as dist
10
11 +import vllm.envs as envs
12 from vllm import LLM, SamplingParams
13 from vllm.distributed.parallel_state import get_world_group
14
13 -# Let PyTorch choose the WORLD backend for the current device type.
14 -dist.init_process_group()
15 +# By default, let PyTorch choose the WORLD backend for the current device
16 +# type (legacy lazy-init path). When VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1,
17 +# use the explicit eager-init pattern required by `split_group` (mixed
18 +# cpu:gloo,cuda:nccl backend + device_id binding).
19 +if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
20 + local_rank = int(os.environ["LOCAL_RANK"])
21 + torch.accelerator.set_device_index(local_rank)
22 + dist.init_process_group(
23 + backend="cpu:gloo,cuda:nccl",
24 + device_id=torch.device(f"cuda:{local_rank}"),
25 + )
26 +else:
27 + dist.init_process_group()
28
29 # Create prompts
30 prompts = [
tests/distributed/test_torchrun_example_moe.py
+15 -2
@@ -5,13 +5,26 @@
5 import os
6 import random
7
8 +import torch
9 import torch.distributed as dist
10
11 +import vllm.envs as envs
12 from vllm import LLM, SamplingParams
13 from vllm.distributed.parallel_state import get_tp_group, get_world_group
14
13 -# Let PyTorch choose the WORLD backend for the current device type.
14 -dist.init_process_group()
15 +# By default, let PyTorch choose the WORLD backend for the current device
16 +# type (legacy lazy-init path). When VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1,
17 +# use the explicit eager-init pattern required by `split_group` (mixed
18 +# cpu:gloo,cuda:nccl backend + device_id binding).
19 +if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
20 + local_rank = int(os.environ["LOCAL_RANK"])
21 + torch.accelerator.set_device_index(local_rank)
22 + dist.init_process_group(
23 + backend="cpu:gloo,cuda:nccl",
24 + device_id=torch.device(f"cuda:{local_rank}"),
25 + )
26 +else:
27 + dist.init_process_group()
28
29 # Create prompts
30 prompts = [
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+10 -1
@@ -10,6 +10,7 @@ import torch
10 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
11 from typing_extensions import ParamSpec
12
13 +import vllm.envs as envs
14 from vllm.config import VllmConfig, set_current_vllm_config
15 from vllm.distributed import (
16 cleanup_dist_env_and_memory,
@@ -60,7 +61,15 @@ def _set_vllm_config(
61 tensor_model_parallel_size=vllm_config.parallel_config.tensor_parallel_size,
62 pipeline_model_parallel_size=vllm_config.parallel_config.pipeline_parallel_size,
63 )
63 - cpu_group = torch.distributed.new_group(list(range(world_size)), backend="gloo")
64 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
65 + cpu_group = torch.distributed.split_group(
66 + split_ranks=[list(range(world_size))],
67 + group_desc="moe_test_cpu",
68 + )
69 + else:
70 + cpu_group = torch.distributed.new_group(
71 + list(range(world_size)), backend="gloo"
72 + )
73 return cpu_group
74
75
tests/kernels/moe/test_deepep_deepgemm_moe.py
+8 -1
@@ -14,6 +14,7 @@ import torch.distributed
14 from torch.distributed import ProcessGroup
15 from typing_extensions import ParamSpec
16
17 +import vllm.envs as envs
18 from vllm.config import VllmConfig, set_current_vllm_config
19 from vllm.forward_context import set_forward_context
20 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -375,7 +376,13 @@ def _test_deepep_deepgemm_moe(
376 w1_scale = w1_scale.to(device=device)
377 w2_scale = w2_scale.to(device=device)
378
378 - pg = torch.distributed.new_group(list(range(pgi.world_size)))
379 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
380 + pg = torch.distributed.split_group(
381 + split_ranks=[list(range(pgi.world_size))],
382 + group_desc="deepep_deepgemm_test",
383 + )
384 + else:
385 + pg = torch.distributed.new_group(list(range(pgi.world_size)))
386 test_tensors = TestTensors.make(config, pgi.rank)
387 block_shape = [w1.size(1) // w1_scale.size(1), w1.size(2) // w1_scale.size(2)]
388
tests/kernels/moe/test_deepep_moe.py
+8 -1
@@ -10,6 +10,7 @@ import pytest
10 import torch.distributed
11 from torch.distributed import ProcessGroup
12
13 +import vllm.envs as envs
14 from tests.kernels.moe.utils import make_dummy_moe_config
15 from vllm import _custom_ops as ops
16 from vllm.config import VllmConfig, set_current_vllm_config
@@ -375,7 +376,13 @@ def _deep_ep_moe(
376 w1_scale = w1_scale.to(device=device_idx)
377 w2_scale = w2_scale.to(device=device_idx)
378
378 - pg = torch.distributed.new_group(list(range(pgi.world_size)))
379 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
380 + pg = torch.distributed.split_group(
381 + split_ranks=[list(range(pgi.world_size))],
382 + group_desc="deepep_test",
383 + )
384 + else:
385 + pg = torch.distributed.new_group(list(range(pgi.world_size)))
386 test_tensors = TestTensors.make(config, low_latency_mode)
387
388 with set_current_vllm_config(VllmConfig()):
vllm/distributed/parallel_state.py
+177 -25
@@ -227,6 +227,67 @@ def patched_fused_scaled_matmul_reduce_scatter_fake(
227 return res
228
229
230 +def _platform_device_type() -> str:
231 + """Return the device-type string (e.g. ``"cuda"``, ``"xpu"``, ``"cpu"``)
232 + for the current platform, in the form expected by
233 + ``torch.distributed.init_process_group(backend=...)``.
234 + """
235 + from vllm.platforms import current_platform
236 +
237 + if current_platform.is_cuda_alike():
238 + return "cuda"
239 + elif current_platform.is_xpu():
240 + return "xpu"
241 + elif current_platform.is_out_of_tree():
242 + return current_platform.device_name
243 + else:
244 + return "cpu"
245 +
246 +
247 +def _device_backend_str(torch_distributed_backend: str | Backend) -> str:
248 + """Normalize ``torch_distributed_backend`` to the ``"<device>:<backend>"``
249 + format required by ``split_group``'s ``backend`` argument.
250 +
251 + Accepts either a bare backend name (e.g. ``"nccl"``) or an already-prefixed
252 + string (e.g. ``"cuda:nccl"``).
253 + """
254 + backend_str = str(torch_distributed_backend)
255 + if ":" in backend_str:
256 + return backend_str
257 + return f"{_platform_device_type()}:{backend_str}"
258 +
259 +
260 +def _create_subgroups_split_group(
261 + group_ranks: list[list[int]],
262 + group_name: str,
263 + torch_distributed_backend: str | Backend,
264 +) -> tuple[ProcessGroup, ProcessGroup]:
265 + """Create the device + CPU subgroups for ``GroupCoordinator`` via
266 + ``torch.distributed.split_group``.
267 +
268 + ``split_group`` is collective on the parent group, so every parent rank
269 + must enter with the same ``split_ranks`` definition. Each rank receives
270 + the subgroup it belongs to.
271 + """
272 + device_backend_str = _device_backend_str(torch_distributed_backend)
273 + self_device_group = torch.distributed.split_group(
274 + split_ranks=group_ranks,
275 + group_desc=f"{group_name}:device",
276 + backend=device_backend_str,
277 + )
278 + # CPU subgroup: split_group requires the requested backend filter to
279 + # include the parent's default device type (= the device the parent PG
280 + # was bound to via ``device_id``), so a cpu-only filter is rejected.
281 + # Include the device backend in the filter; only the gloo backend is
282 + # actually used for CPU collectives on this group.
283 + self_cpu_group = torch.distributed.split_group(
284 + split_ranks=group_ranks,
285 + group_desc=f"{group_name}:cpu",
286 + backend=f"cpu:gloo,{device_backend_str}",
287 + )
288 + return self_device_group, self_cpu_group
289 +
290 +
291 def patched_fused_scaled_matmul_reduce_scatter(
292 A: torch.Tensor,
293 B: torch.Tensor,
@@ -335,26 +396,39 @@ class GroupCoordinator:
396 self_device_group = None
397 self_cpu_group = None
398
338 - from vllm.distributed.utils import get_cpu_distributed_timeout_or_none
399 + # VLLM_DISTRIBUTED_USE_SPLIT_GROUP gates the new ``split_group``
400 + # codepath. Default (False) preserves the legacy ``new_group`` path.
401 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
402 + self_device_group, self_cpu_group = _create_subgroups_split_group(
403 + group_ranks, group_name, torch_distributed_backend
404 + )
405 + for ranks in group_ranks:
406 + if self.rank in ranks:
407 + self.ranks = ranks
408 + self.world_size = len(ranks)
409 + self.rank_in_group = ranks.index(self.rank)
410 + break
411 + else:
412 + from vllm.distributed.utils import get_cpu_distributed_timeout_or_none
413
340 - timeout = get_cpu_distributed_timeout_or_none()
414 + timeout = get_cpu_distributed_timeout_or_none()
415
342 - for ranks in group_ranks:
343 - device_group = torch.distributed.new_group(
344 - ranks, backend=torch_distributed_backend
345 - )
346 - # a group with `gloo` backend, to allow direct coordination between
347 - # processes through the CPU.
348 - with suppress_stdout():
349 - cpu_group = torch.distributed.new_group(
350 - ranks, backend="gloo", timeout=timeout
416 + for ranks in group_ranks:
417 + device_group = torch.distributed.new_group(
418 + ranks, backend=torch_distributed_backend
419 )
352 - if self.rank in ranks:
353 - self.ranks = ranks
354 - self.world_size = len(ranks)
355 - self.rank_in_group = ranks.index(self.rank)
356 - self_device_group = device_group
357 - self_cpu_group = cpu_group
420 + # a group with `gloo` backend, to allow direct coordination between
421 + # processes through the CPU.
422 + with suppress_stdout():
423 + cpu_group = torch.distributed.new_group(
424 + ranks, backend="gloo", timeout=timeout
425 + )
426 + if self.rank in ranks:
427 + self.ranks = ranks
428 + self.world_size = len(ranks)
429 + self.rank_in_group = ranks.index(self.rank)
430 + self_device_group = device_group
431 + self_cpu_group = cpu_group
432
433 assert self_cpu_group is not None
434 assert self_device_group is not None
@@ -1348,6 +1422,62 @@ def set_custom_all_reduce(enable: bool):
1422 _ENABLE_CUSTOM_ALL_REDUCE = enable
1423
1424
1425 +def _init_process_group_for_split_group(
1426 + *,
1427 + backend: str,
1428 + distributed_init_method: str,
1429 + world_size: int,
1430 + rank: int,
1431 + local_rank: int,
1432 + timeout: timedelta | None,
1433 +) -> None:
1434 + """Initialize the default PG with both CPU (gloo) and device (e.g. nccl)
1435 + backends and an eager ``device_id`` binding so that subgroups can be
1436 + created via ``split_group`` (which requires the parent communicator to
1437 + be eagerly initialized). Falls back to ``gloo`` on CPU-only systems.
1438 + """
1439 + if torch.accelerator.is_available() and backend != "gloo":
1440 + init_backend = "cpu:gloo,cuda:nccl"
1441 + device_id: torch.device | None = torch.device(f"cuda:{local_rank}")
1442 + else:
1443 + init_backend = "gloo"
1444 + device_id = None
1445 + torch.distributed.init_process_group(
1446 + backend=init_backend,
1447 + init_method=distributed_init_method,
1448 + world_size=world_size,
1449 + rank=rank,
1450 + timeout=timeout,
1451 + device_id=device_id,
1452 + )
1453 +
1454 +
1455 +def _validate_default_pg_for_split_group() -> None:
1456 + """When an external launcher (e.g. ``torchrun``) initialized the default
1457 + PG, ``GroupCoordinator`` cannot patch in additional backends or change
1458 + the eager-init behavior — ``split_group`` only selects subsets of an
1459 + existing parent. Validate that the parent has both ``device_id`` and a
1460 + CPU (gloo) backend, and emit a descriptive error pointing at the exact
1461 + init call to update otherwise.
1462 + """
1463 + default_pg = torch.distributed.distributed_c10d._get_default_group()
1464 + assert default_pg.bound_device_id is not None, (
1465 + "External launcher initialized the default process group "
1466 + "without device_id. vLLM requires the default PG to be device-"
1467 + "bound for split_group. Pass device_id=torch.device(f'cuda:"
1468 + "{local_rank}') to torch.distributed.init_process_group()."
1469 + )
1470 + try:
1471 + default_pg._get_backend(torch.device("cpu"))
1472 + except RuntimeError as e:
1473 + raise RuntimeError(
1474 + "External launcher initialized the default process group "
1475 + "without a CPU (gloo) backend. vLLM requires both CPU and "
1476 + "device backends. Pass backend='cpu:gloo,cuda:nccl' to "
1477 + "torch.distributed.init_process_group()."
1478 + ) from e
1479 +
1480 +
1481 def _init_elastic_ep_world(
1482 config, local_rank: int, backend: str, rank: int, world_size: int
1483 ) -> None:
@@ -1456,14 +1586,33 @@ def init_distributed_environment(
1586 "Fallback Gloo backend is not available."
1587 )
1588 backend = "gloo"
1459 - # this backend is used for WORLD
1460 - torch.distributed.init_process_group(
1461 - backend=backend,
1462 - init_method=distributed_init_method,
1463 - world_size=world_size,
1464 - rank=rank,
1465 - timeout=timeout,
1466 - )
1589 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
1590 + # split_group needs local_rank early to compute device_id for
1591 + # the eager init. local_rank is not available in torch
1592 + # ProcessGroup, see https://github.com/pytorch/pytorch/issues/122816
1593 + if local_rank == -1:
1594 + local_rank = (
1595 + int(envs.LOCAL_RANK)
1596 + if distributed_init_method == "env://"
1597 + else rank
1598 + )
1599 + _init_process_group_for_split_group(
1600 + backend=backend,
1601 + distributed_init_method=distributed_init_method,
1602 + world_size=world_size,
1603 + rank=rank,
1604 + local_rank=local_rank,
1605 + timeout=timeout,
1606 + )
1607 + else:
1608 + # this backend is used for WORLD
1609 + torch.distributed.init_process_group(
1610 + backend=backend,
1611 + init_method=distributed_init_method,
1612 + world_size=world_size,
1613 + rank=rank,
1614 + timeout=timeout,
1615 + )
1616 if enable_elastic_ep:
1617 tp_pp_cpu_group = torch.distributed.new_group(
1618 backend="gloo", timeout=timeout
@@ -1476,6 +1625,9 @@ def init_distributed_environment(
1625 "Elastic EP is not yet supported with multi-node TP/PP"
1626 )
1627
1628 + if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP and torch.accelerator.is_available():
1629 + _validate_default_pg_for_split_group()
1630 +
1631 # set the local rank
1632 # local_rank is not available in torch ProcessGroup,
1633 # see https://github.com/pytorch/pytorch/issues/122816
vllm/envs.py
+8
@@ -63,6 +63,7 @@ if TYPE_CHECKING:
63 VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
64 VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
65 VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False
66 + VLLM_DISTRIBUTED_USE_SPLIT_GROUP: bool = False
67 VLLM_XLA_USE_SPMD: bool = False
68 VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork"
69 VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets")
@@ -877,6 +878,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
878 "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool(
879 int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "1"))
880 ),
881 + # When True, GroupCoordinator constructs its CPU/device subgroups via
882 + # ``torch.distributed.split_group(backend=...)``
883 + # and ``init_distributed_environment`` initializes the default PG with
884 + # mixed ``cpu:gloo,cuda:nccl`` backend + eager ``device_id`` binding.
885 + "VLLM_DISTRIBUTED_USE_SPLIT_GROUP": lambda: bool(
886 + int(os.getenv("VLLM_DISTRIBUTED_USE_SPLIT_GROUP", "0"))
887 + ),
888 # Use dedicated multiprocess context for workers.
889 # Both spawn and fork work
890 "VLLM_WORKER_MULTIPROC_METHOD": env_with_choices(