[faresobeid]

Builds on #4427 (and so #4008) and on #4430 (native Muon, which removes dion), and lands after both: this branch is #4427's head and #4430's head merged, plus the commits here. With model.offload = "optimizer" and optim.type = "muon", the offloaded step now streams the Muon states through the GPU in chunks, the way #4008 does for AdamW. Before this, Muon took the serial path: all states to the GPU, step(), all states back.

To do that, the chunked path now runs on a small optimizer-agnostic interface. An optimizer opts in by updating a subset of its parameters at a time with the same math as a full step. torch AdamW gets an adapter. #4430's native Muon implements the interface itself, and subset_step_for() checks the optimizer for it before it looks for an adapter. There is no dion adapter, since #4430 removes dion.

Measured with dion's Muon. The profile and numbers below come from this PR's first version, before it was merged with #4430: dion's Muon behind a dion adapter. They have not been re-measured with the native Muon. The chunked step is bound by copy bandwidth, so the copy figures should carry over. The host-side costs in the forward-overlap section were dion's.

Why the Muon step was slow

Qwen3-30B-A3B on 8xH200 (FSDP 8 / EP 8) has 14.5 GiB of Muon state per rank: fp32 momentum, plus AdamW moments for the embeddings, LM head and norms. Main moves all of it to the GPU, steps, and moves it back, one phase after the other. Rank 0 trace, steps 5-7:

phase	no offload	main offload (serial)	this PR (chunked)
H2D state copies	—	558 ms (14.5 GiB at 26.0 GiB/s)	652 ms busy (22.2 GiB/s, runs at the same time as D2H)
D2H state copies	—	562 ms (14.5 GiB at 25.8 GiB/s)	551 ms busy (26.3 GiB/s)
Newton–Schulz	211 ms	209 ms	215 ms
Muon all-to-alls, dense mesh (18 batches × 2)	40 ms	15 ms	45-69 ms
Muon all-to-alls, expert mesh	none (the EP-8 expert mesh has size 1, so experts orthogonalize locally)	none	none
momentum + update (element-wise)	27 ms	27 ms	27 ms
non-Muon group (AdamW)	2 ms	2 ms	2 ms
GPU idle in the step	~35 ms	~50 ms	~8 ms
step (traced)	305 ms	1420 ms	873 ms

Each step moves 14.5 GiB per rank in each direction, 29 GiB in total. At #4426's measured per-GPU bandwidth, copying one direction at a time takes at least 0.56 s. Copying both directions at once (38.8 GiB/s combined) puts the floor at 0.75 s. The chunked step takes 0.85 s untraced. That is within about 0.1 s of the floor, the cost of loading the first chunk and writing back the last one.

Numbers

