[Task][Megatron] 4.2: Expert parallel / MoE communication path

[lvyufeng]

Parent Issue

Part of #206 — Megatron Roadmap, Phase 4 (Advanced Features / Stretch): Task 4.2

Description

Implement expert parallelism communication patterns for Mixture-of-Experts models.

Requirements

• All-to-all dispatch for token routing to experts
• All-to-all combine for gathering expert outputs
• Capacity factor and load balancing support
• Compatible with Megatron's MoE layer implementation
• Works with EP process groups

Blocked By

• [Task][Megatron] 3.3: Extend DeviceMesh beyond 2D #260 (DeviceMesh beyond 2D)
• [Task][Megatron] 1.2: Differentiable collective autograd.Functions #253 (Differentiable collectives)

Blocks

• MoE model training on candle