Pipeline Parallelism Support - Source Code Location

Features
1

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?

2

This function does the layer partitioning and could be your starting point: vllm/vllm/distributed/utils.py at d0cfec7ab919bfab261db419e17c768f08a24dc8 · vllm-project/vllm · GitHub