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

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

Start with stage3.py, tracing the IPG bucket reduce path and _apply_distributed_muon_update call, then compare it with get_flat_partition and the ZeRO-1/2 boundary handling. Run the provided reproduction and inspect the original_muon kernel call counts. Done means ZeRO-3 applies Muon's update once per optimizer step using the accumulated gradient, with gas=1 and gas=4 producing matching results and four calls in the example.