# TAOT: Topology-Aware Optimal Transport for Dynamic Expert Replica Placement in MoE Training

> TAOT jointly optimizes expert-replica placement and communication cost via optimal transport, achieving 42.82% end-to-end training speedup over Megatron-LM with up to 74% lower communication cost.

- **Source:** [arXiv](https://arxiv.org/abs/2608.03676)
- **Published:** 2026-09-05
- **Permalink:** https://picx.dev/p/SUyVjr
- **Whiteboard:** https://picx.dev/p/SUyVjr/image

## Summary

## 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

| Initial Imbalance | Method | EP=16 Improve (pp) | EP=16 Weighted Cost | EP=32 Improve (pp) | EP=32 Weighted Cost |
|---|---|---|---|---|---|
| 10% | ECHO | -9.4 | 24 | -9.5 | 72 |
|  | LLEP | -9.4 | 20 | -9.5 | 44 |
|  | LPLB | -8.8 | 32 | -8.6 | 128 |
|  | **TAOT** | **-9.4** | **17** | **-9.5** | **38** |
| 30% | ECHO | -27.7 | 24 | -28.2 | 75 |
|  | LLEP | -28.7 | 22 | -28.4 | 42 |
|  | LPLB | -28.9 | 32 | -27.6 | 128 |
|  | **TAOT** | -28.7 | **18** | **-28.4** | **37** |
| 70% | ECHO | -68.2 | 42 | -66.4 | 74 |
|  | LLEP | -68.9 | 29 | -66.7 | 40 |
|  | LPLB | -69.5 | 32 | -63.0 | 128 |
|  | **TAOT** | -68.9 | **27** | **-66.7** | **33** |

*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

| Phase 1 flow hint | Phase 2 comm. cost | Final imbalance | Intra transfers | Inter transfers | Weighted cost |
|---|---|---|---|---|---|
| × | × | 1.95% ±0.18 | 4.67 ±1.70 | 18.33 ±1.25 | 59.67 ±2.36 |
| × | √ | 2.00% ±0.82 | 14.67 ±0.94 | 10.00 ±1.41 | 44.67 ±3.86 |
| √ | √ | **1.48% ±0.26** | **15.00 ±0.82** | **9.67 ±1.70** | **44.00 ±4.55** |

*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.

---

_Markdown view of https://picx.dev/p/SUyVjr, served by PicX — AI-generated visual whiteboard summaries of research papers._