What is optimizer-in-backward?

It runs each parameter's optimizer update inside the backward pass, as soon as that parameter's gradient is complete, then frees the gradient. It works for element-wise optimizers such as AdamW, whose update of a tensor needs only that tensor's gradient and state.

Why does Sonar-MoE target the experts?

Experts hold most of an MoE model's parameters (6.44B of 6.92B in OLMoE-1B-7B), so they hold most of the gradient memory. Updating only the experts early frees almost all of it, and the other parameters keep the standard path with global-norm clipping.

What does Sonar-MoE give up?

Global-norm clipping on the expert gradients, because those gradients are freed before the end of the backward pass. With gradient accumulation the hook updates once per microbatch, which is not equivalent to one update on the averaged gradient.

How large is the saving?

In the paper's bf16 fine-tuning runs, peak memory fell 19–22% on OLMoE-1B-7B and Granite-3.1-3B-A800M, with held-out loss within 0.003 of the baselines. These are short fine-tuning runs, so treat the percentage as setup-dependent.

GPU· 2026-11-04

Sonar-MoE frees MoE expert gradients mid-backward — Optimizer-in-backward for experts — What does it mean?

TL;DR

What is it: Sonar-MoE (arXiv 2611.03577, November 2026) applies optimizer-in-backward to batched Mixture-of-Experts expert tensors: each expert tensor is updated the moment its gradient is ready, and the gradient is freed right away.
Why it’s needed: Experts hold most of an MoE model's parameters (6.44B of 6.92B in OLMoE-1B-7B), so their gradients dominate gradient memory (12.9 GB in bf16), and that can decide whether fine-tuning fits on a GPU at all.
vs previous: Standard training keeps every expert gradient until the optimizer step after the whole backward pass, and plain optimizer-in-backward frees every gradient but gives up global-norm clipping; Sonar-MoE frees only the expert gradients, so at most one layer's are resident, and keeps global-norm clipping for the dense parameters.

Jargon

MoE (Mixture-of-Experts): A Transformer whose feed-forward layer is split into many small expert networks; a router sends each token to only a few of them, so the model has many parameters but little compute per token.
Expert gradient: The gradient of the loss with respect to an expert's weights. Training needs one per expert weight, so it takes as much memory as the experts themselves.
AdamW: A common optimizer. It is element-wise: the update of a weight tensor needs only that tensor's gradient and its own two running averages (the optimizer state), nothing from other tensors.
Optimizer-in-backward: Running each parameter's optimizer update inside the backward pass, as soon as that parameter's gradient is complete, instead of after the whole backward pass.
Autograd hook: A function PyTorch calls automatically at a set moment. Sonar-MoE uses register_post_accumulate_grad_hook, which fires once a parameter's gradient has been fully accumulated.
Global-norm clipping: Rescaling all gradients together when their combined size is too large. It needs every gradient to exist at the same time, so it cannot be combined with freeing gradients early.
Batched expert tensor: One tensor that stores all experts of a layer together (for example shape experts × width × hidden size), so a single grouped matrix multiply covers every expert.
bf16: A 16-bit floating-point format, 2 bytes per number. Here weights, gradients and AdamW states are all bf16, with no fp32 master copy.

