[faresobeid]

optim.bf16_momentum = true (Muon only, default off) keeps the Muon momentum in bfloat16 instead of the parameter dtype (fp32). The AdamW group inside Muon is unchanged.

• BF16MomentumMuon creates the momentum in bfloat16. dion's update already casts the gradient to the momentum dtype and the momentum to bfloat16 before Newton-Schulz, so nothing else changes.
• torch.optim.Optimizer.load_state_dict casts floating-point state to the param dtype. After a load the momentum is cast back to bfloat16, so it stays bfloat16 across resume. Checked on CPU: with plain dion Muon it comes back fp32; with this class it stays bfloat16.
• Cost: the momentum is rounded to bfloat16 on every step (no stochastic rounding).

Measurements

GLM-4.5-Air, 2 nodes of 8xH200, CP 4, EP 8, optimizer step only (random grads), steady state, max over ranks:

	step	step peak	resident before step	host RSS/rank
offload = "none", fp32 momentum	0.63 s	76.9 GiB	74.9 GiB	-
offload = "none", bf16 momentum	0.62 s	64.2 GiB	62.7 GiB	-
offload = "optimizer", fp32 momentum	2.77 s	76.9 GiB	49.8 GiB	39 GB
offload = "optimizer", bf16 momentum	1.62 s	64.2 GiB	49.8 GiB	21 GB

Loss A/B: Qwen3-0.6B SFT on PrimeIntellect/Reverse-Text-SFT, 1 GPU, Muon lr 5e-5, batch 32, 150 steps. Mean train loss per window:

run	steps 11-31	steps 32-62
fp32, seed 0	2.2547	0.2857
fp32, seed 0, rerun	2.2538	0.2862
bf16, seed 0	2.2560	0.2849
fp32, seed 1	2.2949	0.3240
bf16, seed 1	2.2929	0.3287

The bf16 runs are within the spread between seeds. This is a small dense model on a short run. It is not evidence for long or large-MoE runs.

Checks

• tests/unit/orchestrator + tests/unit/test_configs.py -m "not gpu": 350 passed. ruff check, ruff format --check.

---

Note

Medium Risk
Changes optimizer state dtype and checkpoint resume behavior for Muon; training dynamics may differ on long or large runs despite short A/B checks.

Overview
Adds an opt-in optim.bf16_momentum flag on Muon (default off) to store Muon momentum in bfloat16 instead of the parameter dtype, cutting optimizer state size and CPU offload traffic at the cost of lower momentum precision.

When enabled, optimizer setup uses a new BF16MomentumMuon subclass that initializes Muon momentum in bf16 and re-casts momentum to bf16 after load_state_dict, so checkpoint resume does not leave momentum in fp32. The embedded AdamW/Lion path for non-Muon params is unchanged.

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