MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866914242378072064 |
|---|---|
| author | Zhang, Jiyuan Liu, Yining Yan, Siqi Deng, Lisen Cao, Jennifer Yang, Shuqi Ni, Min Xue, Bi Li, Shen |
| author_facet | Zhang, Jiyuan Liu, Yining Yan, Siqi Deng, Lisen Cao, Jennifer Yang, Shuqi Ni, Min Xue, Bi Li, Shen |
| contents | The pervasive "memory wall" bottleneck is significantly amplified in modern large-scale Mixture-of-Experts (MoE) architectures. MoE's inherent architectural sparsity leads to sparse arithmetic compute and also introduces substantial activation memory overheads -- driven by large token routing buffers and the need to materialize and buffer intermediate tensors. This memory pressure limits the maximum batch size and sequence length that can fit on GPUs, and also results in excessive data movements that hinders performance and efficient model scaling. We present MoEBlaze, a memory-efficient MoE training framework that addresses these issues through a co-designed system approach: (i) an end-to-end token dispatch and MoE training method with optimized data structures to eliminate intermediate buffers and activation materializing, and (ii) co-designed kernels with smart activation checkpoint to mitigate memory footprint while simultaneously achieving better performance. We demonstrate that MoEBlaze can achieve over 4x speedups and over 50% memory savings compared to existing MoE frameworks. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_05296 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs Zhang, Jiyuan Liu, Yining Yan, Siqi Deng, Lisen Cao, Jennifer Yang, Shuqi Ni, Min Xue, Bi Li, Shen Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing The pervasive "memory wall" bottleneck is significantly amplified in modern large-scale Mixture-of-Experts (MoE) architectures. MoE's inherent architectural sparsity leads to sparse arithmetic compute and also introduces substantial activation memory overheads -- driven by large token routing buffers and the need to materialize and buffer intermediate tensors. This memory pressure limits the maximum batch size and sequence length that can fit on GPUs, and also results in excessive data movements that hinders performance and efficient model scaling. We present MoEBlaze, a memory-efficient MoE training framework that addresses these issues through a co-designed system approach: (i) an end-to-end token dispatch and MoE training method with optimized data structures to eliminate intermediate buffers and activation materializing, and (ii) co-designed kernels with smart activation checkpoint to mitigate memory footprint while simultaneously achieving better performance. We demonstrate that MoEBlaze can achieve over 4x speedups and over 50% memory savings compared to existing MoE frameworks. |
| title | MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2601.05296 |