Muon cannot run under BF16_Optimizer: flat partitions with no Newton-Schulz applied anywhere

[alanhuangyoo]

Split out of #8442 at @delock's suggestion, so the refusal that PR adds has somewhere to be lifted from.

What is refused

{"bf16": {"enabled": true}, "data_types": {"grad_accum_dtype": "fp32"}, "zero_optimization": {"stage": 1}} with Muon now raises ZeRORuntimeException instead of training. That combination selects BF16_Optimizer, and Muon cannot run under it.

Why

MuonWithAuxAdam.step decides what it has been handed by shape:

if p.dim() < 2:
    # A flat ZeRO partition -- the gradient already holds the update
    update = p.grad
else:
    # The weight itself, so nothing has orthogonalized it
    update = muon_update(p.grad, state["momentum_buffer"], ...)

BF16_Optimizer replaces the param groups with flat fp32 partitions, exactly like the ZeRO 1/2 and 3 optimizers, so every parameter reads as 1-D and step takes the first branch. But unlike them it never applies Muon anywhere: DeepSpeedZeroOptimizer and _Stage3 orthogonalize while the partition is filled, and BF16_Optimizer has no such call. So the first branch's premise -- "ZeRO already did the update" -- is false, and the step is plain SGD with momentum.

Nothing in the flat partition records the original 2-D shapes, so step cannot recover them to orthogonalize there.

What lifting it would take

Either:

1. Apply Muon in BF16_Optimizer where the partition is filled, matching what DeepSpeedZeroOptimizer.get_flat_partition does -- it has the parameter list, so the shapes are reachable at that point; or
2. Carry the shapes onto the flat partition so step can orthogonalize per original tensor, which is a wider change and duplicates work the ZeRO optimizers already do in the right place.

(1) looks like the right shape to me, but I have not written it and would rather not guess at the ordering constraints inside BF16_Optimizer's accumulation window.

How it was found

Re-reading my own claim in #8442. That PR's description said Muon runs correctly under "every wrapper", which was generalised from four configurations I had measured; BF16_Optimizer is a fifth and behaves differently. I posted the hole on the PR before it was reviewed rather than letting it stand.

Test pinning the refusal, and the neighbouring config that still works, is in tests/unit/runtime/zero/test_muon_without_zero_optimizer.py::TestMuonRefusesBF16Optimizer.