@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3d2779c29a

Fork 0

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, 09:45 UTC

[Feature] Support Pipeline Parallism in torchrun SPMD offline inference for V1 (#17827)

Signed-off-by: Lucia Fang <fanglu@fb.com>

Lucia Fang committed May 15, 2025 at 22:28 UTC 3d2779c29a9f5003f6fec6ca07205147e2c987d1

9 files changed +55 -27

.buildkite/test-pipeline.yaml

@@ -148,6 +148,8 @@ steps:

148 # test with tp=2 and external_dp=2

149 - VLLM_USE_V1=0 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py

150 - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py

151 + # test with tp=2 and pp=2

152 + - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py

153 # test with internal dp

154 - python3 ../examples/offline_inference/data_parallel.py

155 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py

examples/offline_inference/torchrun_example.py

+14 -9

@@ -8,6 +8,8 @@ the argument 2 should match the `tensor_parallel_size` below.

8 see `tests/distributed/test_torchrun_example.py` for the unit test.

9 """

11 +import torch.distributed as dist

12 +

13 from vllm import LLM, SamplingParams

15 # Create prompts, the same across all ranks

@@ -27,23 +29,26 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

29 # all ranks have the same random seed, so that sampling can be

30 # deterministic across ranks.

31 llm = LLM(

30 - model="facebook/opt-125m",

32 + model="meta-llama/Llama-3.1-8B",

33 tensor_parallel_size=2,

34 + pipeline_parallel_size=2,

35 distributed_executor_backend="external_launcher",

33 - seed=0,

36 + max_model_len=32768,

37 + seed=1,

38 )

40 outputs = llm.generate(prompts, sampling_params)

42 # all ranks will have the same outputs

39 -print("-" * 50)

40 -for output in outputs:

41 - prompt = output.prompt

42 - generated_text = output.outputs[0].text

43 - print(f"Prompt: {prompt!r}\n"

44 - f"Generated text: {generated_text!r}")

43 +if dist.get_rank() == 0:

44 print("-" * 50)

46 -"""

45 + for output in outputs:

46 + prompt = output.prompt

47 + generated_text = output.outputs[0].text

48 + print(f"Prompt: {prompt!r}\n"

49 + f"Generated text: {generated_text!r}\n")

50 + print("-" * 50)

51 + """

52 Further tips:

54 1. to communicate control messages across all ranks, use the cpu group,

tests/distributed/test_torchrun_example.py

+2 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 # unit test for `examples/offline_inference/torchrun_example.py`

4 -

4 +import os

5 import random

7 import torch.distributed as dist

@@ -25,6 +25,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

25 # to test if all ranks agree on the same kv cache configuration.

26 llm = LLM(model="facebook/opt-125m",

27 tensor_parallel_size=2,

28 + pipeline_parallel_size=int(os.getenv("PP_SIZE", 1)),

29 distributed_executor_backend="external_launcher",

30 gpu_memory_utilization=random.uniform(0.7, 0.9),

31 swap_space=random.randint(1, 4),

vllm/config.py

-1

@@ -1695,7 +1695,6 @@ class ParallelConfig:

1695 """Port of the data parallel master."""

1696 enable_expert_parallel: bool = False

1697 """Use expert parallelism instead of tensor parallelism for MoE layers."""

1698 -

1698 max_parallel_loading_workers: Optional[int] = None

1699 """Maximum number of parallel loading workers when loading model

1700 sequentially in multiple batches. To avoid RAM OOM when using tensor

vllm/distributed/device_communicators/custom_all_reduce.py

+4 -2

@@ -265,7 +265,8 @@ class CustomAllreduce:

265

266 def close(self):

267 if not self.disabled and self._ptr:

268 - ops.dispose(self._ptr)

268 + if ops is not None:

269 + ops.dispose(self._ptr)

270 self._ptr = 0

271 self.free_shared_buffer(self.meta_ptrs, rank=self.rank)

272 self.free_shared_buffer(self.buffer_ptrs, rank=self.rank)

@@ -298,4 +299,5 @@ class CustomAllreduce:

299 rank: Optional[int] = 0) -> None:

300 if rank is None:

301 rank = dist.get_rank(group=group)

301 - ops.free_shared_buffer(pointers[rank])

302 + if ops is not None:

303 + ops.free_shared_buffer(pointers[rank])

vllm/engine/arg_utils.py

+3 -2

@@ -1383,9 +1383,10 @@ class EngineArgs:

