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

Introduction

Mixture-of-Experts (MoE) models scale their parameter count without a matching increase in per-token compute, since each token is routed to only a few of many experts ([34, 19, 7, 16, 4]). The cost is paid in memory. A model that activates one billion parameters per token still has to keep the weights, gradients and optimizer state of every parameter during training. In OLMoE-1B-7B (Muennighoff et al. [28]), for example, the experts hold 6.44B of the 6.92B parameters, and their bf16 gradients alone take 12.9 GB. On a single GPU, this gradient memory alone can decide whether fine-tuning fits at all.

Much of this gradient memory is held for no reason. A standard training step runs the whole backward pass and only then calls the optimizer, so the gradient of the last MoE layer waits in memory until the backward pass has reached the first layer. An element-wise optimizer such

as AdamW (Loschilov & Hutter [24]) does not need this wait, since the update of a tensor depends only on that tensor’s gradient and optimizer state. Prior work on dense models fuses the update into the backward pass ((Lv et al., 2024b;a; Lewandowski & Kosson, 2023)), and PyTorch provides a hook that fires as soon as a parameter’s gradient is ready (Paszke et al. [30]). Updating every parameter this way frees all gradient memory, but it gives up global-norm clipping, because no gradient survives to the end of the step. In an MoE the experts hold almost all of the gradient memory, so it is enough to update only the experts early.

We introduce SONAR-MoE, which updates the experts inside the backward pass. It attaches a hook to each batched expert tensor of every MoE layer, and when the tensor’s gradient is ready, the hook applies a fused AdamW update and frees the gradient, so the GPU holds at most one layer of expert gradients at any time (Figure 2). All other parameters keep the standard path with global-norm clipping. The placement of the hooks also matters. Expert weights can be stored per expert or as batched tensors ([8, 35, 10]), and SONAR-MoE hooks the batched tensors. With $E$ experts per layer and $L$ MoE layers, per-expert hooks fire $2EL$ times per step, while one hook per batched weight tensor fires only $2L$ times. On OLMoE ($E=64$, $L=16$) this is 32 hooks per step instead of 2,048, for the same update and the same memory.

In bf16 fine-tuning, SONAR-MoE lowers peak memory from 55.6 to 43.6 GB on OLMoE and by 21.4% on Granite3.1-3B-A800M (Figure 1). Its updates are bitwise identical to those of standard training without expert-gradient clipping, and held-out loss stays within 0.003 of every baseline. Each step is also 19–23% shorter than in standard training with global-norm clipping (Table 1). Two effects shorten the step. The expert gradients skip the global-norm clipping kernels, which gives most of the saving, and on Granite the in-backward updates fill GPU time that would otherwise sit idle, which cuts a further 11.6% at the same update rule. Applying optimizer-in-backward to all parameters reaches the same memory and speed but cannot clip any gradient. SONAR-MOE matches it on both while keeping global-norm clipping for the dense parameters. Our contributions are:

Table 1. Alpaca fine-tuning for 2,000 steps with batch size 2 and seed 42 on two expert-kernel backends; every row uses the same data order. $\Delta$ is relative to standard training with global clipping on the same model and backend. Step time is the median over steps 201–2,000. Same-seed repeats vary by at most 1.4% in step time and 0.0016 in held-out loss. All rows of one model and backend ran on one GPU; the OLMoE SonicMoE rows ran on a different GPU from all other rows, so compare step times only within a model and backend. Bold marks the lowest memory and step time per model and backend.

