PaperlayerView exact revision on arXiv

Source-linked AI summary

Sonar-MoE: In-Backward Expert Updates for Memory-Efficient Mixture-of-Experts Training

Avni Mittal, Avinash Anand, Ashutosh Kumar, Dikshant Kukreja, Kritarth Prasad, Chaitanya Lakhchaura, Rishitej Reddy Vyalla, Erik Cambria, Zhengkui Wang, Rajiv Ratn Shah

arXiv:2611.03577v1cs.LGcs.AIcs.CLcs.PF

TL;DR

MoE training retains gradients for every expert until the optimizer step, even though each gradient is needed for its own update. SONAR-MOE updates batched expert tensors during backward and frees their gradients; fine-tuning results show lower peak memory and shorter steps than standard training with global-norm clipping.

  • Problem

    Standard MoE training keeps every expert gradient in GPU memory until the optimizer step, although each is needed only for its own update.

    Source: §1 · p. 1

  • Method

    SONAR-MOE hooks batched expert tensors to apply fused AdamW updates when their gradients are ready, then frees those gradients while other parameters retain the standard path.

    Source: §1 · p. 1

  • Results

    Peak memory falls by 21.6% on OLMoE and 21.4% on Granite; steps are 19–23% shorter than standard training with global-norm clipping, and held-out loss stays within 0.003 of every baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SONAR-MOE matches all-parameter optimizer-in-backward in memory and speed while keeping global-norm clipping for dense parameters.

    Source: §1 · p. 2

    With A≥2 microbatches, SONAR-MOE applies A updates per step, becomes a different optimizer, and runs 5–35% slower than standard training with global clipping; the authors recommend A=1 or memory-bound settings where that change is acceptable.

    Source: §5.6 · p. 8, LIMITATIONS · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

SONAR-MOE addresses the memory cost of retaining all expert gradients until the optimizer step by updating batched expert tensors during backward and freeing their gradients. In fine-tuning, it reduces peak memory and step time while preserving clipping for dense parameters.

1 INTRODUCTION

  • Standard training retains expert gradients until the optimizer step, although an element-wise optimizer such as AdamW needs only each tensor’s gradient and optimizer state for its update.Experts hold 6.44B of OLMoE-1B-7B’s 6.92B parameters, and their bf16 gradients alone take 12.9 GB.Source: §1 · p. 1
  • SONAR-MOE updates each batched expert tensor as its gradient becomes ready, freeing it so retained expert-gradient memory stays bounded by one layer.Dense parameters remain on the standard path with global-norm clipping.Source: §1 · pp. 1–2
  • Peak memory falls from 55.6 to 43.6 GB on OLMoE and by 21.4% on Granite-3.1-3B-A800M in bf16 fine-tuning.Figure 1 covers Alpaca fine-tuning with SonicMoE and ScatterMoE backends.Source: §1 · p. 1, Figure 1 · p. 2
  • Each step is 19–23% shorter than standard training with global-norm clipping, while held-out loss stays within 0.003 of every baseline.The expert gradients skip global-norm clipping kernels, which gives most of the time saving.Source: §1 · p. 2
  • SONAR-MOE matches all-parameter optimizer-in-backward in memory and speed while retaining global-norm clipping for dense parameters.Updating batched tensors avoids the 20–72% step-time slowdown of per-expert hooks while using the same memory.Source: §1 · p. 2

2 RELATED WORK

Prior MoE systems work improves expert computation, parallelism, communication, or training memory, but does not change when gradients are released. SONAR-MOE applies in-backward updates selectively to batched expert tensors, retaining clipping for the rest of the model.

2 RELATED WORK

  • Prior MoE systems work improves kernels, parallelism, communication, or expert backward computation, while data-parallel all-reduce overlaps backward; none changes when gradients are released.SONAR-MOE hooks update and free materialized gradients of batched expert tensors.Source: §2 · pp. 2–3
  • Other approaches reduce memory through sharding, offloading, checkpointing, quantized optimizers, or fewer trained parameters; SONAR-MOE updates every expert in full, so those savings can add to its own.The related work includes ZeRO, 8-bit optimizers, LoRA, ESFT, GaLore, and BAdam.Source: §2 · p. 2
  • Dense-model optimizer-in-backward work includes LOMO, AdaLomo, and backward-pass optimizer steps; SONAR-MOE does not claim the update itself, but its placement on batched experts.The related-work passage also describes GaLore and FORGE.Source: §2 · pp. 2–3
  • SONAR-MOE’s distinction is where in-backward updates are applied: batched expert tensors are updated and freed, while the remaining model parameters retain global-norm clipping.Updating every parameter frees all gradient memory but rules out global-norm clipping.Source: §2 · p. 3, Figure 2 · p. 3

3 BACKGROUND

