Sonar-MoE: Updating Experts During Backpropagation to Save 20% GPU Memory

When training Mixture-of-Experts (MoE) models, the memory bottleneck often lies not in the attention layers but in the expert parameters. The standard approach keeps all expert gradients resident on the GPU until the end of backpropagation, then performs a unified update — the more experts there are and the larger they are, the more daunting this memory cost becomes. A recent work on Sonar-MoE proposes a more memory-efficient approach: directly updating expert tensors during backpropagation and immediately releasing the gradients once updated.

Where the memory is saved: moving the optimizer into backpropagation

Most parameters in an MoE model are concentrated in the expert networks, and standard training keeps a full gradient copy for each expert until the optimizer step releases it. This means memory usage grows linearly with the number of experts and the size of each expert.

The core idea of Sonar-MoE is to migrate "optimizer-in-backward" from dense models to batched MoE expert tensors. As soon as the gradient of an expert tensor is computed, the update is executed immediately and the gradient is released, so at most one layer's worth of expert gradients is kept on the GPU. This layer-by-layer release strategy brings down peak memory.

Mixture-of-Experts (MoE) models replace the single feed-forward network (FFN) in a traditional Transformer with multiple parallel "expert" networks, and use a lightweight router to select and activate a small subset of experts for each token. Take OLMoE-1B-7B as an example: the model has about 7 billion total parameters, but each token activates only about 1 billion, so the compute of inference and training is close to that of a small model while the capacity is close to that of a large model. The cost lies in memory: all experts' parameters, gradients, and optimizer states must stay resident on the GPU, even though only some of them are used in each forward pass. The more experts there are and the wider each expert is, the more significant this overhead becomes — which is precisely why the memory bottleneck in MoE training is concentrated on the expert side rather than the attention side.

optimizer-in-backward is a training strategy that moves parameter updates earlier into the backpropagation phase. The standard training flow is: forward-compute the loss, backpropagate layer by layer to compute gradients and keep them all, then perform a single unified parameter update with the optimizer after the entire backpropagation finishes. optimizer-in-backward does the opposite: once the gradient of a layer (or a tensor) is computed, the optimizer immediately updates the corresponding parameters and then immediately releases that gradient, without waiting for the entire backpropagation to finish. Since the optimizer states (such as Adam's first and second moments) themselves must still stay resident, what is saved is the memory for gradient copies. This idea was previously mostly used for dense models; Sonar-MoE extends it to batched MoE expert tensors, making the "compute-then-update, update-then-release" pipeline applicable to expert parameters.

Measured data: memory down by a fifth, speed up rather than down

The researchers validated Sonar-MoE on two models: OLMoE-1B-7B and Granite-3.1-3B-A800M, both in bf16 precision, paired with two expert kernel backends.

The results show that peak memory dropped by 19% to 22%. This reduction isn't aggressive, but for MoE fine-tuning scenarios that are tight on memory, it's already enough to make previously infeasible configurations viable.

More critically, correctness: in the absence of expert gradient clipping, Sonar-MoE's updates are bitwise identical to standard training, and the holdout-set loss gap versus all baselines stays within 0.003. In other words, saving memory comes at no cost to model quality.

In terms of speed, each training step is 17% to 23% faster than standard training with global norm clipping. This is somewhat counterintuitive — typically memory-saving approaches slow down training, but Sonar-MoE, by releasing gradients early and reducing memory movement and peak pressure, actually gains an advantage in time.

"Bitwise identical" means that two training runs with the same inputs and the same random seed produce parameter updates and intermediate results that are identical at the binary level, not merely numerically close. In floating-point training, changes in operation order or different reduction methods can introduce tiny errors, so achieving bitwise identity shows that Sonar-MoE's in-backward update is mathematically equivalent to standard training, differing only in execution timing and memory-release strategy. The paper's statement that it is "bitwise identical in the absence of expert gradient clipping" indicates that when gradient clipping is introduced, bitwise identity no longer holds because clipping depends on the global gradient norm and requires aggregating information across tensors — but the holdout-set loss gap is still kept within 0.003, showing that the practical impact is negligible.

Why not update expert by expert

A natural question: if the goal is to save memory, why not simply update and release each expert separately?

The paper gives a clear answer: updating expert by expert causes a 20% to 72% slowdown, while memory usage is not lower than batched updates. Sonar-MoE chooses to update the entire batched tensor precisely to preserve speed. This design trade-off shows that in MoE training, "batching" itself is the key to performance, and splitting too finely is counterproductive.

Comparison with full-parameter optimizer-in-backward

Sonar-MoE performs in-backward updates only for expert parameters, while dense parameters still retain the conventional gradient clipping flow. Experiments show that it matches the "optimizer-in-backward for all parameters" approach in both memory and speed, while still being able to perform clipping normally on dense parameters. This means it does not give up the means needed for training stability just to save memory.

For researchers and engineering teams who need to fine-tune large MoE models on limited GPUs, this kind of "update-on-the-fly during backpropagation" idea offers a memory-optimization option that does not change the model architecture or sacrifice convergence quality. As MoE architectures become more widespread in large models, memory efficiency on the training side will attract increasing attention, and Sonar-MoE's approach is worth adding to the toolbox.

Global norm clipping is a commonly used stability technique when training large models: first compute the global L2 norm of all parameter gradients, and if it exceeds a preset threshold, scale all gradients by the same ratio, thereby limiting the magnitude of a single update step and preventing gradient explosion. A key constraint is that it can only be executed after all gradients have been computed, because the global norm must first be aggregated. This inherently conflicts with the "release-as-soon-as-computed" optimizer-in-backward. Sonar-MoE's compromise: perform in-backward updates only for expert parameters (no clipping on the expert side), while dense parameters still retain the full gradient collection and global norm clipping flow — thereby still capturing most of the memory savings without giving up the means for training stability.

One-line Summary

Sonar-MoE updates experts on the fly during backpropagation and releases gradients, saving about a fifth of memory while speeding up training.

Paragraph Summary

To address the memory bottleneck in MoE training caused by expert-parameter gradients staying resident on the GPU, Sonar-MoE migrates the optimizer-in-backward strategy to batched expert tensors: as soon as an expert tensor's gradient is computed, it is updated and released immediately, so at most one layer's worth of expert gradients is kept on the GPU. On OLMoE-1B-7B and Granite-3.1-3B-A800M, peak memory drops by 19%-22%, and without expert gradient clipping the updates are bitwise identical to standard training, with a holdout-set loss gap within 0.003. Speed is actually 17%-23% faster than standard training with global norm clipping. The paper also points out that updating expert by expert causes a 20%-72% slowdown without saving more memory, so batched updates are the key to preserving speed; dense parameters still retain the conventional clipping flow, balancing training stability.