Model	Backend	Method	Peak alloc. GB ($\Delta$)	Peak res. GB	Step ms ($\Delta$)	Tok/s	Held-out loss
OLMoE-1B-7B	SonicMoE	Standard (global clip)	55.55	56.42	142.8	768	1.2577
OLMoE-1B-7B	SonicMoE	Standard (no expert clip)	55.55 (0.0%)	56.42	115.9 (−18.8%)	958	1.2591
OLMoE-1B-7B	SonicMoE	Optimizer-in-backward (all)	43.35 (−22.0%)	44.24	114.4 (−19.9%)	970	1.2612
OLMoE-1B-7B	SonicMoE	SONAR-MOE	43.57 (−21.6%)	44.24	116.1 (−18.7%)	948	1.2596
OLMoE-1B-7B	ScatterMoE	Standard (global clip)	55.55	57.17	160.2	688	1.2585
OLMoE-1B-7B	ScatterMoE	Standard (no expert clip)	55.55 (0.0%)	57.17	133.3 (−16.7%)	825	1.2601
OLMoE-1B-7B	ScatterMoE	Optimizer-in-backward (all)	43.60 (−21.5%)	44.79	131.0 (−18.2%)	841	1.2610
OLMoE-1B-7B	ScatterMoE	SONAR-MOE	43.81 (−21.1%)	44.80	133.2 (−16.8%)	826	1.2601
Granite-3.1-3B	SonicMoE	Standard (global clip)	26.91	28.03	109.1	1111	1.1625
Granite-3.1-3B	SonicMoE	Standard (no expert clip)	26.91 (0.0%)	28.02	95.5 (−12.4%)	1275	1.1643
Granite-3.1-3B	SonicMoE	Optimizer-in-backward (all)	21.15 (−21.4%)	21.90	86.5 (−20.7%)	1421	1.1630
Granite-3.1-3B	SonicMoE	SONAR-MOE	21.14 (−21.4%)	22.03	84.4 (−22.7%)	1450	1.1616
Granite-3.1-3B	ScatterMoE	Standard (global clip)	26.89	28.49	106.8	1151	1.1591
Granite-3.1-3B	ScatterMoE	Standard (no expert clip)	26.89 (0.0%)	28.49	93.7 (−12.3%)	1230	1.1582
Granite-3.1-3B	ScatterMoE	Optimizer-in-backward (all)	21.77 (−19.1%)	22.40	90.0 (−15.8%)	1362	1.1585
Granite-3.1-3B	ScatterMoE	SONAR-MOE	21.75 (−19.1%)	22.34	88.2 (−17.4%)	1390	1.1582

• SONAR-MOE, which updates each batched expert tensor inside the backward pass and frees its gradient at once. Retained expert-gradient memory is bounded by one layer, peak memory falls by 21.6% (OLMoE) and 21.4% (Granite), and the updates are bitwise identical to standard training without expert-gradient clipping (§ 4, § 5.2).

• A speedup over standard training, with steps 19–23% shorter than under global-norm clipping, and a kernel-level profile that explains it. The expert gradients skip the clipping kernels, and on Granite the AdamW kernels run in otherwise idle GPU time (§ 5.3). Hooking batched tensors rather than single experts keeps this speed, since per-expert hooks make steps 20–72% slower (§ 4.3).

• A direct comparison with all-parameter optimizer-in-backward, which SONAR-MOE matches in memory and speed while keeping global-norm clipping, and a measured operating range across batch sizes, gradient checkpointing, 8-bit AdamW and gradient accumulation (§ 5.5, § 5.6).

Related Work

MoE training systems. Systems work on MoE training has made the expert computation faster through better kernels ([8, 35, 10]), better parallelism and communication across devices ([32, 14, 17, 41] (He et al. 2022)), and a cheaper expert backward pass ([22, 40]). Data-parallel training overlaps the gradient all-reduce with the backward pass (Li et al., 2020). None of this changes when gradients are released.

Training memory. ZeRO partitions optimizer state, gradients and parameters across ranks ([31]), and ZeRO-Offload moves them to CPU memory ([33]). Activation checkpointing and reversible blocks reduce activation memory ([1, 23]), 8-bit optimizers shrink optimizer state ([5]), and Collage trains fully in bf16 without fp32 master weights (Yu et al., 2024). Other methods train fewer parameters or a smaller subspace: LoRA and QLoRA train low-rank adapters ([6] (Hu et al. 2021)), ESFT fine-tunes only the experts most relevant to a task ([37]), GaLore projects gradients to a low-rank subspace ([42]), and BAdam updates one block at a time ([25]). SONAR-MOE updates every expert in full and changes only when its gradient is freed, so these savings add to ours (§ 5.5).

