Summary (Overview)

• TAOT (Topology-Aware Optimal Transport) is a novel method for dynamic expert-replica placement in Mixture-of-Experts (MoE) training that jointly optimizes load balance and communication cost across multi-node GPU clusters.
• The method models replica placement as an entropy-regularized optimal transport problem solved via Sinkhorn-Knopp iterations, producing a soft topology preference that guides integer replica matching.
• TAOT achieves a 42.82% end-to-end training speedup over Megatron-LM (155.4 ms → 108.8 ms per iteration) on Qwen3-30B-A3B with 4×8 A800 GPUs.
• The system attains the lowest weighted expert-communication cost across all test configurations, with up to 74% reduction compared to existing state-of-the-art methods.
• TAOT introduces a communication-overlapping execution design that hides guest-weight transfer overhead behind home-expert computation.

---

Introduction and Theoretical Foundation

Background

Mixture-of-Experts (MoE) has become a foundational architecture for scaling large language models, replacing dense feed-forward layers with multiple experts where each token activates only a few. This architecture has been adopted by major models including GLaM, Mixtral, DeepSeek-V3, Qwen3, Kimi-K2, and Llama 4.

The Load Imbalance Problem

In expert-parallel (EP) training, different experts are placed on different GPU ranks, with All-to-All communication dispatching tokens to target ranks. However, dynamic routing causes highly skewed token counts across experts, creating stragglers that slow down synchronous training:

$$\rho = \frac{\max_r \ell_r - \bar{\ell}}{\bar{\ell}} \tag{1}$$

where $\ell_r$ is the computation load on rank $r$ and $\bar{\ell}$ is the global mean load.

Existing Approaches and Their Limitations

• Algorithm-level methods (auxiliary losses, capacity factors, expert-choice routing) improve statistical balance but cannot handle micro-batch-level instantaneous imbalance.
• System-level methods (hot-expert replication, token scheduling, weight migration) keep routing fixed but often ignore where replicas should be placed.

Key Insight

Existing replica placement strategies treat spare capacity as homogeneous, ignoring that intra-node NVLink bandwidth differs markedly from inter-node InfiniBand/RDMA. Two schemes with identical balancing quality can incur completely different communication costs.

This motivates the central thesis: replica planning must trade off peak-shaving capability against the cost of moving expert weights across the communication topology.

---

Methodology

Problem Formulation

The replica-placement objective jointly minimizes residual imbalance and weighted communication cost:

$$\min_{\mathbf{z} \in \{0,1\}} \underbrace{\rho(\mathbf{z})}_{\text{residual imbalance}} + \mu \cdot \underbrace{\sum_{e,r} z_{er} \cdot W_{r_e, r}}_{\text{weighted replica comm. cost}}$$ $$\text{s.t.} \quad \sum_e z_{er} \leq K, \quad \forall r \in \mathcal{C} \tag{2}$$

where $\mu > 0$ trades off balance against communication cost, and $W_{r_e,r}$ is 1 for intra-node and $\lambda$ for inter-node transfers.

Three-Stage Planning Algorithm

Phase 1: Sinkhorn-Knopp Topology-Aware Flow Planning

Sets up a rank-level balanced optimal transport problem:

$$T^* = \arg\min_{T \geq 0} \langle T, W \rangle$$ $$\text{s.t.} \quad T\mathbf{1} = \mathbf{s}, \quad T^\top \mathbf{1} = \mathbf{d} \tag{3}$$

Adding a negative-entropy regularizer yields a Gibbs-kernel structure $T^*_\varepsilon = \text{diag}(\mathbf{u}) M \text{diag}(\mathbf{v})$ where $M_{ij} = \exp(-W_{ij}/\varepsilon)$, solved via alternating Sinkhorn-Knopp iterations. When $\varepsilon = \lambda$, the intra-node kernel value ratio is $\exp((\lambda-1)/\lambda) > 1$, forming a soft topology preference.

