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

TL;DR

MoEBlaze tackles the memory wall in Mixture-of-Experts training by eliminating large intermediate activation buffers during token routing and by co-designing end-to-end training kernels with activation-aware strategies. It introduces lightweight data structures (e.g., per-expert and per-token index lists) and a three-step dispatch method that obviates global sorts, enabling on-the-fly token gathering and reduction. The approach is augmented with activation checkpointing and fused SwiGLU/MoE kernels, reducing activation traffic and improving GPU utilization. Empirically, MoEBlaze delivers over 4x end-to-end speedups and more than 50% memory savings versus state-of-the-art MoE frameworks on multiple configurations and activation functions, with especially large gains for SwiGLU architectures on modern GPUs.

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.

Figures (6)

• Figure 1: MoE in conventional approach vs. MoEBlaze. Left illustrates the conventional MoE computation, comprising token dispatch, expert computation, and weighted aggregation (details in section \ref{['sec:background']}). Right presents the proposed MoEBlaze framework, which introduces memory-efficient token-routing and expert computation (details in section \ref{['sec:memory_efficient_token_dispatch']}).
• Figure 2: Data structures for the memory-efficient MoE training.
• Figure 3: Activation memory footprint comparison between MoEBlaze and Megablocks across the set of MoE configs defined in Table \ref{['sample-table']} using SiLU activation function.
• Figure 4: Speedups of MoEBlaze w.r.t to Megablocks on the set of configurations in Table \ref{['sample-table']} using SiLU as the activation function.
• Figure 5: Activation memory footprint comparison between MoEBlaze and Megablocks across the set of MoE configs defined in Table \ref{['sample-table']} using SwiGLU activation function.