Optimizer updates in the backward pass. The closest work fuses the optimizer update into the backward pass of dense models. LOMO does this with SGD (Lv et al., 2024b), AdaLomo adds adaptive learning rates (Lv et al., 2024a), Lewandowski & Kosson [20] run the optimizer step during back-propagation, and GaLore uses per-layer updates to avoid storing weight gradients ([42]). FORGE fuses gradient computation with optimizer updates in registers ([18]), whereas SONAR-MOE uses hooks to update and free materialized gradients of batched expert tensors. PyTorch’s post-accumulate hooks make this available for every parameter ([30]), which we call all-parameter optimizer-in-backward. We do not claim the in-backward update itself, only the choice of where to apply it. Hooking every parameter frees all gradient memory but rules out global-norm clipping ([29]), while hooking only the batched expert tensors frees almost as much and keeps clipping for the rest of the model.

Background

MoE layers. Let $d$ be the hidden size, $n$ the expert intermediate size, $E$ the number of experts, $K$ the number of experts used per token and $L$ the number of MoE layers. A router scores each token $x_t$, keeps the top-$K$ experts $S_t$ and mixes their outputs with weights $w_e(x_t)$ from the router softmax:

Each expert is a gated MLP with a fused gate/up projection of size $2n \times d$ and a down projection of size $d \times n$. Existing MoE implementations store the experts of a layer in two batched tensors, $W_{\mathrm{up}}^{(l)} \in\mathbb{R}^{E \times2n \times d}$ and $W_{\mathrm{down}}^{(l)} \in\mathbb{R}^{E \times d \times n}$, and compute each weight gradient with one grouped GEMM over all experts instead of $E$ separate GEMMs ([8, 35, 10]). SONAR-MOE keeps this layout and attaches its hooks to these two tensors.

Memory in a training step. A training step keeps four kinds of data on the GPU. Parameters and optimizer state stay resident across steps. Activations are freed during backward as soon as they have been used. Gradients stay in memory until the optimizer step, so by the end of backward every gradient of the model is held at once. In our setting, weights, gradients and AdamW states are all kept in bf16, so the gradients take as much memory as the weights. This gradient memory is what SONAR-MOE releases early.

Method

Gradient Memory in Standard Training

Write $|W_e^{(l)}|$ for the number of weights of expert $e$ in layer $l$. A standard step keeps every expert gradient until the optimizer runs, so at the end of backward the GPU holds

gradient elements for the experts alone. For OLMoE ($L=16$, $E=64$, $d=2048$, $n=1024$) this is 6.44B elements, 12.9 GB in bf16. None of these gradients is needed after its own update, and for an element-wise optimizer that update could have happened the moment the gradient was complete.

In-Backward Expert Updates

SONAR-MOE removes the floor with PyTorch’s register_post_accumulate_grad_hook, which calls a function as soon as a parameter’s gradient has been fully accumulated. We register one hook on each batched expert tensor, $W_{\mathrm{up}}^{(l)}$ and $W_{\mathrm{down}}^{(l)}$, and remove these tensors from the end-of-step optimizer. When the hook fires, it applies a fused AdamW update to the tensor using its gradient and optimizer state, and then sets the gradient to None so its memory can be reused. All other parameters, such as attention, embeddings and routers, keep the usual end-of-step update. Algorithm 1 summarizes the training Algorithm 1 SONAR-MOE training step.

1. Setup: for each MoE layer $l$, register a hook on $W_{\mathrm{up}}^{(l)}$ and $W_{\mathrm{down}}^{(l)}$ ($2L$ hooks); exclude them from the end-of-step optimizer.

2. Hook($W$): fused AdamW update of $W$ with $W.\mathrm{grad}$; set $W.\mathrm{grad} \leftarrow\mathrm{None}$.

