Muon under AutoTP orthogonalizes each shard, so the update depends on the TP degree

[alanhuangyoo]

What happens

With autotp_size > 1, Muon's Newton-Schulz runs on each rank's column-parallel shard of an attention projection. Orthogonalizing a block of rows is not the same computation as taking that block out of the orthogonalization of all the rows, so the update Muon applies under AutoTP is not the update Muon applies without it.

This is independent of #8384; it is the current behaviour of the whole-matrix path.

Measurement

Llama, hidden_size=256, 8 heads of 32, 1 layer, fp32, ZeRO-1, Muon (lr=0.05, momentum=0.95), same seed and the same batch on both sides, one step, model.layers.0.self_attn.q_proj.weight. The tp=2 shards are concatenated in rank order to compare against the tp=1 matrix.

initial weights identical:  True
loss:                       4.874537 both

gradient   tp=1 vs tp=2   relative difference 7.0e-07     (fp32 all-reduce ordering)
update     tp=1 vs tp=2   relative difference 3.9e-01     cosine 0.936
update norm                tp=1 0.0145      tp=2 0.0161   (+11%)

The gradients agree to 7e-7. The updates do not.

Isolating the kernel from the training loop, on one random 256 x 256 gradient, 5 steps:

whole matrix, tp=2, rank 0  rel diff to the same rows of the unsharded result = 0.5273
whole matrix, tp=2, rank 1                                                      0.5234

Both Newton-Schulz kernels behave the same way here; gram and standard give 0.53 and 0.53.

Why it happens

zeropower_via_newtonschulz5 and zeropower_via_gram_newtonschulz normalize by the matrix norm and then iterate on X @ X.mT. Both depend on every row of the input. A column-parallel shard is a strict subset of the rows, so the iteration converges to the orthogonalization of the shard, which is a different matrix from the corresponding rows of the orthogonalization of the whole.

What could be done

Three options, none free, listed in increasing cost:

1. Document it. Say that Muon under AutoTP orthogonalizes per shard, and that the update is therefore tp-dependent. Cheapest, and at least stops it being a surprise.
2. Gather before orthogonalizing. Restores the semantics exactly, at the cost of an all-gather of the projection and NS on the full matrix on every rank — the redundancy AutoTP exists to avoid.
3. Split on an axis the orthogonalization respects. This is what per-head Newton-Schulz does. A column-parallel shard holds whole attention heads, and per-head NS batches over exactly that axis, so per-head on a shard is bit-identical to the corresponding blocks of per-head on the whole matrix — measured as 0.00e+00 for both kernels at tp=2 and tp=4. That only covers head-blocked projections; the MLP matrices in the same model are still split on an axis NS does not respect.

I have (3) implemented and tested in #8384 as a consequence of the per-head split, and it is what led me here. It is not a general answer, because it says nothing about gate_proj/up_proj/down_proj, which are the larger share of the Muon parameters.

Happy to take whichever of these is wanted, or to write up (1) on its own if the semantics are considered acceptable and just undocumented.

Reproduction

# same script run under `deepspeed --num_gpus 1` and `--num_gpus 2`, the second with
# "tensor_parallel": {"autotp_size": 2} added to the config
cfg = transformers.LlamaConfig(hidden_size=256, num_attention_heads=8, num_key_value_heads=8,
                               num_hidden_layers=1, intermediate_size=512, vocab_size=128)
model = transformers.AutoModelForCausalLM.from_config(cfg).float()
# ... deepspeed.initialize with {"optimizer": {"type": "Muon", "params": {"lr": 0.05}}}
# one step on an identical batch, then compare (w_after - w_before) across the two runs

[alanhuangyoo]

Quantifying what option 3 can and cannot reach, since it changes how much of the problem it solves.

Newton-Schulz on a shard versus the corresponding block of Newton-Schulz on the whole matrix, for a 512 x 1024 gradient, 5 steps (a different shape from the 256 x 256 in the issue body, so the magnitudes differ; the comparison between the two columns is the point):

kernel	tp	dim-0 split (column-parallel: q/k/v)	dim-1 split (row-parallel: o_proj, down_proj)
gram	2	0.3435	0.6831
gram	4	0.3821	0.6646
standard	2	0.3418	0.6758
standard	4	0.3809	0.6602

Row-parallel is about twice as far off as column-parallel, and per-head does nothing for it. The head axis is dim 0, so a per-head split lines up with a column-parallel shard exactly — that is why it comes out bit-identical. A row-parallel shard cuts dim 1, which is the input axis: o_proj is [hidden, num_heads * head_dim] with its heads on the input side, and gate_proj / up_proj / down_proj have no head structure at all.

So for an AutoTP'd transformer, option 3 makes q/k/v exact and leaves the other four matrices per layer where they are — and those four are the larger share of both the parameters and the Newton-Schulz cost. Worth knowing before treating per-head as the answer to this issue rather than as one part of it.

Options 1 and 2 in the issue body still cover everything; they just cost a gather or the documented semantics.