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 14, 05:30 UTC
Enable symmetric memory all reduce by default only enabling for TP (#25070)
Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Ilya Markov committed
Sep 23, 2025 at 21:53 UTC
8bdd8b5c5155114f22af6d99f97176a15b6f8bf5
4 files changed
+56
-16
.buildkite/test-pipeline.yaml
+2
@@ -164,6 +164,7 @@ steps:
164
- tests/v1/test_internal_lb_dp.py
165
- tests/v1/test_hybrid_lb_dp.py
166
- tests/v1/engine/test_engine_core_client.py
167
+ - tests/distributed/test_symm_mem_allreduce.py
168
commands:
169
# test with torchrun tp=2 and external_dp=2
170
- torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
@@ -188,6 +189,7 @@ steps:
189
- pytest -v -s compile/test_basic_correctness.py
190
- pytest -v -s distributed/test_pynccl.py
191
- pytest -v -s distributed/test_events.py
192
+ - pytest -v -s distributed/test_symm_mem_allreduce.py
193
# TODO: create a dedicated test section for multi-GPU example tests
194
# when we have multiple distributed example tests
195
- pushd ../examples/offline_inference
tests/distributed/test_symm_mem_allreduce.py
+47
-12
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import queue
5
import random
6
import typing
7
@@ -10,26 +11,31 @@ import torch.distributed as dist
11
import torch.multiprocessing as mp
12
13
import vllm.envs as envs
14
+from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
15
from vllm.distributed import cleanup_dist_env_and_memory
16
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
17
from vllm.distributed.device_communicators.cuda_communicator import (
18
CudaCommunicator)
17
-from vllm.distributed.parallel_state import (get_tensor_model_parallel_group,
18
- get_tp_group,
19
+from vllm.distributed.parallel_state import (get_tp_group,
20
init_distributed_environment,
21
initialize_model_parallel)
22
+from vllm.engine.arg_utils import EngineArgs
23
+from vllm.engine.llm_engine import LLMEngine
24
from vllm.platforms import current_platform
25
from vllm.utils import update_environment_variables
26
27
torch.manual_seed(42)
28
random.seed(44)
29
27
-test_size_elements = 4 * 1024 * 1024
30
+test_size_elements = 1024 * 1024
31
32
30
-def symm_mem_allreduce_worker(local_rank: int, world_size: int):
33
+def symm_mem_allreduce_worker(local_rank: int, world_size: int, q: mp.Queue):
34
monkeypatch = pytest.MonkeyPatch()
32
- with monkeypatch.context() as m:
35
+ config = VllmConfig(parallel_config=ParallelConfig(
36
+ tensor_parallel_size=world_size))
37
+
38
+ with monkeypatch.context() as m, set_current_vllm_config(config):
39
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
40
dtype = torch.bfloat16
41
device = torch.device(f"cuda:{local_rank}")
@@ -51,22 +57,26 @@ def symm_mem_allreduce_worker(local_rank: int, world_size: int):
57
get_tp_group().device_communicator)
58
symm_mem_comm = cuda_communicator.symm_mem_comm
59
if symm_mem_comm is None or symm_mem_comm.disabled:
54
- pytest.skip("SymmMemCommunicator is not available or disabled.")
60
+ # can't use skip under multiprocessing
61
+ q.put("SymmMemCommunicator is not available or disabled.")
62
+ return
63
64
inp_direct_symm_mem = torch.randint(1,
65
23, (test_size_elements, ),
66
dtype=dtype,
67
device=device)
68
if not symm_mem_comm.should_use_symm_mem(inp_direct_symm_mem):
61
- pytest.skip(
69
+ # can't use skip under multiprocessing
70
+ q.put(
71
"SymmMemCommunicator isn't used for this world and input size."
72
)
73
+ return
74
75
original_inp_direct_symm_mem = inp_direct_symm_mem.clone()
76
out_direct_symm_mem = symm_mem_comm.all_reduce(inp_direct_symm_mem)
77
assert out_direct_symm_mem is not None
78
69
- group = get_tensor_model_parallel_group().device_group
79
+ group = get_tp_group().device_group
80
dist.all_reduce(original_inp_direct_symm_mem, group=group)
81
torch.testing.assert_close(out_direct_symm_mem,
82
original_inp_direct_symm_mem,
@@ -100,9 +110,34 @@ def test_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, tp_size,
110
world_size = tp_size * pipeline_parallel_size
111
if world_size > torch.cuda.device_count():
112
pytest.skip("Not enough GPUs to run the test.")
113
+ q = mp.get_context('spawn').Queue()
114
+ mp.spawn(symm_mem_allreduce_worker,
115
+ args=(world_size, q),
116
+ nprocs=world_size)
117
+ try:
118
+ val = q.get(timeout=1)
119
+ except queue.Empty:
120
+ val = None
121
+ finally:
122
+ cleanup_dist_env_and_memory()
123
+ if val is not None:
124
+ pytest.skip(val)
125
104
- # Enable SymmMemCommunicator
105
- monkeypatch.setenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1")
126
107
- mp.spawn(symm_mem_allreduce_worker, args=(world_size, ), nprocs=world_size)
108
- cleanup_dist_env_and_memory()
127
+@pytest.mark.skipif(
128
+ not current_platform.is_cuda(),
129
+ reason="SymmMemAllreduce is only available for CUDA platforms.")
130
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
131
+ reason="Only test on CUDA")
132
+def test_dp_with_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch):
133
+ world_size = 4
134
+ if world_size > torch.cuda.device_count():
135
+ pytest.skip("Not enough GPUs to run the test.")
136
+ # Verify that the DataParallel runs without error
137
+ engine_args = EngineArgs(model="distilbert/distilgpt2",
138
+ enforce_eager=True,
139
+ enable_prefix_caching=True,
140
+ data_parallel_size=2,
141
+ tensor_parallel_size=2,
142
+ data_parallel_backend="mp")
143
+ LLMEngine.from_engine_args(engine_args)
vllm/distributed/device_communicators/cuda_communicator.py
+5
-2
@@ -30,18 +30,21 @@ class CudaCommunicator(DeviceCommunicatorBase):
30
unique_name: str = ""):
31
super().__init__(cpu_group, device, device_group, unique_name)
32
if "tp" not in unique_name:
33
- # only tp uses custom allreduce
33
+ # custom allreduce or torch symm mem can be used only by tp
34
use_custom_allreduce = False
35
+ use_torch_symm_mem = False
36
else:
37
from vllm.distributed.parallel_state import (
38
_ENABLE_CUSTOM_ALL_REDUCE)
39
use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
40
+ use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
41
42
# ep does not use pynccl
43
use_pynccl = "ep" not in unique_name
44
45
self.use_pynccl = use_pynccl
46
self.use_custom_allreduce = use_custom_allreduce
47
+ self.use_torch_symm_mem = use_torch_symm_mem
48
49
# lazy import to avoid documentation build error
50
from vllm.distributed.device_communicators.custom_all_reduce import (
@@ -65,7 +68,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
68
self.ca_comm: Optional[CustomAllreduce] = None
69
self.qr_comm: Optional[QuickAllReduce] = None
70
self.symm_mem_comm: Optional[SymmMemCommunicator] = None
68
- if envs.VLLM_ALLREDUCE_USE_SYMM_MEM and current_platform.is_cuda():
71
+ if use_torch_symm_mem and current_platform.is_cuda():
72
self.symm_mem_comm = SymmMemCommunicator(
73
group=self.cpu_group,
74
device=self.device,
vllm/envs.py
+2
-2
@@ -182,7 +182,7 @@ if TYPE_CHECKING:
182
VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
183
VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
184
VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
185
- VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False
185
+ VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
186
VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
187
VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False
188
VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
@@ -1370,7 +1370,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
1370
1371
# Whether to use pytorch symmetric memory for allreduce
1372
"VLLM_ALLREDUCE_USE_SYMM_MEM":
1373
- lambda: bool(int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "0"))),
1373
+ lambda: bool(int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1"))),
1374
1375
# Allows vllm to find tuned config under customized folder
1376
"VLLM_TUNED_CONFIG_FOLDER":