3. Set the learning rate $\eta_t$ of the hooked and end-of-step optimizers.

4. Forward pass as usual.

5. for each MoE layer $l = L, \ldots, 1$ in backward do

6. $\nabla W_{\mathrm{down}} \leftarrow$ one grouped GEMM over all experts

7. Hook$(W_{\mathrm{down}}^{(l)})$ $\triangleright$ update, then free

8. $\nabla W_{\mathrm{up}} \leftarrow$ one grouped GEMM over all experts

9. Hook$(W_{\mathrm{up}}^{(l)})$ $\triangleright$ update, then free

10. pass the input gradient to layer $l - 1$

11. end for

12. Optimizer step on non-expert parameters only.

step.

Because the backward pass visits the MoE layers one at a time, each expert gradient is freed before the next layer’s gradients are computed. The retained expert-gradient memory is therefore bounded by one layer:

For OLMoE this is a 16-fold reduction, from 12.9 GB to 0.8 GB (one layer). Without gradient accumulation, this is the same update that standard training applies when its clipping excludes the expert tensors, and we verify that the two are bitwise identical (Table 4). Because the hooks read the learning rate during the backward pass, the learning rate for step $t$ must be set before backward rather than after it.

Table 4. Weights that differ between two runs after three steps on OLMoE (6.44B expert weights, 477M other weights).

Runs compared	Expert weights	Other weights
Same baseline, run twice	0	6
SONAR-MoE vs. standard (no expert clip)	0	6
SONAR-MoE: per-expert vs. batched hooks	0	5
Opt.-in-backward (all) vs. standard (no clip)	0	6
Standard: global clip vs. no expert clip	3.04B (47%)	91.5M (19%)

Hook Granularity

The same schedule can be built on per-expert parameters, with one hook and one optimizer per expert slice, so each step makes $2EL$ optimizer calls instead of $2L$. Since AdamW is element-wise, one fused update on an $(E, 2n, d)$ tensor gives exactly the same result as $E$ fused updates on its slices. What changes is the number of hook calls per step:

Every call has a fixed cost, because autograd has to schedule a Python callback and the callback launches its own small kernels. If $\tau_{\mathrm{hook}}$ is this fixed cost, the per-expert layout pays $2EL\,\tau_{\mathrm{hook}}$ per step and the batched layout pays $2L\,\tau_{\mathrm{hook}}$. For OLMoE this means 2,048 calls instead of 32, each doing $E$ times less work. Modern MoE designs favour many small experts ([3, 4]), which makes the gap grow with $E$. On OLMoE ($E=64$), per-expert hooks add 23.6 ms per step (+20%), about 12 $\mu$s per extra optimizer call. On Granite ($E=40$, 32 layers) they add 61.1 ms (+72%), about 24 $\mu$s per extra call. Both layouts use the same memory and the per-expert layout is slower (Table 3), so SONAR-MOE attaches the hooks to the batched tensors.

Table 3. Hook granularity. Both layouts hold the same memory; per-expert hooks make $2EL$ optimizer calls per step instead of $2L$.

Model	Hooks	Calls	Peak GB	Step ms	Held-out
OLMoE	batched	32	43.57	116.1	1.2596
OLMoE	per-expert	2048	43.58	139.7	1.2600
Granite	batched	64	21.14	84.4	1.1616
Granite	per-expert	2560	21.37	145.5	1.1606

Clipping and Gradient Accumulation

Updating during backward changes two things that training loops usually take for granted, and our evaluation keeps both visible.

Global-norm clipping. Clipping by the global gradient norm needs every gradient before any update, but SONAR-MOE has already updated and freed the expert gradients by then. We therefore compare with two standard baselines, one that clips all gradients and one that clips only the non-expert gradients. The second performs exactly our update (§ 5.1).

