Skip to content

[0.9.1][Feature]Moe alltoallv communication optimization for unquantized RL training sence & alltoallv support dpo #1547

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 41 commits into
base: v0.9.1-dev
Choose a base branch
from
Open

[0.9.1][Feature]Moe alltoallv communication optimization for unquantized RL training sence & alltoallv support dpo #1547

wants to merge 41 commits into from
+2,991 −32

Conversation

weijinqian0
Copy link
Contributor

@weijinqian0 weijinqian0 commented Jul 1, 2025 •
edited
Loading

[Feature]Moe alltoallv communication optimization for unquantized RL training sence & alltoallv support dpo

Introduction

This PR introduces two key optimizations for MoE model performance:

  1. Efficient Token Dispatcher:

    • Implements an optimized alltoallv_seq token dispatcher (adopted from NVIDIA Megatron and Ascend MindSpeed)
    • Significantly more efficient than current alltoall implementation when using token_permute/unpermute fusion
    • Enable with: VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ=1

  2. DBO Support for alltoallv_seq:

    • Builds upon the alltoallv_seq dispatcher to support DBO (Dual Batch Overlap)
    • Enables overlapping of alltoallv communication during the prefilling stage
    • Enable with both:
      • VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ=1
      • VLLM_ASCEND_ENABLE_DBO=1

Performance Improvements

Testing on Qwen3-30B-A3B shows nearly 2x throughput improvement compared to the original alltoall implementation.

weijinqian_v1 added 12 commits July 1, 2025 09:51
[Feature]Moe alltoallv communication optimization for unquantized RL …
7ff288e 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
6d7b5b4 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
6a8e1a9 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
4805c5a 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
d68ce07 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
0aff693 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
f6ab19e 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
a94c094 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
91570d8 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
e7c0d2d 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
47439e8 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
cf3f1c8 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
weijinqian_v1 added 3 commits July 1, 2025 14:03
[Feature]Moe alltoallv communication optimization for unquantized RL …
a4126f3 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
807aaf0 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
[Feature]Moe alltoallv communication optimization for unquantized RL …
6f6efc1 
…training sence & alltoallv support dpo

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@github-actions GitHub Actions
Copy link

github-actions bot commented Jul 3, 2025

This pull request has conflicts, please resolve those before we can evaluate the pull request.

handle conflict
305a0eb 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@github-actions GitHub Actions
Copy link

github-actions bot commented Jul 8, 2025

This pull request has conflicts, please resolve those before we can evaluate the pull request.

weijinqian_v1 and others added 5 commits July 9, 2025 16:25
add st:qwen3
5411ed6 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
add st for moe token dispatcher
3f88769 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
fix bug
854c149 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
add st for moe token dispatcher
d0bd006 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
add moe_block: AscendSparseMoeBlock
49e9771 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
harygo22 added 6 commits July 9, 2025 16:28
@harygo22
revert
b02ad40 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
fix bug
66807e0 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
fix a bug
d24758e 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
fix a bug
d76c4fb 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
ut test
f883902 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@harygo22
liscens & fix dsk dbo.
d5656f4 
* ut test

* liscense & fix dsk dbo.

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
weijinqian_v1 added 2 commits July 9, 2025 16:41
handle conflict
df52070 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
handle code clean
adf3f74 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@github-actions GitHub Actions
Copy link

github-actions bot commented Jul 9, 2025

This pull request has conflicts, please resolve those before we can evaluate the pull request.

handle code clean
5956ef0 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
weijinqian_v1 added 2 commits July 9, 2025 23:48
handle code clean
af85566 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
handle code clean
d4ad734 
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
@github-actions GitHub Actions
Copy link

github-actions bot commented Jul 9, 2025

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@harygo22
fix comment
3b7269a 
Signed-off-by: duyangkai <duyangkai@huawei.com>
@wangxiyuan wangxiyuan changed the title [Feature]Moe alltoallv communication optimization for unquantized RL training sence & alltoallv support dpo [0.9.1][Feature]Moe alltoallv communication optimization for unquantized RL training sence & alltoallv support dpo Jul 10, 2025
ganyi1996ppo
ganyi1996ppo reviewed Jul 10, 2025
View reviewed changes
vllm_ascend/ascend_forward_context.py
elif envs_ascend.VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ:
# MC2 Dispatch/Combine performs better than alltoall_seq in decoding stage.
return FusedMoEState.All2AllSeq if (
ep_size < 16 or with_prefill) else FusedMoEState.MC2
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why there is this restriction ep_size <16 ?

Copy link

@harygo22 harygo22 Jul 10, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MC2 Dispatch/Combine is still faster than alltoall_seq in decoding stage. so when ep_size >= 16, use MC2 for better performance.

@github-actions GitHub Actions
Copy link

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@harygo22
fix init
a8b3e15 
Signed-off-by: duyangkai <duyangkai@huawei.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ganyi1996ppo ganyi1996ppo ganyi1996ppo left review comments

@harygo22 harygo22 harygo22 left review comments

Assignees
No one assigned
Labels
module:core module:ops module:tests
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
7 participants
@weijinqian0 @ganyi1996ppo @harygo22 @whx-sjtu @JC-ut0 @weiguihua2 @songshanhu07