ZeRO-3 applies Muon's Newton-Schulz once per micro-batch, so gradient accumulation changes the optimizer

[alanhuangyoo]

Summary

Under ZeRO-3, Muon's Newton-Schulz runs once per micro-batch instead of once per optimizer step. With gradient_accumulation_steps: n the momentum buffer is advanced n times per step and the orthogonalization is applied to partial gradients rather than to the accumulated one. ZeRO-1 and ZeRO-2 are correct.

Measurement

Two Linear layers (so two Muon matrices), two optimizer steps, counting calls into the Newton-Schulz kernels in original_muon. The correct count is 2 steps x 2 matrices = 4 in every row.

stage	gas	Newton-Schulz calls	expected
1	1	4	4
1	4	4	4
2	1	4	4
2	4	4	4
3	1	4	4
3	4	16	4

Stage 3 with gas=4 does 4x the work, one orthogonalization per micro-batch.

That is not only wasted compute. Same data, same seed, three optimizer steps, momentum=0.95, gradient_clipping: 0, and the same 8 samples per optimizer step either way — one micro-batch of 8 with gas=1, four of 2 with gas=4. The mean gradient is identical by construction, so the two runs must agree:

stage 2:  gas=1 vs gas=4   relative difference in the weights = 3.3e-04
stage 3:  gas=1 vs gas=4   relative difference in the weights = 1.0e-01

3.3e-04 is the half-precision Newton-Schulz noise floor — stage 2 is the same training run either way. Stage 3 is 300x that: a different one.

Why

stage3.py:

    def _apply_distributed_muon_update(self, communication_data_type, buffer_to_reduce):
        if not self.use_muon:
            return
        ...

and its only call site is inside the IPG bucket reduce path:

        dist.all_reduce(buffer_to_reduce, group=process_group)
        ...
        self._apply_distributed_muon_update(communication_data_type, buffer_to_reduce)
        for param in params_in_bucket:
            grad = param.grad

That path runs on every micro-batch. There is no accumulation-boundary guard, and use_muon is the only condition. ZeRO-1/2 does the same work in get_flat_partition, which ipg_epilogue calls only under if self.is_gradient_accumulation_boundary(): — which is why those two stages come out right.

The consequences, in order of how much they matter:

1. The momentum decay is wrong. momentum.lerp_(grad, 1 - beta) runs n times per optimizer step, so the effective retention is beta**n rather than beta. At beta=0.95 and gas=4 that is 0.81, and at gas=16 it is 0.44 — the momentum the user configured is not the momentum they get, and the discrepancy moves with an unrelated knob.
2. The orthogonalization sees partial gradients. Newton-Schulz is not linear, so a sum of orthogonalized micro-batch gradients is not the orthogonalization of their sum. Since Muon's update is scale-invariant, each micro-batch contributes a unit-scale update regardless of how few samples it saw, which weights noisy micro-batches equally with the rest.
3. n times the Newton-Schulz cost, which for large matrices is the expensive part of the step.

ZeRO-3 with gradient accumulation is the standard configuration for a model large enough to need ZeRO-3, so this is the common path rather than a corner.

What the fix needs

is_gradient_accumulation_boundary() already exists on the class, but guarding the call with it is not enough on its own: the current code takes the full-shape gradient out of buffer_to_reduce inside the reduce path, and on the boundary micro-step that buffer holds only that micro-step's contribution. Applying the update once, at the boundary, on the accumulated gradient means reading the accumulated partition instead — which is what ZeRO-1/2 does in get_flat_partition.

Happy to implement it if the shape of that is agreed. Flagging first because the change is inside the ZeRO-3 reduce path and is more than a guard.

Reproduction

import torch, deepspeed
import deepspeed.runtime.zero.muon.original_muon as om

NS = {"n": 0}
for name in ("zeropower_via_gram_newtonschulz", "zeropower_via_newtonschulz5"):
    f = getattr(om, name)
    setattr(om, name, (lambda f: (lambda *A, **K: (NS.__setitem__("n", NS["n"] + 1), f(*A, **K))[1]))(f))

model = torch.nn.Sequential(torch.nn.Linear(64, 64, bias=False), torch.nn.Linear(64, 64, bias=False))
GAS = 4
engine, _, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(),
    config={"train_micro_batch_size_per_gpu": 2, "gradient_accumulation_steps": GAS,
            "gradient_clipping": 0.0,
            "zero_optimization": {"stage": 3, "reduce_scatter": False},
            "optimizer": {"type": "Muon", "params": {"lr": 0.02}}})

for _ in range(2 * GAS):                      # two optimizer steps
    x = torch.randn(2, 64, device=engine.device)
    engine.backward(engine(x).square().sum())
    engine.step()
print("newton_schulz calls:", NS["n"])        # 16 on stage 3, 4 on stages 1 and 2

[alanhuangyoo]

A note on what the fix would cost, because I assumed it would add communication and it does the opposite.

_apply_distributed_muon_update already gathers the momentum every time it runs — _partitioned_buffers_all_gather(params, momentum_buffer, ...) plus the per-parameter dist.all_gather of gradients and momentums in the round-robin below it. Since it runs per micro-batch, at gas=n those collectives happen n times per optimizer step today. Moving the update to the accumulation boundary makes it once, so the change reduces ZeRO-3 Muon's communication by a factor of gas, on top of the gasx Newton-Schulz saving.

What it needs in exchange: Newton-Schulz is over the whole matrix, so the update needs the full-shape gradient. The current code gets that for free because it sits in the reduce path where buffer_to_reduce still holds unpartitioned gradients. At the boundary the accumulated gradient lives in the partitions, so it has to be gathered — which is the same collective the momentum already uses, so the machinery is there.

Sketch, for whoever picks this up:

1. Guard the call site with self.is_gradient_accumulation_boundary() so non-boundary micro-steps only reduce and accumulate, as they do for every other optimizer.
2. On the boundary, gather the accumulated gradient partitions for the Muon sub-groups the same way the momentum is gathered, run muon_update on the full-shape result, and write it back into the partitions.

Step 2 is the real work, and it is why I would rather not guess at the shape. Happy to write it if that is the direction, or to review someone else's.

One more thing worth checking whoever does it: with the update at the boundary, beta finally means what it says, so any tuned ZeRO-3 Muon run will need its momentum revisited — at gas=4 the effective retention today is 0.95**4 = 0.81. That is a behaviour change, but in the direction of the configured value.

[alanhuangyoo]

Reproduced with different seeds and a different width, since a single measurement of a divergence is worth confirming.

Three optimizer steps, same data per step either way, gradient_clipping: 0, momentum=0.95:

	stage 2	stage 3
seed 0, dim 64 (issue body)	3.3e-04	1.0e-01
seed 11, dim 128	6.660e-04	6.720e-02
seed 22, dim 128	6.663e-04	6.641e-02

Stage 2 sits at its half-precision floor and barely moves with the seed — the two runs are the same training. Stage 3 is two orders of magnitude above that in every case. The separation is the finding, not the exact figure.