Gradient accumulation. Gradient accumulation sums the gradients of $A$ microbatches and applies one optimizer step to their mean $\bar{g} = \frac{1}{A}\sum_a g_a$. A post-accumulate hook fires as soon as each microbatch’s gradient is ready, so the hook schedule applies one AdamW update per microbatch instead. These are not equivalent to a single update using the combined gradients, because both the weights and AdamW’s optimizer state change between microbatches. The first and second moments can both differ between these schedules. The standard schedule updates $v_t$ once per step with $\bar{g}^2$, while the hook schedule updates it $A$ times with the individual $g_a^2$:

For independent microbatch gradient noise of equal variance, averaging $A$ microbatches divides the noise variance in $\bar{g}$ by $A$, whereas the hook schedule uses the individual noisy gradients in its second-moment updates. The two schedules coincide at $A=1$. We therefore use $A=1$ in all main experiments and evaluate $A>1$ separately in § 5.6. Table 8 gives the measured cost of these per-microbatch updates.

$A$	Std. (global)	Std. (no exp.)	SONAR-MOE	Peak GB
1	143.4	117.3	116.4	55.55 / 43.57
2	217.8	191.9	229.4	57.21 / 44.31
4	368.5	341.7	451.7	57.21 / 44.31
8	665.8	639.6	895.6	57.21 / 44.32

Experiments

Experimental Setup

