MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs

TL;DR

MoEBlaze introduces a memory-efficient training framework for large-scale MoE models on GPUs, significantly reducing memory usage and increasing training speed by optimizing data structures and kernels.

Contribution

It presents a co-designed system approach that eliminates intermediate buffers and reduces memory overhead in MoE training on modern GPUs.

Findings

Over 4x speedup in training performance

More than 50% reduction in memory usage

Effective scaling for large MoE architectures

Abstract

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.