[QUESTION] Possible communication reduction in MoE training

[laze44]

Hi Megatron team:
During MoE training, there might be opportunities to combine communication operators when simultaneously using sequence and expert parallelism, though I'm uncertain if this hypothesis is accurate.
In the original sequence parallelism + expert parallelism process, assuming sp=ep=4, the activations undergo an all-gather phase after dropout, followed by gating operations for all sequences, which are then selectively routed to different GPUs' experts. However, if I move the gating function before the all-gather phase, each GPU would perform the gating operations on its sequences, followed by an all-to-all communication based on the gating results. The rationale behind this approach is that, unlike FFN tensor parallelism, during the MoE forward process, each GPU only needs to handle a subset of sequences. Theoretically, this could eliminate the need for the all-gather phase. （As shown in Figure below)
I want to understand whether this approach is correct, as the MoE training process typically incorporates some load-balancing-related loss functions. Does altering the order affect the backward process? Moreover, if this approach is correct, does Megatron-LM support this communication concept?

[Advait251206]

This is a very good question, and it gets at a subtle but important point about how Megatron-style MoE actually uses sequence parallelism vs expert parallelism.

Short answer:

❌ You cannot safely move gating before the sequence-parallel all-gather in standard MoE implementations, and Megatron-LM does not support that alternative execution path.

Now the reasoning.

---

1. Why the all-gather exists in MoE + sequence parallelism

In Megatron-LM MoE with sequence parallelism (SP), the key idea is:

• Each GPU holds only a partition of hidden dimension (H) per token
• Not the full hidden representation

So before routing decisions:

the gating network needs a full representation per token

That is why Megatron does:

Step A (SP context rebuild)

all-gather → reconstruct full hidden state [BS, H]

Step B

router / gating uses full H

Step C

top-k expert assignment
all-to-all dispatch to experts

---

2. Why your proposed reordering looks attractive

Your idea:

“run gating before all-gather so each GPU routes its local shard”

This is appealing because:

• avoids all-gather (good for communication)
• seems like each GPU could independently route tokens
• reduces memory traffic

But there is a fundamental issue.

---

3. Why it is NOT equivalent mathematically

The gating function in MoE is:

gate(x) = softmax(Wx)

But in sequence parallelism:

• each GPU holds only x_shard
• not full x

So if you do gating locally:

gate(x_shard)

you are computing:

a different function

because:

• Wx is not decomposable across sequence-parallel shards
• you cannot reconstruct logits from partial hidden states without communication

So:

❌ local gating ≠ global gating

This breaks:

• routing consistency
• load balancing correctness
• expert selection distribution

---

4. What breaks in backward pass?

If you reorder as proposed:

You break gradient consistency in:

• routing probabilities
• expert assignment decisions
• capacity balancing loss

Why?

Because top-k routing is:

• discrete / non-smooth
• depends on full logits

Changing inputs changes:

• which expert processes which token
• therefore completely different computation graph

So backward pass is no longer equivalent.

---

5. Why load-balancing loss does NOT fix this

You mentioned:

“MoE uses load balancing loss anyway”

Yes, but:

• load balancing loss assumes correct global routing decisions
• it does NOT compensate for incorrect routing inputs

So:

• it stabilizes expert usage
• but does not fix wrong gating inputs

---

6. Does Megatron-LM support your proposed optimization?

❌ No

Current Megatron MoE design assumes:

• SP → all-gather first
• then routing
• then EP all-to-all dispatch

There is no supported mode where:

gating is done on sharded hidden states before reconstruction

Because it would require:

• redesign of router
• redesign of loss
• different parallelism semantics

---

7. Why this optimization is fundamentally hard

To remove all-gather, you would need:

Option A: replicated gating weights + full tokens everywhere

→ defeats SP memory savings

Option B: distributed router (compute logits jointly across GPUs)

→ requires all-reduce of logits per token (still communication-heavy)

Option C: approximate routing

→ hurts model quality significantly

So there is no free win here.

---

8. When communication reduction is possible in MoE

Megatron and related systems do optimize communication, but via:

• token dropping / top-k sparsity tuning
• expert capacity tuning
• fused all-to-all kernels
• overlapping communication + computation
• grouped dispatch optimizations

Not by changing gating order.

---

9. Bottom line

• ❌ Your proposed reordering is not mathematically equivalent
• ❌ It breaks routing correctness and training dynamics
• ❌ Megatron-LM does not support it
• ✅ SP all-gather before gating is required for correctness
• ⚠️ Communication can be optimized, but not removed this way

---

If you're interested, I can also explain:

• how DeepSpeed-MoE differs here
• why Switch Transformer avoids some of this overhead
• or what modern systems (e.g. DeepSpeed-Ulysses / Megablocks) do to reduce all-to-all cost

Those are closer to the optimization direction you're thinking about.