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.

翻译：普遍存在的"内存墙"瓶颈在现代大规模专家混合（MoE）架构中被显著放大。MoE固有的架构稀疏性不仅导致稀疏算术计算，还引入了巨大的激活内存开销——这源于庞大的令牌路由缓冲区以及对中间张量进行实例化和缓冲的需求。这种内存压力限制了GPU上可容纳的最大批处理大小和序列长度，同时导致过多的数据移动，从而阻碍了性能表现和高效的模型扩展。本文提出MoEBlaze，一种内存高效的MoE训练框架，通过协同设计的系统方法解决上述问题：（i）采用端到端令牌调度与MoE训练方法，配合优化的数据结构以消除中间缓冲区和激活实例化；（ii）通过协同设计的内核与智能激活检查点技术，在降低内存占用的同时实现更优性能。实验表明，相较于现有MoE框架，MoEBlaze可实现超过4倍的加速比和超过50%的内存节省。