[BUG] Shared expert overlap delays independent expert wgrad until input-gradient merge
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**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`:
1. [MoELayer.dispatch](https://github.com/NVIDIA/Megatron-LM/blob/bb5dfd08f09ce06c5925af453fef06b3129f199d/megatron/core/transformer/moe/moe_layer.py#L588-L597) inserts `_RegisterDelayedWgradForExperts` before calling the token dispatcher.
2. [Flex token_dispatch](https://github.com/NVIDIA/Megatron-LM/blob/bb5dfd08f09ce06c5925af453fef06b3129f199d/megatron/core/transformer/moe/token_dispatcher.py#L2099-L2106) feeds that same tensor into dispatch and the shared-expert branch.
3. The registration node's backward therefore waits for the gradients from both branches before calling `backward_dw`.
```mermaid
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:
```python
moe_token_dispatcher_type = "flex"
moe_shared_expert_overlap = True
overlap_dispatch_backward_with_experts_wgrad = True
```
Run 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
Contributor guide
Research direction
Start with MoELayer.dispatch in megatron/core/transformer/moe/moe_layer.py and Flex token_dispatch in megatron/core/transformer/moe/token_dispatcher.py at the referenced revision. Reproduce a forward/backward pass with the two overlap settings enabled, then inspect the autograd and gradient-hook ordering around _RegisterDelayedWgradForExperts and MoELayer.backward_dw. Done means routed-expert wgrad submission no longer waits for the shared/routed input-gradient merge while its existing completion wait and gradient-processing hooks remain intact.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100