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
create is_in_the_same_node on cpu (#26832)
Co-authored-by: Lunwen He <lunwenh@meta.com>
Lunwen He committed
Oct 20, 2025 at 19:04 UTC
0eb8f2b88004157da4cf7b41a72182a7c91d08a2
4 files changed
+26
-8
.buildkite/test-amd.yaml
+1
@@ -1081,6 +1081,7 @@ steps:
1081
- pytest -v -s ./compile/test_basic_correctness.py
1082
- pytest -v -s ./compile/test_wrapper.py
1083
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
1084
+ - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
1085
- pytest -v -s distributed/test_sequence_parallel.py
1086
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
1087
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
.buildkite/test-pipeline.yaml
+1
@@ -977,6 +977,7 @@ steps:
977
- pytest -v -s ./compile/test_basic_correctness.py
978
- pytest -v -s ./compile/test_wrapper.py
979
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
980
+ - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
981
- pytest -v -s distributed/test_sequence_parallel.py
982
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
983
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
tests/distributed/test_same_node.py
+21
-7
@@ -3,12 +3,25 @@
3
4
import os
5
6
+import torch
7
import torch.distributed as dist
8
9
from vllm.distributed.parallel_state import in_the_same_node_as
10
from vllm.distributed.utils import StatelessProcessGroup
11
from vllm.utils.network_utils import get_ip, get_open_port
12
13
+
14
+def _run_test(pg):
15
+ test_result = all(in_the_same_node_as(pg, source_rank=0))
16
+
17
+ expected = os.environ.get("VLLM_TEST_SAME_HOST", "1") == "1"
18
+ assert test_result == expected, f"Expected {expected}, got {test_result}"
19
+ if pg == dist.group.WORLD:
20
+ print("Same node test passed! when using torch distributed!")
21
+ else:
22
+ print("Same node test passed! when using StatelessProcessGroup!")
23
+
24
+
25
if __name__ == "__main__":
26
dist.init_process_group(backend="gloo")
27
@@ -25,11 +38,12 @@ if __name__ == "__main__":
38
stateless_pg = StatelessProcessGroup.create(ip, port, rank, dist.get_world_size())
39
40
for pg in [dist.group.WORLD, stateless_pg]:
28
- test_result = all(in_the_same_node_as(pg, source_rank=0))
29
-
30
- expected = os.environ.get("VLLM_TEST_SAME_HOST", "1") == "1"
31
- assert test_result == expected, f"Expected {expected}, got {test_result}"
32
- if pg == dist.group.WORLD:
33
- print("Same node test passed! when using torch distributed!")
41
+ if os.environ.get("VLLM_TEST_WITH_DEFAULT_DEVICE_SET", "0") == "1":
42
+ default_devices = ["cpu"]
43
+ if torch.cuda.is_available():
44
+ default_devices.append("cuda")
45
+ for device in default_devices:
46
+ torch.set_default_device(device)
47
+ _run_test(pg)
48
else:
35
- print("Same node test passed! when using StatelessProcessGroup!")
49
+ _run_test(pg)
vllm/distributed/parallel_state.py
+3
-1
@@ -1526,7 +1526,9 @@ def in_the_same_node_as(
1526
ranks = list(range(world_size))
1527
1528
# local tensor in each process to store the result
1529
- is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32)
1529
+ is_in_the_same_node = torch.tensor(
1530
+ [0] * world_size, dtype=torch.int32, device="cpu"
1531
+ )
1532
1533
magic_message = b"magic_message"
1534
shm = None