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, 07:15 UTC
[BugFix] Fix multi-node offline data-parallel (#18981)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Yizhou Liu <liu_yizhou@outlook.com>
Nick Hill committed
May 31, 2025 at 08:34 UTC
9a1b9b99d7d21a65280fc68c82f5bb6152fdf9dd
2 files changed
+11
-8
examples/offline_inference/data_parallel.py
+8
-4
@@ -97,10 +97,14 @@ def main(
97
# with DP, each rank should process different prompts.
98
# usually all the DP ranks process a full dataset,
99
# and each rank processes a different part of the dataset.
100
- promts_per_rank = len(prompts) // dp_size
101
- start = global_dp_rank * promts_per_rank
102
- end = start + promts_per_rank
103
- prompts = prompts[start:end]
100
+ floor = len(prompts) // dp_size
101
+ remainder = len(prompts) % dp_size
102
+
103
+ # Distribute prompts into even groups.
104
+ def start(rank):
105
+ return rank * floor + min(rank, remainder)
106
+
107
+ prompts = prompts[start(global_dp_rank) : start(global_dp_rank + 1)]
108
if len(prompts) == 0:
109
# if any rank has no prompts to process,
110
# we need to set a placeholder prompt
vllm/v1/engine/core_client.py
+3
-4
@@ -363,6 +363,7 @@ class MPClient(EngineCoreClient):
363
local_engine_count = parallel_config.data_parallel_size_local
364
local_start_index = parallel_config.data_parallel_rank_local
365
dp_size = parallel_config.data_parallel_size
366
+ dp_rank = parallel_config.data_parallel_rank
367
368
# SPMD mode is where there is an LLM instance per DP rank and
369
# one core engine per LLM, see
@@ -370,11 +371,9 @@ class MPClient(EngineCoreClient):
371
spmd_mode = local_start_index is not None
372
if spmd_mode:
373
assert local_engine_count == 1
373
- self.core_engines = [
374
- CoreEngine(index=local_start_index, local=True)
375
- ]
374
+ self.core_engines = [CoreEngine(index=dp_rank, local=True)]
375
else:
377
- assert parallel_config.data_parallel_rank == 0
376
+ assert dp_rank == 0
377
local_start_index = 0
378
self.core_engines = [
379
CoreEngine(index=i, local=(i < local_engine_count))