[BUG] Shared expert overlap delays independent expert wgrad until input-gradient merge
Describe the bug
When moe_shared_expert_overlap=True and overlap_dispatch_backward_with_experts_wgrad=True are enabled together, the Flex dispatcher's autograd graph makes routed-expert weight-gradient (wgrad) submission depend on the shared/routed input-gradient merge. This introduces a dependency on shared-expert backward before wgrad can be submitted, undermining the intended dispatch-backward/wgrad overlap.
The dependency is visible in dev commit bb5dfd08f09ce06c5925af453fef06b3129f199d:
- MoELayer.dispatch inserts
_RegisterDelayedWgradForExpertsbefore calling the token dispatcher. - Flex token_dispatch feeds that same tensor into dispatch and the shared-expert branch.
- The registration node's backward therefore waits for the gradients from both branches before calling
backward_dw.
flowchart LR
D["Dispatch backward"] --> M["Shared + routed input-gradient merge"]
S["Shared FC1 backward"] --> M
M --> W["Submit routed expert wgrad"]This dependency is in the common Flex graph structure; it is not inherently specific to DeepEP V2.
Steps/Code to reproduce bug
On the dev revision above, use a supported Flex backend and a MoE layer with shared experts, enabling:
moe_token_dispatcher_type = "flex"
moe_shared_expert_overlap = True
overlap_dispatch_backward_with_experts_wgrad = TrueRun a forward/backward pass with gradients enabled. To inspect the scheduling dependency:
- Register an autograd node post-hook on the dispatched hidden states.
- Register a tensor gradient hook on the input received by
shared_experts.pre_forward_comm; this observes the shared/routed input-gradient merge. - Record entry into
MoELayer.backward_dw.
The graph requires the merge before wgrad submission, even though routed-expert dgrad has its own completion event.
Expected behavior
Submit routed-expert wgrad once expert dgrad is ready and dispatch backward has submitted communication, without waiting for the shared-expert input-gradient merge. Preserve the completion wait and gradient-processing hooks at their existing boundary.
Additional context
This report concerns an avoidable scheduling dependency. Whether removing it improves latency depends on workload size and the available communication-overlap window.
cc @NVIDIA/mcore-oncall
Source: NVIDIA/Megatron-LM