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 11, 03:45 UTC
[CI] Fix H200 Distributed test (#31054)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Dec 20, 2025 at 16:48 UTC
ae0770fa6bcc58892692f0d1b2a3f0a300dbe82a
4 files changed
+51
-129
.buildkite/test-amd.yaml
+3
-3
@@ -1254,13 +1254,13 @@ steps:
1254
- # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
1255
- VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
1256
- NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
1257
- - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1257
+ - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1258
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py
1259
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py
1260
- # the following commands are for the second node, with ip 192.168.10.11 (ray environment already set up)
1261
- VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
1262
- NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
1263
- - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1263
+ - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1264
1265
- label: Distributed Tests (2 GPUs) # 68min
1266
timeout_in_minutes: 90
@@ -1508,7 +1508,7 @@ steps:
1508
- "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1509
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1510
- pytest -v -s tests/distributed/test_context_parallel.py
1511
- - HIP_VISIBLE_DEVICES=0,1 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048 --all2all-backend deepep_high_throughput
1511
+ - HIP_VISIBLE_DEVICES=0,1 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
1512
- pytest -v -s tests/v1/distributed/test_dbo.py
1513
1514
##### B200 test #####
.buildkite/test-pipeline.yaml
+3
-3
@@ -1109,13 +1109,13 @@ steps:
1109
- # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
1110
- VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
1111
- NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
1112
- - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1112
+ - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1113
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py
1114
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py
1115
- # the following commands are for the second node, with ip 192.168.10.11 (ray environment already set up)
1116
- VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
1117
- NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
1118
- - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1118
+ - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
1119
1120
- label: Distributed Tests (2 GPUs) # 68min
1121
timeout_in_minutes: 90
@@ -1334,7 +1334,7 @@ steps:
1334
- "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1335
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1336
- pytest -v -s tests/distributed/test_context_parallel.py
1337
- - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048 --all2all-backend deepep_high_throughput
1337
+ - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
1338
- pytest -v -s tests/v1/distributed/test_dbo.py
1339
1340
##### B200 test #####
.buildkite/test_areas/distributed.yaml
+2
-2
@@ -145,7 +145,7 @@ steps:
145
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
146
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
147
- pytest -v -s tests/distributed/test_context_parallel.py
148
- - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048 --all2all-backend deepep_high_throughput
148
+ - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
149
- pytest -v -s tests/v1/distributed/test_dbo.py
150
151
- label: Distributed Tests (2 GPUs)(B200)
@@ -171,7 +171,7 @@ steps:
171
- tests/distributed/
172
- tests/examples/offline_inference/data_parallel.py
173
commands:
174
- - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:0bec63fa317e1fbd62e19b0fc31c43c81bf89077 "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code"
174
+ - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:0bec63fa317e1fbd62e19b0fc31c43c81bf89077 "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --nnodes=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code"
175
176
- label: Distributed NixlConnector PD accuracy (4 GPUs)
177
timeout_in_minutes: 30
examples/offline_inference/data_parallel.py
+43
-121
@@ -5,25 +5,25 @@ Usage:
5
Single node:
6
python examples/offline_inference/data_parallel.py \
7
--model="ibm-research/PowerMoE-3b" \
8
- --dp-size=2 \
9
- --tp-size=2
8
+ -dp=2 \
9
+ -tp=2
10
11
Multi-node:
12
Node 0 (assume the node has ip of 10.99.48.128):
13
python examples/offline_inference/data_parallel.py \
14
--model="ibm-research/PowerMoE-3b" \
15
- --dp-size=2 \
16
- --tp-size=2 \
17
- --node-size=2 \
15
+ -dp=2 \
16
+ -tp=2 \
17
+ --nnodes=2 \
18
--node-rank=0 \
19
--master-addr=10.99.48.128 \
20
--master-port=13345
21
Node 1:
22
python examples/offline_inference/data_parallel.py \
23
--model="ibm-research/PowerMoE-3b" \
24
- --dp-size=2 \
25
- --tp-size=2 \
26
- --node-size=2 \
24
+ -dp=2 \
25
+ -tp=2 \
26
+ --nnodes=2 \
27
--node-rank=1 \
28
--master-addr=10.99.48.128 \
29
--master-port=13345
@@ -32,103 +32,40 @@ Multi-node:
32
import os
33
from time import sleep
34
35
-from vllm import LLM, SamplingParams
35
+from vllm import LLM, EngineArgs, SamplingParams
36
from vllm.platforms import current_platform
37
+from vllm.utils.argparse_utils import FlexibleArgumentParser
38
from vllm.utils.network_utils import get_open_port
39
40
40
-def parse_args():
41
- import argparse
41
+def create_parser():
42
+ parser = FlexibleArgumentParser(description="Data Parallel Inference")
43
43
- parser = argparse.ArgumentParser(description="Data Parallel Inference")
44
- parser.add_argument(
45
- "--model",
46
- type=str,
47
- default="ibm-research/PowerMoE-3b",
48
- help="Model name or path",
49
- )
50
- parser.add_argument("--dp-size", type=int, default=2, help="Data parallel size")
51
- parser.add_argument("--tp-size", type=int, default=2, help="Tensor parallel size")
52
- parser.add_argument(
53
- "--node-size", type=int, default=1, help="Total number of nodes"
54
- )
55
- parser.add_argument(
56
- "--node-rank", type=int, default=0, help="Rank of the current node"
57
- )
58
- parser.add_argument(
59
- "--master-addr", type=str, default="", help="Master node IP address"
60
- )
61
- parser.add_argument("--master-port", type=int, default=0, help="Master node port")
62
- parser.add_argument(
63
- "--enforce-eager", action="store_true", help="Enforce eager mode execution."
64
- )
65
- parser.add_argument(
66
- "--trust-remote-code", action="store_true", help="Trust remote code."
67
- )
68
- parser.add_argument(
69
- "--max-num-seqs",
70
- type=int,
71
- default=64,
72
- help=("Maximum number of sequences to be processed in a single iteration."),
73
- )
74
- parser.add_argument(
75
- "--max-model-len",
76
- type=int,
77
- help=("Maximum number of tokens to be processed in a single iteration."),
44
+ # Add all engine args
45
+ EngineArgs.add_cli_args(parser)
46
+ parser.set_defaults(
47
+ model="ibm-research/PowerMoE-3b",
48
+ enable_expert_parallel=True,
49
)
50
+
51
+ # Add timeout (not in EngineArgs)
52
parser.add_argument(
53
"--timeout",
54
type=int,
55
default=300,
83
- help=("Number of seconds before unresponsive process is killed."),
84
- )
85
- parser.add_argument(
86
- "--gpu-memory-utilization",
87
- type=float,
88
- default=0.8,
89
- help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
90
- )
91
- parser.add_argument(
92
- "--enable-dbo",
93
- action="store_true",
94
- help=("Enable microbatched execution"),
95
- )
96
- parser.add_argument(
97
- "--compilation-config",
98
- type=int,
99
- help=("Compilation optimization (O) mode 0-3."),
100
- )
101
- parser.add_argument(
102
- "--quantization",
103
- type=str,
104
- )
105
- parser.add_argument(
106
- "--disable-expert-parallel",
107
- dest="enable_expert_parallel",
108
- action="store_false",
109
- help="Disable expert parallel (default: enabled).",
56
+ help="Number of seconds before unresponsive process is killed.",
57
)
111
- parser.set_defaults(enable_expert_parallel=True)
112
- return parser.parse_args()
58
+
59
+ return parser
60
61
62
def main(
116
- model,
63
dp_size,
64
local_dp_rank,
65
global_dp_rank,
66
dp_master_ip,
67
dp_master_port,
122
- GPUs_per_dp_rank,
123
- enforce_eager,
124
- enable_expert_parallel,
125
- trust_remote_code,
126
- max_num_seqs,
127
- max_model_len,
128
- compilation_config,
129
- gpu_memory_utilization,
130
- enable_dbo,
131
- quantization,
68
+ engine_args,
69
):
70
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
71
os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
@@ -173,19 +110,7 @@ def main(
110
)
111
112
# Create an LLM.
176
- llm = LLM(
177
- model=model,
178
- tensor_parallel_size=GPUs_per_dp_rank,
179
- enforce_eager=enforce_eager,
180
- enable_expert_parallel=enable_expert_parallel,
181
- trust_remote_code=trust_remote_code,
182
- max_num_seqs=max_num_seqs,
183
- max_model_len=max_model_len,
184
- gpu_memory_utilization=gpu_memory_utilization,
185
- enable_dbo=enable_dbo,
186
- quantization=quantization,
187
- compilation_config=compilation_config,
188
- )
113
+ llm = LLM(**engine_args)
114
outputs = llm.generate(prompts, sampling_params)
115
# Print the outputs.
116
for i, output in enumerate(outputs):
@@ -204,22 +129,29 @@ def main(
129
130
131
if __name__ == "__main__":
207
- args = parse_args()
132
+ parser = create_parser()
133
+ args = vars(parser.parse_args())
134
+
135
+ # Extract DP-specific args
136
+ dp_size = args.pop("data_parallel_size")
137
+ nnodes = args.get("nnodes", 1)
138
+ node_rank = args.get("node_rank", 0)
139
+ master_addr = args.get("master_addr", "")
140
+ master_port = args.get("master_port", 0)
141
+ timeout = args.pop("timeout")
142
209
- dp_size = args.dp_size
210
- tp_size = args.tp_size
211
- node_size = args.node_size
212
- node_rank = args.node_rank
143
+ # Remaining args are engine args
144
+ engine_args = args
145
214
- if node_size == 1:
146
+ if nnodes == 1:
147
dp_master_ip = "127.0.0.1"
148
dp_master_port = get_open_port()
149
else:
218
- dp_master_ip = args.master_addr
219
- dp_master_port = args.master_port
150
+ dp_master_ip = master_addr
151
+ dp_master_port = master_port
152
221
- assert dp_size % node_size == 0, "dp_size should be divisible by node_size"
222
- dp_per_node = dp_size // node_size
153
+ assert dp_size % nnodes == 0, "dp_size should be divisible by nnodes"
154
+ dp_per_node = dp_size // nnodes
155
156
from multiprocessing import Process
157
@@ -235,29 +167,19 @@ if __name__ == "__main__":
167
proc = Process(
168
target=main,
169
args=(
238
- args.model,
170
dp_size,
171
local_dp_rank,
172
global_dp_rank,
173
dp_master_ip,
174
dp_master_port,
244
- tp_size,
245
- args.enforce_eager,
246
- args.enable_expert_parallel,
247
- args.trust_remote_code,
248
- args.max_num_seqs,
249
- args.max_model_len,
250
- args.compilation_config,
251
- args.gpu_memory_utilization,
252
- args.enable_dbo,
253
- args.quantization,
175
+ engine_args,
176
),
177
)
178
proc.start()
179
procs.append(proc)
180
exit_code = 0
181
for proc in procs:
260
- proc.join(timeout=args.timeout)
182
+ proc.join(timeout=timeout)
183
if proc.exitcode is None:
184
print(f"Killing process {proc.pid} that didn't stop within 5 minutes.")
185
proc.kill()