MoE Overlap: Add 2 CUDA Events to Synchronize Computation & Communication Operations for Forward and Backward Respectively - #2630
Closed
yanminjia wants to merge 1 commit into
Closed
Conversation
… forward and backward pass respectively in case of MoE overlap
sequenceDiagram
participant Host as Host (CPU)
participant S1 as Stream 1
participant S2 as Stream 2
participant E as Event (e)
Note over Host, S1: 1. cudaStreamWaitEvent(s1, e) <br/>(wait for the former event)
Host->>S1: enqueue cudaStreamWaitEvent
Note over Host, S1: 2. launch(post_attn_fwd)
Host->>S1: Async Launch: post_attn_fwd
activate S1
Note over Host, S1: 3. cudaEventRecord(e, s1)
Host->>S1: Record event e
S1-->>E: (Event e is triggered after kernel completion)
Note over Host, S1: 4. cudaStreamWaitEvent(s1, e)<br/>(Same-stream wait, no actual blocking)
Host->>S1: enqueue cudaStreamWaitEvent
Note over Host, S1: 5. launch(mlp_bwd)
Host->>S1: Async Launch: mlp_bwd
deactivate S1
activate S1
Note over Host, S1: 6. cudaEventRecord(e, s1)<br/>(The trigger point of e is updated to here)
Host->>S1: Record Event e（overwriting the previous event）
S1-->>E: (Event e is triggered after kernel completion)
deactivate S1
Note over Host, S2: 7. cudaStreamWaitEvent(s2, e)<br/>(Critical synchronization point)
Host->>S2: enqueue cudaStreamWaitEvent
E-->>S2: Block S2 until mlp_bwd completes
Note over Host, S2: 8. launch(moe_dispatch_fwd)
Host->>S2: Sync Launch: moe_dispatch_fwd (alltoall)
activate S2
Note over Host, S2: 9. cudaEventRecord(e, s2)
Host->>S2: Record event e
deactivate S2
gantt
dateFormat X
axisFormat %s
title GPU Kernel Execution Timeline
section Stream 1
post_attn_fwd :a1, 0, 10s
Event Record (e) :milestone, after a1
mlp_bwd :a2, after a1, 10s
Event Record (e) :crit, milestone, after a2
section Stream 2
Wait for Event (e) :active, after a2, 0s
moe_dispatch_fwd :b1, after a2, 10s
Event Record (e) :milestone, after b1
gantt
dateFormat X
axisFormat %s
title Expected result: moe_dispatch_fwd depends only on post_attn_fwd
section Stream 1
post_attn_fwd (k1) :done, a1, 0, 10s
Event Record (e) :crit, milestone, after a1
mlp_bwd (k2) :active, a2, after a1, 10s
section Stream 2
Wait for Event (e) :milestone, after a1, 0s
moe_dispatch_fwd (k3) :active, b1, after a1, 10s
%% This row is only for visually emphasizing parallelism, not actual timing
section Parallel State
Overlap of S1 and S2 :crit, after a1, 10s
|
Contributor
|
@wujiahao15 Current design for ep overlap assigns 1 CUDA event for each microbatch and the overlapping always happen between 2 different microbatches, which is 1 event for fwd microbatch and 1 event for bwd microbatch. |
|
@Wohox Thanks for pointing this out. You are correct and I misunderstood the
|
Contributor
|
@wujiahao15 Can you share the nsys timeline file, it would be easier for me to get information~ |
Author
|
@Wohox @wujiahao15 Hello guys, many thanks for your kindly response. We are trying to maximize perforamance of MoE training & inference. Basically, if MoE computation & communication overlap is secured on the megatron side, we will refine NCCL AlltoAll collective communication based on mechanism such as sm-free & symmetric memory to save SM resource on the one side and reduce communication latency on the other side. Anyway, I suggest we could have an on-line or even off-line (if possible) meeting to discuss some kind of technology details. Many thanks. |
Contributor
2 issues I found
|
Author
|
Hello @Wohox, I tried Mixtral-8x7B with below MoE overlap settings:
Unfortunately, print error message as follows. Please refer to the configuration as attached. Would be highly appreciated if any clue. Many thanks. By the way, with the attached configuration, TFLOPs can go up to 490+ even if no MoE overlap enabled. |
This was referenced Dec 22, 2025
Contributor
|
@yanminjia It seems this error is because of local expert number being 1 (EP8 & num_experts=8), you should be able to resolve the issue by either of these WARs:
The official fix should be soon, thanks~ |
Contributor
|
@yanminjia update: fix PR - #3163 |
Member
|
Fixed by #3164. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem Description
Please refer to issue #2180. MoE communication & computation cannot overlap completely.
Root Cause Analysis
Basically, forward computation of a micro-batch overlaps with backward communication of a different micro-batch and vice versa during interleaving phase. Therefore, the forward or backward process with respect to a transformer layer is split into a couple of communication and computation modules which are managed by TransformerLayerSchedulePlan (megatron/core/model/common/mode_trunk_schedule_plan.py). Additionally, based on TransformerLayerSchedulePlan, the schedule plan of a model trunk is generated by TransformerModelChunkSchedulePlan (megatron/core/model/common/mode_trunk_schedule_plan.py).
Roughly, in the interleaving phase, a forward process of a micro-batch over a transformer layer moves forward side by side with a backward process of a different micro-batch over a different transformer layer within a model trunk. For example, backward combine (communication) is scheduled to go in parallel with forward attention (computation) on different CUDA streams, forward dispatch (communication) is conducted wth backward mlp.
To synchronize the operations on communication stream and computation stream, one CUDA event is used to manage the dependencies of the sub modules in forward pass or backward pass. The computation and communication in forward pass are independent of the computation and communication in backward pass and vice versa. With only one CUDA event for synchronization, a computation operation in forward pass may wait for a communication operation in backward pass mistakenly.
Solution
Add 2 CUDA events for communication & computation synchronization with regard to forward pass and backward pass respectively. A CUDA event is used for forward synchronization and a different CUDA is used for backward synchronization.
Test
As shown by below screen shot, forward attention overlaps with moe_combine backward.