[faresobeid]

Summary

model.offload = "full" now supports Muon, so Muon no longer needs more trainer nodes than AdamW or SignSGD. GLM-4.5-Air Muon now trains on one trainer node.

How it works:

• CPU state. FP32 masters, accumulated gradients, and FP32 momentum stay in CPU RAM. The layout is Dion's (momentum; momentum + variance for Dion's AdamW group; the step count lives on the group), so checkpoints keep Dion's state keys.
• During backward. Muon param groups get their own CPU chunks. As each chunk's gradients arrive, the CPU computes M = mu * M + scale * G. Dion's AdamW group (embeddings, lm_head, norms) runs on the existing native CPU AdamW kernel.
• After backward. step() walks Dion's batches of same-shaped matrices with two batches in flight. For each batch:
  1. BF16 momentum goes H2D.
  2. Dion's muon_update_batch_async runs on GPU. This is the same Newton-Schulz path as offload=none, including local NS for EP-sharded experts and the all-to-all for FSDP-sharded matrices. It gets a zero gradient, unit momentum, and a zero parameter, so it returns exactly -adjusted_lr * NS(M).
  3. The update goes D2H as BF16.
  4. The CPU applies weight decay and the update to the FP32 master (native add_bfloat16_).
  5. The new BF16 compute weights go H2D.
• Matrix partitions (packed qkv / gate_up) are taken from the model's parameters, because the optimizer owns CPU master shards.
• Removed: the config validators that rejected Muon under full offload.
• Not supported: Nesterov raises under full offload.

Results

Correctness: Qwen3-30B-A3B, 2 layers, EP8, 1 node, lr 1e-3

Each row is the relative error of the cumulative update against offload=none after steps 1 / 2 / 3.

comparison	step 1	step 2	step 3
Muon full vs Muon none	1.03%	5.5%	8.0%
Muon full vs Muon none, wd 0.1	1.02%	5.5%	8.1%
AdamW full vs AdamW none (existing path, reference)	0.71%	8.8%	11.3%
Muon full vs Dion on GPU, same BF16 grads and state (per step)	0.055%	0.076%	0.31%

• The full-vs-none gap comes from full offload's BF16 gradients. It is the same size as AdamW's.
• Against Dion on GPU with identical inputs, the gap is about 0.17% per matrix. That matches BF16 rounding of the transferred update.
• Checkpoint save at step 2 and resume to step 4 works.

GLM-4.5-Air: 1 node vs 2 nodes

Settings: real weights, fake data, swe-budget.toml trainer settings (cp 4 ulysses, AC + AC offload, compile, fused LM head 1024), 8 x 131k-token samples per step.

config	nodes	step time	GPU peak (rank 0)	node RAM used
SignSGD, offload=full	1	208.1 s	71.8 GiB	947 GiB
Muon, offload=full	1	218.7 s	72.8 GiB	1426 / 1574 GiB
Muon, offload=optimizer (swe.toml trainer)	2	109.0 s	91.7 GiB	475 GiB per node

• The 2-node row uses the default fused LM head chunk size, as in swe.toml.
• Muon full is 5% slower than SignSGD full. The extra time is the post-backward Muon phase, about 9.4 s per step, which is not overlapped with backward.
• Muon full on 1 node costs the same GPU-seconds per step as Muon on 2 nodes (1750 vs 1744) and uses half the nodes.

---

Note

High Risk
Changes the distributed full-offload optimizer step and Muon’s GPU Newton–Schulz path; mistakes affect training correctness, checkpoints, and multi-node collectives.

Overview
model.offload = "full" now works with Muon, not just AdamW and SignSGD, so Muon training can fit on fewer trainer nodes when optimizer state and gradients live on CPU.

FullCPUOffloadOptimizer gains a Muon path: CPU chunks are split per param group; during backward Muon chunks only fold gradients into FP32 momentum on CPU; Dion’s AdamW subgroup still uses the existing native CPU AdamW kernel with Dion’s state keys. After backward, step() runs batched GPU Newton–Schulz via Dion’s muon_update_batch_async (BF16 momentum H2D, orthogonalized update D2H, FP32 master + compute weight refresh), with matrix partitions passed from the model for fused packed weights. Nesterov Muon is rejected under full offload.

Config/runtime guards that raised on Muon + full offload are removed (setup_optimizer, SFT/Trainer validators); docs and the start-run skill document Muon support and host RAM per parameter. The unit test expecting that validation error is deleted.

^{Reviewed by Cursor Bugbot for commit 55be9e7. Bugbot is set up for automated code reviews on this repo. Configure here.}