[faresobeid]

Stacked on #4008.

With model.offload = "optimizer", Muon now uses #4008's chunked state offload instead of moving all of its momentum to the GPU at once. Each chunk loads on a side stream while the previous chunk is updated, and the chunk before that is written back.

• Chunk units are Muon's own batches. dion's create_param_batches gives batches of same-shape matrices, one per rank of the group's mesh. The AdamW group inside Muon is one unit, so it stays a single foreach update. Chunks are cut on the mean shard size (numel / mesh size), not on the local size. So every rank cuts at the same batches even with uneven shards, and Muon's all-to-alls match across ranks.
• Step counter. dion increments group["step"] once per step() call, and its AdamW bias correction reads it. Each chunk therefore starts from the same count, so the count advances by one per optimizer step, as before.
• AdamW is unchanged, except that its chunks are also cut on mean shard bytes.

Checks

• Qwen3-30B-A3B, 1 node of 8xH200, EP 8, offload = "optimizer", optimizer step only (random grads), steady state, max over ranks:

	step	step peak
main	1.38 s	43.8 GiB
this PR	0.83 s	31.3 GiB

• GLM-4.5-Air, 2 nodes, CP 4, EP 8, same measurement:

	step	step peak
main	2.77 s	76.9 GiB
this PR	1.52 s	53.1 GiB

• The same setup with 4 layers over 3 steps: the parameter update norms equal main's exactly.
• tests/unit/train/test_state_offload.py: 6 passed on an H200. The new test checks that Muon (matrices plus an AdamW group) with one unit per chunk is bitwise equal to plain Muon after 4 steps.
• ruff check, ruff format --check.

---

Note

Medium Risk
Touches distributed optimizer stepping and chunk synchronization for Muon collectives; incorrect batching or step counters could change training numerics, though a parity unit test mitigates this.

Overview
Extends CPU offload chunked stepping (previously AdamW-only) to dion Muon, so optimizer state is streamed in batches instead of loading all momentum onto GPU each step.

Chunking now groups step units rather than single parameters: Muon matrix groups are split with create_param_batches (same-shape batches per mesh rank); embedded AdamW groups stay one unit. Chunk boundaries use mean shard bytes (numel / mesh size) so every rank splits at the same units despite uneven local shards, keeping Muon's collectives aligned.

Before each chunk update, group["step"] is reset to the value at the start of the full step() call so dion's per-group counter still advances once per optimizer step across multiple chunks.

Adds test_chunked_muon_matches_plain_muon (forces tiny CHUNK_BYTES) to verify bitwise parity vs plain Muon for matrix + AdamW bias groups.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 4091e72. Configure here.}

[cursor]

Packed params can be updated twice

High Severity

Chunk units come from create_param_batches with _matrix_partitions, so one packed parameter can land in several same-shape matrix batches. Planning and _step_params still key on the physical tensor, so each chunk that contains that tensor runs a full optimizer.step() on it. Fused QKV or gate-up weights can be updated more than once per optimizer step once those batches no longer fit in a single CHUNK_BYTES window.

Additional Locations (2)

• src/prime_rl/trainer/optim/state_offload.py#L120-L138
• src/prime_rl/trainer/optim/state_offload.py#L143-L153

 

^{Reviewed by Cursor Bugbot for commit 4091e72. Configure here.}