Phase 2: Column-First Iterative Matching

Converts the continuous flow into binary replica decisions using a scoring function:

$$\text{score}_{er} = \underbrace{\min(\text{spill}_e, \text{spare}_r)}_{\text{main: balance gain}} + \alpha \underbrace{B_{er}}_{\text{second: topology pref.}} + 0.1\alpha \underbrace{(T_{er})_{\text{norm}}}_{\text{third: OT flow hint}} \tag{4}$$

where $\alpha = \bar{\ell} R/E$ is a scale-alignment factor ensuring balance takes priority.

Phase 3: Lagrange Auction Token Assignment

Uses a Lagrange multiplier (price) $p_{r'}$ for rank-capacity constraints. Each spare slot bids for source ranks by net gain $B_{r's} - p_{r'}$, with prices increasing monotonically after winning to spread load evenly.

Communication Overlapping

Guest-weight transfer is overlapped with home-expert computation, hiding communication overhead during both forward and backward passes.

---

Empirical Validation / Results

Experimental Setup

• Hardware: 4×8 A800 GPUs
• Model: Qwen3-30B-A3B MoE
• Dataset: Pile-test
• Baselines: Megatron-LM, ECHO, LPLB, LLEP

End-to-End Performance

TAOT reduces iteration time from 155.4 ms to 108.8 ms (42.82% speedup). Loss consistency is maintained with mean absolute relative error of 0.297‰, confirming no numerical precision loss.

Balance Quality and Communication Cost

Table 1: Comprehensive comparison across EP scales and initial balance conditions. Best values in bold.

Key findings:

• TAOT achieves the lowest weighted communication cost in all 10 configurations
• At EP=32, TAOT is best or tied for best in balance quality
• Communication cost reduction: up to 53% lower than LPLB at EP=16, up to 74% lower at EP=32

Scalability Results

• EP scalability: Speedup grows from EP4 to EP16, reaching 1.79×
• Imbalance scalability: Speedup grows from 1.21× (30% imbalance) to 1.75× (90% imbalance)
• Guest-slot sensitivity: Optimal at 2 slots per rank (balance of performance and memory)
• Planning overhead: Below 1% of forward time

Ablation Study

Table 2: Ablation of TAOT planning algorithm.

The ablation confirms both components are essential: Phase 2 reduces inter-node transfers by 45%, while Phase 1 further improves balance quality by 26%.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Joint optimization framework: First to explicitly model both peak-shaving gain and topology-dependent communication cost in a unified objective function
2. Soft topology preference: Unlike LPLB's hard graph constraints, TAOT uses a continuous cost matrix that allows flexible cross-node placement when necessary
3. GPU-friendly algorithm: Sinkhorn-Knopp iterations are naturally parallelizable, enabling micro-batch-level planning

Practical Implications

• Scalability: TAOT's advantage grows with EP scale, making it increasingly valuable for larger clusters
• Cost efficiency: Significant reduction in cross-node communication translates directly to faster training and lower infrastructure costs
• Generality: Works with standard EP settings without requiring specialized hardware (unlike FEPLB) or fixed communication paradigms (unlike LAER-MoE)

---

Conclusion

TAOT addresses the critical but underrated question of where to place replicated experts in MoE training. By unifying peak-shaving gain and communication cost into an entropy-regularized optimal transport framework, it achieves:

• 42.82% end-to-end speedup over Megatron-LM
• Best-in-class communication efficiency with up to 74% reduction
• Competitive or superior balance quality that improves with scale

The three-phase planning algorithm (Sinkhorn flow hints → column-first matching → Lagrange auction) provides a practical, GPU-friendly solution with <1% overhead. The communication-overlapping execution design further hides guest-weight transfer costs.

Future directions could include extending to heterogeneous clusters, dynamic topology adaptation, and integration with other load-balancing techniques at the routing level.