Qwen3-30B-A3B, 1 node 8xH200, FSDP 8 / EP 8, 16k seq, micro batch 1, global batch 8, full AC, compile, fake data, optim.type = "muon" (dion's Muon), no activation offloading. tok/s is the mean of the last 10 of 20 steps. Step time and optimizer step time are means/medians over steps 6-20 (synchronized at both ends). "NUMA" means bound to the GPU's socket, as in #4426.

	tok/s	step	optimizer step	peak GPU	host RSS / rank
no offload	51.5k	2.54 s	0.26 s	67.3 GiB	3.4 GiB
main, offload = "optimizer"	34.4k	3.81 s	1.44-1.48 s	52.8 GiB	22.7 GiB
main + #4426 (NUMA)	35.2k	3.72 s	1.38 s	52.8 GiB	22.7 GiB
this PR, NUMA (default config)	41.1-41.2k	3.18 s	0.85 s	52.8 GiB	22.7 GiB
this PR, NUMA, forward overlap forced on for Muon	40.1-40.2k	3.26 s	0.41 s returned + forward 1.15 s instead of 0.64 s	54.3 GiB	22.7 GiB

The chunked step removes 0.53 s of the 1.12 s that offload added to each step: +17% tok/s over main + #4426, +20% over main. The no-offload arm's step shows the remaining cost of the copies (0.64 s per step).

AdamW is unchanged by the refactor. With forward overlap (offload_overlap_forward = true, opt-in since #4427) and NUMA, this PR runs at 37.3k tok/s vs 37.5k for #4427 (3.52 s vs 3.49 s per step), within noise.

Forward overlap (#4427-style) for Muon: measured with dion, and left off

I ran the chunked dion Muon step alongside the next forward through the same interface, and it is bitwise correct. It is slower than the synchronous chunked step, so setup_optimizer keeps offload_overlap_forward AdamW-only. #4430's Muon has no host syncs in its step and batches its Newton–Schulz launches, but overlap with it has not been measured, so it stays off. From the dion trace:

• dion's step costs about 0.3 s of host time. That is about 5.8k Newton–Schulz kernels per step plus compile guards (no offload: 301 ms CPU for a 270 ms GPU step). Under overlap, the 8 chunks queued at step time take 0.41 s, and the rest runs on the forward's launch thread.
• dion syncs the host in the partitioned path. row_lr.unsqueeze(-1).to(X[0].device) is a pageable H2D copy, so the host waits until the chunk's load lands. That happened 14 times per forward, about 116 ms in total.
• The model syncs the host too. The MoE layers do a 128-byte device-to-host copy of token counts once per layer, about 520 ms of host blocking per forward. That is the compute stream waiting for the chunks. This also limits AdamW's overlap in perf(optim): overlap the offloaded AdamW step with the next forward (builds on #4008) #4427. Splitting the state copies into 32 MiB pieces made no difference (40.1k).

Change

• optim/subset_step.py adds the interface and the AdamW adapter (spec below).
• CPUOffloadOptimizer chunks the step for any optimizer for which subset_step_for finds one, not just isinstance(optimizer, AdamW). The chunk planner packs whole units instead of whole parameters, and streams every state tensor with ndim >= 1.
• Chunk boundaries use a rank-invariant size (global numel / number of shards). Uneven FSDP shards therefore can't make ranks cut different chunks and issue different collectives.
• Under forward overlap, units are sorted (stably, the same way on every rank) by the first FSDP group that needs them.
• Forward overlap requires subset-step support and stays AdamW-only in setup_optimizer. It is opt-in (model.offload_overlap_forward, default false, from perf(optim): overlap the offloaded AdamW step with the next forward (builds on #4008) #4427).

Design note

The interface

class SubsetStep(Protocol):
    units: list[list[torch.Tensor]]
    def step_params(self, params: Collection[torch.Tensor]) -> None: ...

class SupportsSubsetStep(Protocol):
    def begin_subset_step(self) -> SubsetStep: ...

• begin_subset_step() replaces step() for one step. It captures what the step reads that the trainer may change before the last chunk runs (group hyperparameters, gradients) and advances step counters.
• units partitions the parameters the step updates into the smallest sets that must update together. If the step communicates, units must be identical on every rank, decided from global shapes.
• step_params(params) updates exactly a union of whole units, bitwise as the full step would. It reads state from optimizer.state at call time (the wrapper has swapped it for GPU staging views), enqueues on the current stream without host syncs, and never creates state.

Implementations

• AdamW: one parameter per unit. Parameter groups are narrowed around one step() (GroupSubsetStep, perf(optim): overlap optimizer state offload copies with the AdamW step #4008's mechanism, moved here).
• Native Muon (feat(optim): native Muon, replacing dion #4430): implements SupportsSubsetStep itself, found by duck typing. begin_subset_step advances the groups' step counters and captures the hyperparameters and gradients. Its units are its buckets of FSDP-sharded matrices (one all-to-all each, the same on every rank), its batches of local matrices, and each AdamW group as a whole (one compiled update). step() is begin_subset_step() followed by step_params(all), so a chunked step runs the same code as a plain one. The step has no host syncs and uses its own NCCL communicators.

What stays outside the optimizer:

• The FSDP2 private-API wrapping (FSDPParamGroup.unshard / wait_for_unshard) belongs to the offload wrapper's forward overlap and does not depend on the optimizer.
• The MoE token-count sync per layer belongs to the model.

Validation

• Bitwise check (dion's Muon, before the feat(optim): native Muon, replacing dion #4430 merge): FSDP2 on 2 and on 8 GPUs, 6 steps of dion Muon, with parameters compared to un-offloaded Muon. Bitwise equal in every case, and the group step counters match.
  • Model: 6 blocks plus root, a muon group, a stacked 3-D "expert" group on the local (no-communication) path, an AdamW group and a Lion group.
  • Covers a fused QKV with matrix partitions, a 255-row matrix (uneven shards, and empty shards on 8 GPUs), an LR schedule stepped between steps, and grad clipping.
  • Paths: main's serial path, the chunked path with 1 GiB, 1 MiB and 64 KiB chunks (down to one batch per chunk), and the chunked path overlapped with the next forward.
  • The same holds for AdamW.
• New unit test: test_chunked_muon_matches_plain_muon (single GPU, native Muon, 1 GiB and 4 KiB chunks, step counters). On the current head, tests/unit/train/test_state_offload.py passes (8 tests, H200), with chunked native Muon bitwise equal to plain native Muon, and feat(optim): native Muon, replacing dion #4430's tests/unit/train/test_muon.py passes too (9 tests).
• Trainer runs above (dion's Muon): loss and grad norm per step stay within run-to-run noise.
  • Two runs with identical math (main with and without NUMA binding) differ by up to 1.5e-3 relative in loss and 2.7e-2 in grad norm, because the MoE kernels are not deterministic.
  • This PR vs main differs by 1.2e-3 to 2.4e-3 in loss and 2.2e-2 to 3.4e-2 in grad norm.
  • These runs use the default LR (1e-6), so parameter updates are tiny; the bitwise check above is the stronger evidence.
• ruff check and ruff format --check pass on the changed files.

🤖 Generated with Claude Code