Multi-Head LatentMoE and Head Parallel: Communication-Efficient and Deterministic MoE Parallelism
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2026
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866915774478680064 |
|---|---|
| author | Cui, Chenwei Jackson, Rockwell Herrera, Benjamin Joseph Tárano, Ana María Kerner, Hannah |
| author_facet | Cui, Chenwei Jackson, Rockwell Herrera, Benjamin Joseph Tárano, Ana María Kerner, Hannah |
| contents | 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 achieving $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.61\times$ faster while having identical performance. With doubled granularity, it achieves higher overall performance while still being $1.11\times$ faster. Our method makes multi-billion-parameter foundation model research more accessible. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_04870 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Multi-Head LatentMoE and Head Parallel: Communication-Efficient and Deterministic MoE Parallelism Cui, Chenwei Jackson, Rockwell Herrera, Benjamin Joseph Tárano, Ana María Kerner, Hannah Machine Learning 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 achieving $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.61\times$ faster while having identical performance. With doubled granularity, it achieves higher overall performance while still being $1.11\times$ faster. Our method makes multi-billion-parameter foundation model research more accessible. |
| title | Multi-Head LatentMoE and Head Parallel: Communication-Efficient and Deterministic MoE Parallelism |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.04870 |