deepspeedai / deepspeedai/DeepSpeed

Default gradient_clipping divides every Muon update by its own norm, shrinking the step by a model-sized factor

Open
#8,439 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Summary

gradient_clipping defaults to 1.0 (GRADIENT_CLIPPING_DEFAULT), and under Muon the quantity it clips is not the gradient — it is the orthogonalized update, whose norm is a property of the model's shape rather than of the gradient. Every Muon step is therefore divided by a constant that grows with the model, and the resulting step size is pinned by the clip threshold instead of by the learning rate. Nothing reports it.

Measurement

Four square Linear layers, one step, ZeRO-1, fp32, lr=0.02, same gradient in every row. Only gradient_clipping differs: 0 disables it, 1.0 is what a config that does not mention it gets.

hidden reported grad_norm |Δw| with clipping off |Δw| at the default 1.0 attenuation
128 4.2182 0.038865 0.009214 4.22x
512 7.1564 0.058267 0.008142 7.16x
2048 22.7588 0.174972 0.007688 22.76x
4096 44.2460 0.339524 0.007674 44.25x
8192 85.6024 0.676925 0.007908 85.60x

Three things to note:

  1. The attenuation equals the reported norm, to three digits, at every size. The clip is dividing the update by its own norm.
  2. Unclipped, the step grows with the model, as Muon intends. Clipped, it is flat at ~0.0077 across a 64x range of model width — the step size is set by the clip threshold, not by Muon or by lr.
  3. The factor grows with the model: 4x on a toy, 86x by hidden 8192, on a 4-layer MLP. A real model has far more matrices.

Identical numbers on ZeRO-1 and ZeRO-2. ZeRO-3 reports the same norm; I did not measure |Δw| there because the parameter is partitioned.

Why it happens

Muon's update is scale-invariant in the gradient. Newton-Schulz normalizes its input and returns a near-orthogonal matrix, so the output depends on the gradient's direction and not its magnitude:

gradient scaled by 1      -> update norm 15.023504
gradient scaled by 0.001  -> update norm 15.024741
gradient scaled by 1000   -> update norm 15.029729
relative difference between the x1 and x1000 updates: 0.0055   (half-precision NS noise)
sqrt(min(n, m)) for this 256 x 512 matrix = 16.0

So the update's norm is set by the matrix, not by the gradient — and clipping the gradient, which is what gradient_clipping names, provably cannot change a Muon step.

DeepSpeed does not clip the gradient. get_flat_partition writes the post-Newton-Schulz update into averaged_gradients (stage_1_and_2.py, the muon_update call), which is how Muon is threaded through the ZeRO pipeline at all. scaled_global_norm then takes the norm of that, and unscale_and_clip_grads divides by it:

clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad
clip = torch.clamp(clip, min=1.0)

total_norm is the norm of the orthogonalized update, so for any model where it exceeds clip_grad, clip is essentially total_norm and the whole update is renormalized to a fixed global norm every step. Muon's per-matrix spectral scaling — the thing the orthogonalization exists to produce — is replaced by one global rescale.

Two consequences worth separating

The step size. A user who never writes gradient_clipping in their config gets an effective learning rate divided by a model-dependent constant. That is a plausible mechanism behind "Muon performs worse than AdamW" reports; #7713 has a chart of exactly that shape, with AdamW unaffected in the same run. AdamW is unaffected because for it averaged_gradients really does hold gradients, so its clipping does what it says.

The reported metric. grad_norm under Muon is the norm of the update, not of the gradient. On the same model and step:

Adam  reported_norm = 0.0298
Muon  reported_norm = 4.2182

Anyone watching grad_norm to judge training health is reading a number with a different meaning under Muon, and one that barely moves because it is dominated by the matrix ranks.

Options

  1. Exclude the Muon groups from clipping, and clip only the Adam half. This matches what gradient clipping means: since Muon is scale-invariant, clipping its gradient is a no-op, so not clipping it is the faithful implementation rather than a special case. It also takes the Muon half out of the reported norm, which makes the metric mean what it says again.
  2. Clip before Newton-Schulz. Faithful, and provably a no-op for the Muon half by the scale-invariance above — so it is option 1 with more machinery.
  3. Warn. If the current behaviour is considered acceptable, it should at least say at deepspeed.initialize that gradient_clipping will rescale Muon updates by the update norm, since the default value alone triggers it.

Whichever is chosen, the default is the part that makes this urgent: it fires without anyone opting in.

I am happy to implement (1) with tests. Because it changes the step size for every existing Muon run that did not set gradient_clipping: 0, I would rather have a maintainer pick the option than choose one myself.

Reproduction

# one step, read |Δw| of one layer with gradient_clipping 0 and then 1.0
cfg = {"train_micro_batch_size_per_gpu": 4, "gradient_accumulation_steps": 1,
       "zero_optimization": {"stage": 1},
       "gradient_clipping": 0.0,          # then 1.0
       "optimizer": {"type": "Muon", "params": {"lr": 0.02}}}

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in stage_1_and_2.py and trace get_flat_partition, the muon_update call, scaled_global_norm, and unscale_and_clip_grads. Reproduce the one-step configuration with gradient_clipping set to 0.0 and 1.0, then confirm the maintainer-selected behavior with regression tests covering Muon and Adam clipping.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.