解决MoE训练内存瓶颈，实现4倍加速和50%内存节省
MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs
- 通过优化数据结构消除中间缓冲区，避免激活值重复存储
- 采用智能激活检查点技术，内存占用减少50%以上
- 适合大规模MoE模型训练，尤其在现代GPU上表现优异
现代大规模混合专家（MoE）架构中，内存墙问题尤为突出。由于架构固有的稀疏性，导致稀疏计算与大量激活内存开销——主要来自大型令牌路由缓冲区及中间张量的显式存储需求。这种内存压力限制了GPU上的最大批大小和序列长度，并引发频繁的数据搬运，阻碍性能提升与模型高效扩展。我们提出MoEBlaze，一种协同设计的内存高效训练框架：(i) 端到端的令牌调度与MoE训练方法，采用优化数据结构消除中间缓冲区与激活值显式存储；(ii) 协同设计的核函数结合智能激活检查点，在降低内存占用的同时提升性能。实验表明，MoEBlaze相比现有MoE框架可实现超过4倍的速度提升和超过50%的内存节省。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。