fix bug: dp+tp warmup

[Tsundoku958]

Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily receiving feedbacks. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.

Motivation

During the warmup phase of LMDeploy when using Data Parallelism (DP) + Tensor Parallelism (TP), the build_dp_meta() function is not invoked.

Reproduction:

GPU: H20

Command:

lmdeploy serve proxy --server-name 0.0.0.0  --server-port 23333  --routing-strategy "min_expected_latency" --serving-strategy Hybrid
export LMDEPLOY_DP_MASTER_ADDR=0.0.0.0
export LMDEPLOY_DP_MASTER_PORT=23337
lmdeploy serve api_server ../deepseek-v2-lite/ --dp 2 --tp 2 --proxy_url http://0.0.0.0:23333 --backend pytorch

Traceback:

Modification

Call inputs.build_dp_meta() before call _forward_impl

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[grimoire]

Note that we might change the behaviour of DP+TP in the future.

[Tsundoku958]

Note that we might change the behaviour of DP+TP in the future.

Could you tell me what it's specifically about?

[grimoire]

Current implementation would pad inputs to the same batch size(input_meta). For each layer, the pipeline would be gather -> forward -> scatter, which requires more computation and memory usage.

I want to decouple DP and TP. DP2 TP2 would use 4 GPU, each DP rank would be single engine with TP2. This is good for DP+TP+EP (less padding and less collective OP).

I have not finish planning yet, any advices are welcome.

[Tsundoku958]

Current implementation would pad inputs to the same batch size(input_meta). For each layer, the pipeline would be gather -> forward -> scatter, which requires more computation and memory usage.

I want to decouple DP and TP. DP2 TP2 would use 4 GPU, each DP rank would be single engine with TP2. This is good for DP+TP+EP (less padding and less collective OP).

I have not finish planning yet, any advices are welcome.

For models utilizing GQA or MHA, this is indeed a superior solution. However, for MLA or MQA architectures, employing TP cannot partition the KV cache based on the number of heads. This results in each GPU storing a full replica of the KV cache. In such cases, DP can be applied specifically to the attention component to reduce the GPU memory footprint of the KV cache on individual cards.

Therefore, I believe that retaining the current implementation approach in future versions remains a viable strategy.
Sglang call this dp attention https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models

[grimoire]

LGTM

[lvhan028]

Hi, @Tsundoku958
Thank you very much for the contribution.
Since we've freezed main last Friday for v0.10.1 testing, this PR will be handled in next version.

[grimoire]

#4004 This is the DP-TP refactor with both TP implementations. Feel free to review the PR.