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 15, 04:45 UTC
fix ci issue distributed 4 gpu test (#20204)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 28, 2025 at 01:50 UTC
d45417b804aaf7f90c9ae70a32f8f07d6b371a8c
1 file changed
+18
examples/offline_inference/data_parallel.py
+18
@@ -64,6 +64,18 @@ def parse_args():
64
parser.add_argument(
65
"--trust-remote-code", action="store_true", help="Trust remote code."
66
)
67
+ parser.add_argument(
68
+ "--max-num-seqs",
69
+ type=int,
70
+ default=64,
71
+ help=("Maximum number of sequences to be processed in a single iteration."),
72
+ )
73
+ parser.add_argument(
74
+ "--gpu-memory-utilization",
75
+ type=float,
76
+ default=0.8,
77
+ help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
78
+ )
79
return parser.parse_args()
80
81
@@ -77,6 +89,8 @@ def main(
89
GPUs_per_dp_rank,
90
enforce_eager,
91
trust_remote_code,
92
+ max_num_seqs,
93
+ gpu_memory_utilization,
94
):
95
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
96
os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
@@ -127,6 +141,8 @@ def main(
141
enforce_eager=enforce_eager,
142
enable_expert_parallel=True,
143
trust_remote_code=trust_remote_code,
144
+ max_num_seqs=max_num_seqs,
145
+ gpu_memory_utilization=gpu_memory_utilization,
146
)
147
outputs = llm.generate(prompts, sampling_params)
148
# Print the outputs.
@@ -181,6 +197,8 @@ if __name__ == "__main__":
197
tp_size,
198
args.enforce_eager,
199
args.trust_remote_code,
200
+ args.max_num_seqs,
201
+ args.gpu_memory_utilization,
202
),
203
)
204
proc.start()