Models and data. We pick two public MoE models of different sizes and expert counts that both run on batched expert kernels. We fine-tune allenai/OLMoE-1B-7B-0924 (Muennighoff et al., [28]), (16 MoE layers, 64 experts, top-8) and `ibm-granite/ granite-3.1-3b-a800m-base (IBM Granite Team, 2024) (32 MoE layers, 40 experts) on 8,000 Alpaca examples (Taori et al., [36]) for 2,000 steps. A held-out set of 500 further examples (indices 8000–8499 of the same seed-42 shuffle) is disjoint from training. We use a batch size of 2, sequences of at most 256 tokens and no gradient accumulation. The learning rate peaks at $2 \times10^{-5}$ after 100 warm-up steps and decays with a cosine schedule toward $2 \times10^{-6}$. AdamW uses $\beta=(0.9, 0.95)$, $\epsilon=10^{-8}$ and no weight decay. Training is in bf16, with weights, gradients and AdamW states all kept in bf16 and no fp32 master weights. All arms use seed 42 and the same data order, which we check by hashing the first batches.

Baselines. Standard training with global clipping runs the backward pass, clips the global gradient norm of all parameters to 1.0 and takes one fused AdamW step. This is the usual default. Standard training without expert-gradient clipping is the same but clips only the non-expert parameters. It performs exactly the SONAR-MOE update, so any speed difference between the two comes from when the update happens rather than from the update itself. All-parameter optimizer-in-backward updates every parameter in its own backward hook, so no clipping is possible. We also run SONAR-MOE with per-expert hooks, one optimizer per expert slice, as an ablation.

Backends. We run every method with two expert-kernel libraries that store experts in batched tensors, SonicMoE (Guo et al., [10]) and ScatterMoE (Tan et al., [35]). Table 1 reports both; the other experiments use SonicMoE.

Metrics and hardware. We report the median step time over steps 201–2,000, with a GPU synchronization at the end of each step, throughput in supervised response tokens per second (about 111 per step), and peak allocated and peak reserved GPU memory. Same-seed repeats vary by at most 1.4% in step time and 0.0016 in held-out loss. Runs use two NVIDIA RTX PRO 6000 Blackwell GPUs (96 GB) with PyTorch 2.11 and CUDA 13.1. All arms of one model and backend ran on the same GPU, and the two GPUs have different power limits, so we compare step times only within a model and backend. Downstream accuracy is measured with lm-evaluation-harness (Gao et al., [9]) on the full test sets of ARC-Challenge (Clark et al., [2]) and HellaSwag (Zellers et al., [39]) (0-shot, acc norm) and MMLU (Hendrycks et al., [12]) (5-shot, accuracy).

layer that is still held. Held-out loss stays within 0.003 of every baseline. Speed needs a more careful reading. Standard training without expert-gradient clipping is as fast as SONAR-MOE on OLMoE (115.9 vs. 116.1 ms), so on OLMoE the whole speedup over standard training comes from not clipping the expert gradients. On Granite SONAR-MOE is a further 11.6% faster than that baseline, which § 5.3 traces to idle GPU time. All-parameter optimizer-in-backward matches SONAR-MOE in both memory and speed and differs only in giving up global-norm clipping on the dense parameters. ScatterMoE shows the same pattern. SONAR-MOE lowers peak memory by 21.1% on OLMoE and 19.1% on Granite, is as fast as standard training without expert-gradient clipping on OLMoE (133.2 vs. 133.3 ms), and is 5.9% faster than that baseline on Granite (88.2 vs. 93.7 ms), a smaller gain than with SonicMoE.

Loss alone does not show whether a fine-tuned model is still useful, so we also evaluate the OLMoE checkpoints after 2,000 steps on downstream tasks (Table 2). Fine-tuning raises ARC-Challenge accuracy by 3–5 points and leaves HellaSwag and MMLU almost unchanged. The two training methods differ by 0.016 on ARC-Challenge, about one standard error, and by at most 0.001 on HellaSwag and MMLU. We evaluate only OLMoE and omit GSM8K, where this model stays near 0.1 accuracy after Alpaca fine-tuning.

Table 2. Downstream accuracy on OLMoE after 2,000 steps (full test sets). Standard errors are 0.015 for ARC-Challenge and 0.004 for HellaSwag and MMLU in every row.

Checkpoint	ARC-C	HellaSwag	MMLU
Pretrained	0.482	0.771	0.527
Standard (global clip)	0.532	0.771	0.521
SONAR-MOE	0.516	0.772	0.521

Main Results

SONAR-MOE saves the memory that the mechanism predicts (Table 1). With SonicMoE, peak memory falls by 12.0 GB on OLMoE and 5.8 GB on Granite, which is the expert-gradient buffer of 12.9 and 6.0 GB minus the one

Sources of the Speedup

Three effects explain the step times (Figure 3). First, global-norm clipping of the expert gradients costs about 25 ms of kernel time per step on OLMoE (26.9 vs. 1.8 ms). It reads the 12.9 GB of gradients to compute the norm, then reads and writes them again to rescale them. This accounts for nearly all of the difference between standard training and SONAR-MoE on OLMoE. The cost scales with the number of expert parameters rather than with the batch, which is why it matters most for small batches (§5.5). Second, on Granite the GPU work is the same with and without in-backward updates (57.7 vs. 57.4 ms busy), but idle time falls from 65 to 42 ms. Granite’s backward pass is limited by kernel-launch overhead, and the in-backward AdamW kernels run in these gaps. OLMoE leaves little idle time, so it gains almost nothing from this. Third, attaching hooks per expert keeps the memory saving but adds 23.6 ms per step on OLMoE, about 12 $\mu$s per extra optimizer call Table 3.

Exactness and Training Quality

Update exactness. Without gradient accumulation, SONAR-MoE performs the same arithmetic as standard training that clips only the non-expert gradients. We train OLMoE for three steps with a learning rate of $10^{-4}$, large enough that bf16 updates are visible, and compare the weights of two runs element by element Table 4. SONAR-MoE and standard training without expert-gradient clipping give exactly the same 6.44B expert weights. They differ in 6 of the 477M other weights, and two runs of the same baseline also differ in 6 elements, so this is ordinary run-to-run noise, not an effect of the schedule. Per-expert and batched hooks also give the same expert weights, and all-parameter optimizer-in-backward gives exactly the same weights as standard training without clipping.

Clipping the expert gradients, as standard training does, changes 47% of the expert weights after three steps, yet changes held-out loss on OLMoE by only 0.0014 after 2,000 steps Table 1. AdamW divides each update by a running estimate of the gradient scale, which largely cancels a uniform rescaling.

Pure-bf16 updates. Because weights are stored in bf16, an update smaller than half a unit in the last place of a weight is rounded away. Table 5 shows that after 2,000 steps most weights still hold their pretrained values, and that the two methods change the same share of weights.

Method	Expert weights	Other weights
Standard (global clip)	19.3%	22.4%
SONAR-MOE	19.3%	22.5%

Batch Size and Memory Composition

The cost of clipping expert gradients is fixed per step (22–29 ms), so the speedup over standard training shrinks from 18% at batch size 2 to 8% at batch size 32 (Figure 4). With the same update rule, SONAR-MOE and standard training are within ±1.1% at batch sizes 2–16, and at batch size

32 SONAR-MOE is 4.0% faster. The memory saving also shrinks with batch size, because activations, not gradients, set the peak at large batches. The saving in peak allocated memory is 12.0 GB at batch size 2, 9.4 GB at 8, 5.2 GB at 16 and none at 32, where only 2.2 GB of reserved memory is saved. SONAR-MOE helps most when batches are small relative to the model. In practice the saving can be read off Eq. (2): it equals the expert-gradient buffer as long as gradients, not activations, set the peak.

Gradient checkpointing does not lower peak memory at batch size 2, because activations are small, and only costs time (+10–16%, Table 6). 8-bit AdamW and SONAR-MOE save different memory, optimizer states and gradients, so they combine. Together they need about 30 GB and run under a 30.5 GB cap, where 8-bit AdamW alone runs out of memory. All-parameter optimizer-in-backward gives the same savings. It also runs under the 46 GB cap, and with 8-bit AdamW its peak allocated memory matches ours (29.7 vs. 29.9 GB). 8-bit AdamW is slower per step (186 vs. 141 ms).

Setup	Alloc.	Res.	ms	46 GB	30.5 GB
Standard	55.55	56.40	140.8	OOM	–
Standard + GC	55.48	56.16	155.6	–	–
SONAR-MOE	43.56	44.22	111.8	✓	–
SONAR-MOE + GC	43.23	44.03	130.1	–	–
Opt.-in-bwd	43.35	44.24	109.9	✓	–
8-bit AdamW	41.93	43.13	186.0	–	OOM
8-bit + SONAR-MOE	29.94	30.52	165.4	–	✓
8-bit + SONAR-MOE + GC	29.61	29.78	189.5	–	✓
8-bit + opt.-in-bwd	29.74	30.69	165.0	–	OOM

8-bit AdamW quality. For the combination that fits OLMoE fine-tuning on a 32 GB GPU, we also check that quality is preserved. We use the setup of Table 1 (2,000 steps, batch size 2, seed 42) but store the AdamW states in 8 bits with bitsandbytes AdamW8bit ([5]). With SONAR-MOE, each expert tensor gets its own 8-bit optimizer inside its hook, and the other parameters use one 8-bit optimizer at the end of the step with global-norm clipping, as in Table 1. Held-out loss stays within 0.003 of standard training (Table 7), so the memory saving of 8-bit AdamW and SONAR-MOE comes at no visible cost in quality.

Configuration	Held-out loss
Standard (global clip)	1.2577
8-bit AdamW	1.2574
8-bit AdamW + SONAR-MOE	1.2601

Gradient Accumulation

As described in § 4.4, hooks fire once per microbatch, so with $A$ microbatches SONAR-MOE takes $A$ optimizer steps where the baseline takes one. Table 8 sweeps $A \in\{1, 2, 4, 8\}$. At $A=1$ SONAR-MOE is as fast as standard training without expert-gradient clipping (116.4 vs. 117.3 ms). With accumulation it runs the optimizer $A$ times per step and is slower than standard training with global clipping, by 5% at $A=2$, 23% at $A=4$ and 35% at $A=8$, while still saving 12.9 GB. Each extra update is a full AdamW pass over the 6.44B expert weights, about 37 ms at every $A$, so the overhead grows linearly with $A$. The method also becomes a different optimizer, so we recommend SONAR-MOE for $A=1$, or for settings where memory is the binding constraint and the change of update rule is acceptable.

Discussion

Practical memory gains. The most direct benefit is that a run fits where it did not before. Fine-tuning OLMoE-1B-7B in our setup peaks at 55.6 GB with standard training and at 43.6 GB with SONAR-MOE, which fits on a 48 GB GPU, since it runs under a 46 GB allocator cap where standard training runs out of memory. With 8-bit AdamW it fits a 32 GB GPU (30.5 GB cap), where 8-bit AdamW alone does not. All-parameter optimizer-in-backward reaches the same peak memory and also fits under the 46 GB cap. Gradient checkpointing does not reduce peak memory at batch size 2 (§ 5.5), and the saving shrinks as batches grow, because activations then set the peak.

Operating conditions. Our results point to a simple rule. Early release is a good default when four conditions hold: the optimizer is element-wise, gradients are not accumulated across microbatches, the experts are stored in batched tensors, and the expert gradients are not needed for global-norm clipping. The memory saving is largest at small batch sizes. A speed benefit appears only when the baseline clips expert gradients or when the backward pass leaves the GPU idle. With accumulation the method becomes a different and slower optimizer (§ 5.6), and with per-expert hooks it keeps the memory saving but loses time (§ 5.3). We think this kind of explicit operating range is more useful to practitioners than a single headline number, and we have tried to report every regime we measured, including the ones where the method does not help.

Why hook granularity is overlooked. Optimizer-in-backward methods were developed for dense models, where each weight matrix is one parameter and the question of hook granularity never arises. In MoE code it does arise. Expert weights are often stored as a list of per-expert modules for readability, or are split into per-expert views when a checkpoint is loaded. Both cases silently multiply the hook count by $E$. Our ablation shows that this keeps the memory saving intact but adds 20% (OLMoE) and 72% (Granite) step time (Table 3), so the problem is easy to overlook if only memory is measured.

Conclusion

In MoE training, the gradients of all experts usually sit in memory until the optimizer step, even though each one is needed only for its own update. SONAR-MOE updates every batched expert tensor as soon as its gradient is ready and frees the gradient right away, which bounds the retained expert-gradient memory by one layer. At batch size 2 it lowers peak memory by 21.6% on OLMoE and 21.4% on Granite, and its updates are bitwise identical to standard training without expert-gradient clipping. Its speed advantage over standard training comes from not clipping the expert gradients and, on models whose backward pass leaves the GPU idle, from running the optimizer kernels in that idle time. Hooks attached per expert save the same memory but make each step 20–72% slower, so early release should be attached to batched expert tensors. All-parameter optimizer-in-backward matches the method in memory and speed. Natural next steps are to support global-norm clipping without keeping all gradients in memory, and to support data parallelism by updating after the gradient all-reduce. We are also extending SONAR-MoE to inference, where expert computation can be started speculatively before the router’s decision is final and reconciled against the exact routing, so that the same early-release principle speeds up both training and inference of MoE models.

Limitations

Gradient clipping. Standard global-norm clipping needs all gradients before any update, while SONAR-MOE updates experts during backward. Compatibility can be restored by delaying the update and keeping, offloading or recomputing the gradients after the norm is known, or by clipping each expert tensor by its own norm inside the hook. We plan to evaluate these options in the next version. Gradient accumulation. With $A \ge2$ the method applies $A$ updates per step and becomes a different optimizer that is 5–35% slower than standard training with global clipping (§ 5.6), so our main results use $A=1$. The schedule is exact for element-wise optimizers such as AdamW, SGD or Lion, and extending it to optimizers that need statistics across tensors is left to future work. Scope of the evidence. All results use one GPU, and under data parallelism the update would run after the gradient all-reduce, which we will add by integrating with the communication hooks. The memory saving is largest at small batch sizes and shrinks as activations dominate. We train in bf16 and test models up to 7B parameters on Alpaca fine-tuning with one seed, using same-seed repeats to measure noise; larger models, more datasets and more seeds are planned for the next experiments.