Pipeline Parallelism Support

[egebeysel]

I was looking to get inspired by how the pipeline parallelism support is implemented in vLLM. Therefore, I was looking for the source code that adjusts the model pipeline to be compatible with pipeline-parallel inference. Any pointers on where the code lives that enables this would be super welcome.

In other words, where does the model get partitioned into pipeline stages when we run vllm serve gpt2 --tensor-parallel-size 4 --pipeline-parallel-size 2? Or does the model already have to be in a compatible form?

[Shyboy0499]

Here's the map of where vLLM actually partitions the model for pipeline parallelism:

1. Layer partitioning — get_pp_indices() in vllm/distributed/utils.py
This is the core function that decides the split:

def get_pp_indices(num_hidden_layers, pp_rank, pp_size) -> tuple[int, int]:

It returns the (start_layer, end_layer) range for a given rank, distributing layers as evenly as possible — and deliberately giving the last partition one fewer layer, because that rank usually carries an extra final norm. Every model's base class (vllm/model_executor/models/transformers/base.py) calls this in __init__, so each PP rank only instantiates its own slice of the transformer stack.

2. Config — ParallelConfig in vllm/config/model.py
pipeline_parallel_size (CLI: --pipeline-parallel-size) is where the topology is declared; pp_rank is inferred per worker from it.

3. Group setup — vllm/distributed/parallel_state.py
After workers start, this builds the pipeline-parallel groups from ParallelConfig, wiring the send/recv peers between stages.

4. Execution — the worker's model runner (vllm/worker/model_runner.py)
Each PP rank's worker loads only its partition (the model constructed in step 1), and the runner drives the forward pass across stages.

Best starting points: get_pp_indices() together with tests/distributed/test_pipeline_partition.py — the test exercises the layer distribution in isolation, so you can see the expected split behavior (including the remainder-handling edge cases) without tracing the full runtime.

One caveat: vLLM has two engines (V0 and V1) with different distributed orchestration. The layer-slicing helpers above are the shared foundation used by the model classes, but stage wiring differs per engine — worth confirming which engine you're running (VLLM_USE_V1) before tracing deeper into the runner.