Sonar-MoE cuts memory use and speeds up large AI model training
Sonar-MoE: In-Backward Expert Updates for Memory-Efficient Mixture-of-Experts Training
Machine LearningArtificial IntelligenceComputation and LanguagePerformance
Summary
Training large AI models that split work across many parts takes a lot of computer memory. The authors introduce Sonar-MoE, a method that updates these parts one at a time during training, freeing memory quickly. This reduces the peak memory needed by about 20% and makes training faster, without hurting the results. Sonar-MoE achieves this by updating grouped parts instead of individual ones, avoiding slowdowns.
What this means in practice
- •For machine learning engineers: Reduce memory use and speed up training of large Mixture-of-Experts AI models by updating expert parameters in backward passes.
- •For ai infrastructure teams: Optimize GPU resources when scaling MoE models for production deployment with more efficient memory management during training.
Authors
Avni Mittal, Avinash Anand, Ashutosh Kumar, Dikshant Kukreja, Kritarth Prasad, Chaitanya Lakhchaura, Rishitej Reddy Vyalla, Erik Cambria, Zhengkui Wang, Rajiv Ratn Shah
Abstract
Most parameters of a Mixture-of-Experts (MoE) model belong to its experts, and standard training keeps every expert gradient in GPU memory until the optimizer step. This memory grows with the number and size of the experts, so it rises as MoE models scale. We introduce Sonar-MoE, which updates each expert tensor during the backward pass and frees its gradient right away. It brings optimizer-in-backward from dense models to batched MoE expert tensors, so the GPU never holds more than one layer of expert gradients. When we fine-tune OLMoE-1B-7B and Granite-3.1-3B-A800M in bf16 with two expert-kernel backends, Sonar-MoE lowers peak memory by 19% to 22%. Its updates are bitwise identical to standard training without expert gradient clipping, and held-out loss stays within 0.003 of all baselines. Each step also takes 17% to 23% less time than standard training with global-norm clipping. Updating whole batched tensors rather than individual experts is what keeps Sonar-MoE fast, as it avoids a 20% to 72% slowdown while using the same memory. Sonar-MoE matches optimizer-in-backward over all parameters in memory and speed while still clipping the dense parameters.