1383 return False

1384

1385 if (self.pipeline_parallel_size > 1

1386 - and self.distributed_executor_backend not in ["ray", "mp"]):

1386 + and self.distributed_executor_backend

1387 + not in ("ray", "mp", "external_launcher")):

1388 name = "Pipeline Parallelism without Ray distributed executor " \

1388 - "or multiprocessing executor"

1389 + "or multiprocessing executor or external launcher"

1390 _raise_or_fallback(feature_name=name, recommend_to_remove=False)

1391 return False

1392

vllm/executor/uniproc_executor.py

-3

@@ -86,9 +86,6 @@ class ExecutorWithExternalLauncher(UniProcExecutor):

86 def _init_executor(self) -> None:

87 """Initialize the worker and load the model.

88 """

89 - assert self.vllm_config.parallel_config.pipeline_parallel_size == 1, \

90 - ("ExecutorWithExternalLauncher does not "

91 - "support pipeline parallelism.")

89 assert self.vllm_config.scheduler_config.delay_factor == 0.0, \

90 ("ExecutorWithExternalLauncher needs deterministic "

91 "execution, so it"

vllm/v1/worker/gpu_model_runner.py

+27 -6

@@ -22,7 +22,8 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,

22 has_kv_transfer_group)

23 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

24 from vllm.distributed.parallel_state import (

25 - get_pp_group, graph_capture, prepare_communication_buffer_for_model)

25 + get_pp_group, get_tp_group, graph_capture,

26 + prepare_communication_buffer_for_model)

27 from vllm.forward_context import get_forward_context, set_forward_context

28 from vllm.logger import init_logger

29 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

@@ -1162,13 +1163,32 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1163 hidden_states, aux_hidden_states = model_output

1164 else:

1165 hidden_states = model_output

1165 -

1166 + # Broadcast PP output for external_launcher (torchrun)

1167 + # to make sure we are synced across pp ranks

1168 + # TODO: Support overlapping mirco-batches

1169 + # https://github.com/vllm-project/vllm/issues/18019

1170 + broadcast_pp_output = \

1171 + self.parallel_config.distributed_executor_backend \

1172 + == "external_launcher" and len(get_pp_group().ranks) > 0

1173 if not get_pp_group().is_last_rank:

1174 # For mid-pipeline stages, return the hidden states.

1168 - return hidden_states

1169 -

1170 - sample_hidden_states = hidden_states[logits_indices]

1171 - logits = self.model.compute_logits(sample_hidden_states, None)

1175 + if not broadcast_pp_output:

1176 + return hidden_states

1177 + assert isinstance(hidden_states, IntermediateTensors)

1178 + get_pp_group().send_tensor_dict(hidden_states.tensors,

1179 + all_gather_group=get_tp_group())

1180 + logits = None

1181 + else:

1182 + sample_hidden_states = hidden_states[logits_indices]

1183 + logits = self.model.compute_logits(sample_hidden_states, None)

1184 + if broadcast_pp_output:

1185 + model_output_broadcast_data = {

1186 + "logits": logits.contiguous(),

1187 + } if logits is not None else {}

1188 + model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(

1189 + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)

1190 + assert model_output_broadcast_data is not None

1191 + logits = model_output_broadcast_data["logits"]

1192

1193 # Apply structured output bitmasks if present

1194 if scheduler_output.grammar_bitmask is not None:

@@ -1186,6 +1206,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1206 # creates a new tensor with separate storage from the original

1207 # logits tensor. This means any in-place operations on bonus_logits

1208 # won't affect the original logits tensor.

1209 + assert logits is not None

1210 bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]

1211 sampler_output = self.sampler(

1212 logits=bonus_logits,

vllm/v1/worker/gpu_worker.py

+3 -3

@@ -275,13 +275,13 @@ class Worker(WorkerBase):

275

276 output = self.model_runner.execute_model(scheduler_output,

277 intermediate_tensors)

278 -

279 - if not get_pp_group().is_last_rank:

278 + parallel_config = self.vllm_config.parallel_config

279 + if parallel_config.distributed_executor_backend != "external_launcher" \

280 + and not get_pp_group().is_last_rank:

281 assert isinstance(output, IntermediateTensors)

282 get_pp_group().send_tensor_dict(output.tensors,

283 all_gather_group=get_tp_group())

284 return None

284 -

285 assert isinstance(output, ModelRunnerOutput)

286 return output if self.is_driver_worker else None

287