MoE layers route each token to a small selected set of experts, whose outputs are mixed using router weights. Implementations store layer experts in batched tensors, and training ordinarily retains gradients until the optimizer step.

3 BACKGROUND

  • A router scores each token, selects its top-K experts, and mixes their outputs with router-softmax weights.The setup defines d as hidden size, n as expert intermediate size, E as number of experts, K as experts used per token, and L as MoE layers.Source: §3 · p. 3
  • Each expert is a gated MLP with a fused gate/up projection of size 2n × d and a down projection of size d × n..Source: §3 · p. 3
  • MoE implementations store each layer’s experts in two batched tensors and compute each weight gradient with one grouped GEMM over all experts rather than E separate GEMMs.SONAR-MOE attaches its hooks to those two tensors.Source: §3 · p. 3
  • Parameters and optimizer state stay resident across steps, activations are freed during backward, and gradients ordinarily remain until the optimizer step.With weights, gradients, and AdamW states kept in bf16, gradients take as much memory as weights; SONAR-MOE releases expert gradients early.Source: §3 · p. 3

4 METHOD

SONAR-MOE updates batched expert tensors as their gradients become ready during backward, freeing them immediately and bounding retained expert-gradient memory to one layer. Non-expert parameters retain the standard end-of-step path, while hook granularity and gradient accumulation determine important trade-offs.

4.1 Gradient Memory in Standard Training

  • Standard training retains every expert gradient until the optimizer runs, although each gradient could be updated once complete.OLMoE expert gradients alone total 6.44B elements, or 12.9 GB in bf16.Source: §4.1 · p. 3

4.2 In-Backward Expert Updates

  • Hooks apply fused AdamW updates to batched expert tensors as their gradients become ready, then free those gradients; other parameters keep the end-of-step update.Retained expert-gradient memory is bounded by one layer because backward visits MoE layers one at a time.Source: §4.2 · pp. 3–4
  • For OLMoE, retained expert-gradient memory falls 16-fold, from 12.9 GB to 0.8 GB, and updates are bitwise identical to standard training without expert-gradient clipping.The learning rate for step t must be set before backward because the hooks read it during that pass.Source: §4.2 · p. 4

4.3 Hook Granularity

  • On Granite, per-expert hooks add 61.1 ms per step (+72%) versus batched hooks, despite using the same memory.Batched layout requires 2L hook calls rather than 2EL, avoiding fixed callback and kernel-launch costs.Source: §4.3 · p. 4

4.4 Clipping and Gradient Accumulation

  • Global-norm clipping requires every gradient before any update, so the comparison distinguishes clipping all gradients from clipping only non-expert gradients.The non-expert-only clipping baseline performs exactly SONAR-MOE's update.Source: §4.4 · p. 4
  • With gradient accumulation above one microbatch, hooks update once per microbatch, which is not equivalent to one update using the combined gradients; the main experiments use A=1.The differing schedules can change AdamW's first and second moments.Source: §4.4 · p. 4

5 EXPERIMENTS

The experiments evaluate SONAR-MOE against standard and optimizer-in-backward baselines on two MoE models, measuring memory, speed, exactness, and training quality. Results show lower peak memory and shorter steps, with important qualifications for batch size and gradient accumulation.

5.1 Experimental Setup

  • The evaluation fine-tunes OLMoE-1B-7B and Granite-3.1-3B-A800M on 8,000 Alpaca examples for 2,000 steps, with a disjoint 500-example held-out set.Runs use bf16, batch size 2, no gradient accumulation, and two expert-kernel backends; timing comparisons are limited to the same model and backend.Source: §5.1 · pp. 4–5

5.2 Main Results

  • Peak memory falls by 12.0 GB on OLMoE and 5.8 GB on Granite with SonicMoE, while held-out loss stays within 0.003 of every baseline.SONAR-MOE is also as fast as standard training without expert-gradient clipping on OLMoE and faster on Granite; Table 1 reports both expert-kernel backends.Source: §5.2 · p. 5, Table 1 · p. 6

5.3 Sources of the Speedup

  • Skipping expert-gradient clipping kernels accounts for most of the speedup, while on Granite in-backward AdamW updates also use otherwise idle GPU time.The profiler’s absolute times are higher than Table 1 because profiling adds CPU overhead, so only its bars should be compared.Source: Figure 3 · p. 5, §5.3 · p. 6

5.4 Exactness and Training Quality

  • Without gradient accumulation, SONAR-MOE’s expert weights are bitwise identical to standard training without expert-gradient clipping; per-expert and batched hooks also match.The two methods differ in six of 477M other weights, the same count seen between two runs of the baseline.Source: §5.4 · p. 6
  • Clipping expert gradients changes 47% of OLMoE expert weights after three steps, yet held-out loss differs by only 0.0014 after 2,000 steps; downstream methods differ by 0.016 on ARC-Challenge and at most 0.001 on HellaSwag and MMLU.After bf16 training, most weights still retain their pretrained values, and the two methods change the same share of weights.Source: §5.2 · p. 5, §5.4 · pp. 6–7

