[V1][PP] Support PP for MultiprocExecutor

[bigPYJ1151]

By offloading MQ reading to a IO thread, MultiprocExecutor can receive multiple batches in flight and support V1 PP seamlessly.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[comaniac]

Thanks for adding this!
Please see the comments to align the v1 PP interface and design.
Meanwhile, do you have any benchmark numbers?

[comaniac]

Also cc @ruisearch42

[ruisearch42]

update comment above

[youkaichao]

since @comaniac and @ruisearch42 do not want to add non_block at the collective_rpc level, you can have a _run_workers function to have the ability to be non_blocking, and then in the execute_model level of this executor, call _run_workerswith non-blocking if pp is enabled.

[bigPYJ1151]

Good idea :)

[youkaichao]

why is it calculated this way?

[bigPYJ1151]

I thought the driver should be the first rank in the last TP group.

Now changed is_driver checking to self.rank % tp_size == 0 as the definition in V0.

[youkaichao]

also cc @njhill and @tlrmchlsmth for the mp-related change.

[bigPYJ1151]

Thanks for @comaniac @ruisearch42 @youkaichao 's comments, I have resolved them all.

I have some initial benchmark results. On a 4xL40 platform:

# TP=2 PP=1
VLLM_USE_V1=1 python3 benchmark_throughput.py --backend=vllm --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=meta-llama/Meta-Llama-3-8B-Instruct --n=1 --num-prompts=1000 --trust-remote-code --disable-log-stats -tp=2 -pp=1 --max-model-len=8192 
Throughput: 5.60 requests/s, 2314.67 total tokens/s, 1110.17 output tokens/s

# TP=1 PP=2
VLLM_USE_V1=1 python3 benchmark_throughput.py --backend=vllm --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=meta-llama/Meta-Llama-3-8B-Instruct --n=1 --num-prompts=1000 --trust-remote-code --disable-log-stats -tp=1 -pp=2 --max-model-len=8192 
Throughput: 5.74 requests/s, 2374.38 total tokens/s, 1138.81 output tokens/s

# TP=1 PP=4
VLLM_USE_V1=1 python3 benchmark_throughput.py --backend=vllm --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=meta-llama/Meta-Llama-3-8B-Instruct --n=1 --num-prompts=1000 --trust-remote-code --disable-log-stats -tp=4 -pp=1 --max-model-len=8192 
Throughput: 7.93 requests/s, 3277.62 total tokens/s, 1572.03 output tokens/s

# TP=1 PP=4
VLLM_USE_V1=1 python3 benchmark_throughput.py --backend=vllm --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=meta-llama/Meta-Llama-3-8B-Instruct --n=1 --num-prompts=1000 --trust-remote-code --disable-log-stats -tp=1 -pp=4 --max-model-len=8192 
Throughput: 9.28 requests/s, 3838.89 total tokens/s, 1841.22 output tokens/s

# TP=2 PP=2
VLLM_USE_V1=1 python3 benchmark_throughput.py --backend=vllm --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=meta-llama/Meta-Llama-3-8B-Instruct --n=1 --num-prompts=1000 --trust-remote-code --disable-log-stats -tp=2 -pp=2 --max-model-len=8192 
Throughput: 9.69 requests/s, 4005.70 total tokens/s, 1921.23 output tokens/s

[comaniac]

nit: We should be able to eliminate non_block here and just have

def _run_workers(self, ...):
    non_block = self.max_concurrent_batches > 1
    ...

[bigPYJ1151]

Now collective_rpc is using _run_workers. I think this will make all collective_rpc return future objects when using PP 🤔

[comaniac]

I see. It's a bit confusing tho. Could we try to add more comments to help future developments? For example:

    def _run_workers(self,
                     method: Union[str, Callable],
                     timeout: Optional[float] = None,
                     args: tuple = (),
                     kwargs: Optional[dict] = None,
                     non_block: bool = False) -> list[Any]:
        """Run the method on workers of all ranks and get their responses.
        Note that when non_block=True (PP > 1), this function immediately returns
        future objects to unblock pipeline.

        Args:
            ...
        
        Return:
            A list of responses from workers of all ranks. The response will be
            future objects if non_block=True.
        """

[comaniac]

This calculation may have issues when other parallelism is used (e.g., EP, DP attention)?

Suggested change

	# Note: only returns ModelRunnerOutput from the driver worker with the
	# last PP rank
	return output[self.world_size -
	self.parallel_config.tensor_parallel_size]
	# Only returns ModelRunnerOutput from TP rank=0 and PP rank=-1
	# (the first TP worker of the last PP stage).
	return output[self.world_size -
	self.parallel_config.tensor_parallel_size]

[youkaichao]

self.world_size - self.parallel_config.tensor_parallel_size is tp rank 0 with pp rank = 1. i think it is still not correct.

[comaniac]

Assuming TP=8, PP=4, then the world size is 32. IIUC, in this case the PP ranks would be
0-7, PP rank 0
8-15, PP rank 1
16-23, PP rank 2
24-31, PP rank 3

so world_size - tp_size = 32 - 8 = 24 should be PP rank = -1 (i.e. 3)?

[bigPYJ1151]

EP is using TP group so it will not have issue.

DP will have issue because the rank map shape is (DP, PP, TP). But right now the MP executor suppose there are only TP and PP and checking it in initialization.

[youkaichao]

oh my bad, this is indeed correct then. we should add more comment here, e.g. including the calculation from @comaniac . it is quite easy to mess things up here.

[comaniac]

LGTM.
Since we are in the process of making v1 as the default and this PR changes the executor, perhaps it's better to wait for a few more days until other issues are resolved.

Also cc @njhill @robertgshaw2-redhat @WoosukKwon for awareness.

[comaniac]

I see. It's a bit confusing tho. Could we try to add more comments to help future developments? For example:

    def _run_workers(self,
                     method: Union[str, Callable],
                     timeout: Optional[float] = None,
                     args: tuple = (),
                     kwargs: Optional[dict] = None,
                     non_block: bool = False) -> list[Any]:
        """Run the method on workers of all ranks and get their responses.
        Note that when non_block=True (PP > 1), this function immediately returns
        future objects to unblock pipeline.

        Args:
            ...
        
        Return:
            A list of responses from workers of all ranks. The response will be
            future objects if non_block=True.
        """

[WoosukKwon]

Hi @bigPYJ1151, can you please rebase the PR and resolve merge conflicts?

[bigPYJ1151]

Hi @bigPYJ1151, can you please rebase the PR and resolve merge conflicts?

@WoosukKwon Sure, updated, also verified the unit tests. Please take a look :)

[WoosukKwon]

@ruisearch42 @comaniac @youkaichao Can you please take a final look by any chance?

[ruisearch42]

Otherwise LGTM. Thanks for the PR.
Would be nice if the MP backend owners could review too.

[ruisearch42]

We only support "ray" and "mp", so should have the else branch and assert it's "mp"

[ruisearch42]

nit: use a single word pipeline_parallel_size in the message

[ruisearch42]

Please fix the other minor issue as well

[WoosukKwon]

@bigPYJ1151 I've just started the CI test. Will merge once it becomes green.

[bigPYJ1151]

@WoosukKwon All required became green, please help to merge, thanks :)