The news. On November 2, 2026, a research team posted Sonar-MoE (arXiv 2611.03577). In bf16 fine-tuning of OLMoE-1B-7B and Granite-3.1-3B-A800M it lowered peak GPU memory by 19–22% and cut step time by roughly 17–23% across the four model-and-backend pairs in its main table (the paper's abstract says 17–23%, other passages say 19–23%), against standard training with global-norm clipping, with held-out loss within 0.003 of every baseline. Read the paper →

Picture a restaurant kitchen where every dirty plate must stay on the counter until the dinner service ends, because the manager wants to count the whole pile first. A standard training step works this way: the backward pass produces one gradient per parameter, and all of them stay in GPU memory until the optimizer runs at the very end. In a Mixture-of-Experts model most plates are expert plates: in OLMoE-1B-7B the experts hold 6.44B of the 6.92B parameters, so their bf16 gradients alone take 12.9 GB. The gradient of the final MoE layer in the forward pass is computed first during backward, then waits while the backward pass walks back to layer one.

That wait is not needed. AdamW updates each tensor from that tensor's own gradient and state, so nothing forces the update to come late. Sonar-MoE washes each plate as it leaves the table: a hook fires when an expert tensor's gradient is complete, applies a fused AdamW update, and sets the gradient to None so its memory can be reused. The backward pass visits MoE layers one at a time, so at most one layer's expert gradients sit on the counter at once. The idea is not new for dense models; the paper's contribution is making it work for the batched expert tensors of an MoE feed-forward layer.

Backward reaches layer lexpert gradients computed

Hook firesgradient fully accumulated

Fused AdamW updatewhole batched tensor

Gradient set to Nonememory reusable

Walk back to layer l-1one layer resident

Two choices decide whether this pays off. The first is hook granularity. Attach one hook to the whole batched expert tensor, not one per expert. Expert kernels such as SonicMoE and ScatterMoE store a layer's experts in one tensor, and because AdamW is element-wise, one update on that tensor gives exactly the same result as many updates on its slices. Only the number of callbacks changes, and each autograd callback has a fixed cost. With per-expert hooks the paper measured +23.6 ms per step (+20%) on OLMoE and +61.1 ms (+72%) on Granite, at identical memory.

The second is clipping, and it is the price of the method. Global-norm clipping needs every gradient before any update, but the expert gradients have already been updated and freed by then. Sonar-MoE therefore leaves experts out of the clip, and keeps the normal end-of-step update, with clipping, for attention, embeddings and routers. Without expert clipping, the paper reports updates bitwise identical to standard training.

Worked example (OLMoE-1B-7B, SonicMoE backend; figures from the paper). The experts hold 6.44B parameters. At 2 bytes each in bf16, keeping every layer's gradient until the optimizer step costs 6.44B × 2 B ≈ 12.9 GB. OLMoE has 16 MoE layers, so freeing each layer's gradient right after its update leaves at most one layer resident: 12.9 GB ÷ 16 ≈ 0.8 GB. Measured peak allocated memory fell from 55.55 GB to 43.57 GB, a drop of about 12.0 GB, or 21.6% lower peak memory. The hook count follows the same arithmetic: per-expert hooks fire 2 × 64 experts × 16 layers = 2,048 times per step, batched hooks 2 × 16 = 32 times per step.

Method (OLMoE, SonicMoE)	Peak memory	Step time	Held-out loss	Gives up
Standard, global clip	55.55 GB	142.8 ms	1.2577	nothing
Standard, no expert clip	55.55 GB	115.9 ms	1.2591	clipping on experts
Optimizer-in-backward, all parameters	43.35 GB	114.4 ms	1.2612	all clipping
Sonar-MoE	43.57 GB	116.1 ms	1.2596	clipping on experts

Read the speed numbers with care. On OLMoE, standard training without expert clipping already runs in 115.9 ms against 116.1 ms for Sonar-MoE, so the speedup over the clipped baseline there comes from dropping expert clipping, not from the hook. The memory saving is the part Sonar-MoE adds. Gradient accumulation is the other caveat: the hook fires once per microbatch, so AdamW's running averages evolve differently from a single update on the averaged gradient, and the main experiments use one microbatch per step. The reported runs are 2,000-step Alpaca fine-tuning runs with batch size 2 and sequences up to 256 tokens, with weights, gradients and optimizer states all in bf16.

Goes deeper in: GPU & CUDA → Memory Hierarchy → HBM: Where Your Model Lives

Related explainers

TACO fine-tunes 32B models on one H100 — another way to shrink optimizer-side memory during fine-tuning.

Frequently Asked Questions

Check what you knowMap your AI & GPU knowledge across every track — free, role-based