5.5 Batch Size and Memory Composition

  • The speedup over standard training shrinks from 18% at batch size 2 to 8% at batch size 32, while peak allocated-memory savings fall from 12.0 GB to none.At batch size 32, activations set the peak; SONAR-MOE is 4.0% faster, and memory savings are limited to 2.2 GB of reserved memory.Source: §5.5 · p. 7
  • At batch size 2, gradient checkpointing adds 10–16% time without lowering peak memory, while combining SONAR-MOE with 8-bit AdamW fits under a 30.5 GB cap.The combination preserves held-out loss within 0.003 of standard training; 8-bit AdamW alone runs out of memory under that cap.Source: §5.5 · p. 7, Table 7 · p. 8

5.6 Gradient Accumulation

  • With A microbatches, SONAR-MOE takes A optimizer steps and becomes a different optimizer; it is 5%, 23%, and 35% slower than globally clipped standard training at A=2, 4, and 8.Each extra update adds a full AdamW pass over the expert weights, while the method still saves 12.9 GB.Source: §5.6 · p. 8, Table 8 · p. 8

6 DISCUSSION

Sonar-MoE’s practical gains and speed depend on its operating regime and hook granularity. Memory savings can enable runs on smaller GPUs, while per-expert hooks preserve memory savings but slow steps.

6 DISCUSSION

  • 55.6 to 43.6 GB peak memory lets OLMoE fine-tuning fit under a 46 GB allocator cap on a 48 GB GPU.With 8-bit AdamW, it fits a 32 GB GPU under a 30.5 GB cap; the saving shrinks as activations set the peak at larger batches.Source: §6 · p. 8
  • Early release is a good default with an element-wise optimizer, no microbatch gradient accumulation, batched expert tensors, and no need for expert gradients in global-norm clipping.Memory savings are largest at small batch sizes; speed gains appear when expert gradients would otherwise be clipped or the backward pass leaves GPU time idle.Source: §6 · p. 8
  • 20% on OLMoE and 72% on Granite step-time increases result when per-expert hooks multiply the hook count by the number of experts, despite preserving memory savings.Hooks on batched expert tensors avoid this slowdown.Source: §6 · p. 8

7 CONCLUSION

Sonar-MoE releases expert gradients during backward, bounding retained expert-gradient memory by one layer. It reduces peak memory while matching standard updates without expert-gradient clipping, but its speed and use depend on the training regime.

7 CONCLUSION

  • At batch size 2, peak memory falls by 21.6% on OLMoE and 21.4% on Granite, with updates bitwise identical to standard training without expert-gradient clipping.Per-expert hooks save the same memory but make steps 20–72% slower; all-parameter optimizer-in-backward matches Sonar-MoE in memory and speed.Source: §7 · p. 8
  • Sonar-MoE updates each batched expert tensor as soon as its gradient is ready and frees it, bounding retained expert-gradient memory by one layer.Its speed advantage comes from avoiding expert-gradient clipping and, when backward leaves GPU time idle, running optimizer kernels during that time.Source: §7 · p. 8
  • The authors identify future work on global-norm clipping without retaining all gradients, data-parallel updates after all-reduce, and speculative inference for MoE models.The inference direction would start expert computation before routing is final and reconcile it with exact routing.Source: §7 · p. 8

LIMITATIONS

The method’s clipping compatibility, optimizer scope, distributed use, and empirical evidence remain bounded by the settings evaluated so far.

LIMITATIONS

  • Global-norm clipping is not compatible with immediate expert updates unless updates are delayed or gradients are retained, offloaded, recomputed, or clipped per expert tensor.These compatibility options are planned for evaluation in a future version.Source: LIMITATIONS · p. 9
  • With gradient accumulation A≥2, the method applies A updates per step, changes the optimizer, and runs 5–35% slower than standard training with global clipping.The main results therefore use A=1.Source: LIMITATIONS · p. 9
  • The update schedule is exact for element-wise optimizers such as AdamW, SGD, and Lion; optimizers requiring statistics across tensors remain future work.Source: LIMITATIONS · p. 9
  • All reported results use one GPU; data-parallel use would require running the update after gradient all-reduce, an integration planned for future work.Source: LIMITATIONS · p. 9
  • Evidence is limited to bf16 Alpaca fine-tuning of models up to 7B parameters with one seed, with larger models, more datasets, and more seeds planned.Same-seed repeats were used to measure noise.Source: LIMITATIONS · p. 9
  • Memory savings are largest at small batch sizes and shrink when activations dominate.Source: LIMITATIONS · p. 9

Related papers

Loading 2611.03577v1…