HierMoE System

• HierMoE is an advanced system that optimizes distributed training of MoE models using topology-aware communication and hierarchical token deduplication.
• It employs a multi-dimensional AlltoAll algorithm and expert swap to reduce redundant transfers and balance loads across GPU clusters.
• Evaluations on 32-GPU clusters with Megatron-LM show significant speedups and scalability improvements over traditional MoE training systems.

HierMoE is an advanced system designed to accelerate the distributed training of large-scale Mixture-of-Experts (MoE) models by employing topology-aware communication optimization techniques. The architecture specifically addresses the communication and load imbalance bottlenecks that occur when tokens routed to experts must traverse complex GPU cluster interconnects (such as NVLink, PCIe, or QPI). Implemented atop Megatron-LM and evaluated on 32-GPU clusters with models such as DeepSeek-V3 and Qwen3-30B-A3B, HierMoE demonstrates substantial improvements in training speed and communication efficiency over existing MoE training systems, including Tutel-2DH, SmartMoE, and Megatron-LM (Lin et al., 13 Aug 2025).

1. Topology-Aware Hierarchical Communication Architecture

HierMoE restructures the standard MoE communication paradigm by mapping the model's expert and token routing onto the physical hierarchy of the GPU interconnect topology. GPUs are organized into multiple nested communication groups reflecting real-world network layers (such as node, intra-node, and link-specific subgroups).

The system employs "HierD-AlltoAll", a multi-dimensional AlltoAll communication algorithm that eliminates redundant token transfers at each level of the hierarchy. Rather than naively sending each token copy to every routed expert (across arbitrary GPUs), HierMoE deduplicates tokens within tightly coupled groups and only transfers unique tokens across higher-bandwidth-limited links. This substantially reduces the aggregate communication volume that would otherwise saturate interconnects.

2. Hierarchical Token Deduplication Technique

HierMoE's hierarchical token deduplication identifies all situations in which multiple experts co-located in the same communication group require the same token. Before the system conducts the expensive AlltoAll exchange across a network boundary, it collates these duplicates and sends only one copy, leveraging a bitwise OR operation across the expert routing masks to determine token destinations.

The theoretical communication time for ordinary AlltoAll is given by:

$t_1 = \alpha_{a2a} + n_{a2a} \cdot \beta_{a2a}$

where $n_{a2a} = G \cdot \max(p) \cdot M \cdot v$ ($G$: GPU count, $\max(p)$: max deduped token count per GPU, $M$: embedding dimension, $v$: bytes per dim). In the $d$-dimensional hierarchical setting, the communication cost is:

$$t_d = \sum_{i=1}^{d-1} [ n_{a2a}^{Inter(i)} \cdot \beta_{a2a}^{Inter(i)} + \alpha_{a2a}^{Inter(i)} ] + n_{a2a}^{Intra(d-1)} \cdot \beta_{a2a}^{Intra(d-1)} + \alpha_{a2a}^{Intra(d-1)}$$

Here, $n_{a2a}^{Inter(i)}$ is computed post-deduplication for the $i$th hierarchy boundary, drastically minimizing inter-node traffic and congestion.

3. Hierarchical Expert Swap for Load Balancing

Token routing functions can imbalance the token distribution among experts, causing communication hotspots. HierMoE introduces hierarchical expert swap, which exchanges expert positions within the layout to optimize the aggregate communication pattern.

The swap candidates are evaluated via a time–cost model:

$$\mathcal{Q}_d[r, c] = \sum_{i=1}^{d-1} (\mathcal{N}_{a2a}^{Inter(i)}[r, c] \cdot \beta_{a2a}^{Inter(i)} + \alpha_{a2a}^{Inter(i)}) + \mathcal{N}_{a2a}^{Intra(d-1)}[r, c] \cdot \beta_{a2a}^{Intra(d-1)} + \alpha_{a2a}^{Intra(d-1)}$$

The swap pair $(r^*, c^*) = \arg\min \mathcal{Q}_d[r, c]$ minimizes communication cost post-swap. Where abrupt cost changes must be smoothed, the smooth-max function:

$$\text{smooth-max}(x, \gamma) = \max(x) \cdot \left( \sum_i \left( \frac{x[i]}{\max(x)} \right)^\gamma \right)^{1/\gamma}$$

is used to regularize the optimization.

4. Theoretical Performance Models

HierMoE’s closed-form models for communication prediction are parameterized using real hardware benchmarks. For token deduplication, the communication volume per inter-group at level $i$ is expressed as:

$$n_{a2a}^{Inter(i)} = \frac{U[i]}{U[i-1]} \cdot \max( p_{a2a}^{Inter(i)} ) \cdot M \cdot v$$

Optimal hierarchical dimension $d^*$ is selected as

$$d^* = \begin{cases} 1, & t_1 < t_d,\ \forall\ d > 1\ \arg\min_{1 < d \leq D} t_d, & \text{otherwise} \end{cases}$$

Expert swap effects are dynamically modeled per iteration using actual routing patterns, allowing precise adaptation to changing token–expert assignments and cluster parameters.

5. Implementation and Experimental Evaluation

The HierMoE system is integrated into Megatron-LM, modifying the dispatch/combine stages of MoE layers to include hierarchical token deduplication and swap logic. Hardware parameters ($\alpha, \beta$) are microbenchmarked per hierarchy layer (e.g., node, link, device).

Empirical tests on 32-GPU clusters (organized as 4 nodes × 8 GPUs per node) with models DeepSeek‑V3 (6 layers, halved hidden size due to memory constraints) and Qwen3-30B-A3B (32 layers) demonstrate:

• Communication speedup of 1.55× to 3.32× over Tutel‑2DH and SmartMoE baselines
• End-to-end training speedup of 1.18× to 1.27× compared to Megatron-LM (Lin et al., 13 Aug 2025)
• Stability and adaptivity to model and hardware changes, with optimal communication hierarchy and swap configuration selected via the performance model

6. Impact on Scalability and Distributed Training

HierMoE significantly reduces redundant data transfer by deduplicating tokens within communication subgroups, minimizing inter-node bandwidth demand. The expert swap algorithm ensures that the token-expert workload is well balanced across the device and network hierarchy, avoiding communication bottlenecks typical of naive MoE routing.

Because both deduplication and swap are parameterized in terms of the topology of the hardware, the techniques generalize across different physical cluster organizations and model configurations. The result is improved scalability, allowing efficient distributed training of much larger MoE models across GPU clusters.

A plausible implication is that as clusters scale in GPU count and interconnect complexity, HierMoE’s hierarchical model allows near-optimal communication and load balance, avoiding superlinear scaling of bandwidth requirements and iteration time.

7. Conclusion and Future Significance

HierMoE realigns MoE training with awareness of the underlying hardware topology. Its hierarchical token deduplication algorithm tangibly reduces communication volume, while expert swap prevents load imbalance that can stall distributed systems. Both are operationalized via theoretical models and evaluated on production-grade clusters and models, confirming speedup and efficiency over established baselines.

The demonstrated improvement in communication (up to 3.32×) and total training speed (up to 1.27×) (Lin et al., 13 Aug 2025) is notable given the computational cost of large-scale distributed LLM training. HierMoE’s techniques are generalizable and parameterized for extensibility to larger clusters and diverse hardware. This suggests that HierMoE provides a foundation for continued scaling of sparsely activated expert-based LLMs with minimal communication and load bottlenecks in future multi-device systems.