Multi-Head LatentMoE and Head Parallel: Communication-Efficient and Deterministic MoE Parallelism

Large language models have transformed many applications but remain expensive to train. Sparse Mixture of Experts (MoE) addresses this through conditional computation, with Expert Parallel (EP) as the standard distributed training method. However, EP has three limitations: communication cost grows linearly with the number of activated experts $k$, load imbalance affects latency and memory usage, and data-dependent communication requires metadata exchange. We propose Multi-Head LatentMoE and Head Parallel (HP), a new architecture and parallelism that achieve $O(1)$ communication cost regardless of $k$, completely balanced traffic, and deterministic communication, all while remaining compatible with EP. To accelerate Multi-Head LatentMoE, we propose IO-aware routing and expert computation. Compared to MoE with EP, Multi-Head LatentMoE with HP trains up to $1.82\times$ faster while having better performance. With double the granularity, the performance is even better while being $1.08\times$ faster. Our method makes multi-billion-parameter foundation model research more accessible.

Training large language models requires many computers working together, but communication between them creates a major bottleneck. In "Mixture of Experts" models, only a small portion of the network activates for each input, making training more efficient. However, the standard approach sends data between computers after deciding which portions to activate. This creates problems: increasing active portions leads to more communication, uneven workloads cause some computers to wait for others, and unpredictable communication patterns require extra coordination steps. We propose a new architecture called Multi-Head LatentMoE and a training method called Head Parallel. The key insight is to split each input into smaller pieces and distribute them across computers before making activation decisions. This way, communication happens once in a fixed, predictable pattern, regardless of how the model routes its computations internally. Our approach trains up to 1.82 times faster than existing methods while achieving better model quality. By reducing communication overhead to 25 percent of standard methods, our work makes research on billion-parameter language models more accessible to academic researchers with limited computing resources.