NVIDIA / NVIDIA/Megatron-LM

[BUG] Shared expert overlap delays independent expert wgrad until input-gradient merge

Open
#7,389 0 comments 0 reactions 0 assignees View on GitHub
community-request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.