[BUG] ChainedOptimizer applies global grad-norm clipping to Muon (orthogonalizing) param groups, silently stalling training

[yuchenwang3]

Describe the bug

ChainedOptimizer.step() computes a single global grad_norm across all chained sub-optimizers and applies clip_grad_by_total_norm_fp32(..., total_norm=grad_norm) to every sub-optimizer's parameters, including a Muon / orthogonalizing group (dist_muon). For an orthogonalizing optimizer this is (a) meaningless — Muon discards gradient magnitude by design (Newton–Schulz orthogonalization is scale-invariant) — and (b) harmful: when the global grad_norm is large, the clip coefficient c = clip_grad/grad_norm is tiny and scales the per-matrix gradients fed into Newton–Schulz below its internal F.normalize(eps=1e-7) floor, so orthogonalization degenerates (near-zero, non-orthogonal update). Affected layers stop updating, loss stalls, grad_norm grows, c shrinks further, more layers collapse — a silent positive-feedback stall.

Steps/Code to reproduce bug

We hit this fine-tuning a GatedDeltaNet model whose gradients are inherently large and per-layer-imbalanced (grad_norm ≈ 5e7, growing to ~2e11). Minimal mechanism repro (torch-only) showing the clip coefficient pushing a per-matrix gradient below the Newton–Schulz floor:

import torch, torch.nn.functional as F
clip_grad, grad_norm = 1.0, 5e7
c = clip_grad / (grad_norm + 1e-6)                 # ≈ 2e-8

def newton_schulz(x, steps=6, eps=1e-7):
    a, b, cc = 3.4445, -4.7750, 2.0315
    X = F.normalize(x, p=2, dim=(-2, -1), eps=eps)
    for _ in range(steps):
        A = X @ X.transpose(-2, -1); X = a*X + (b*A + cc*(A@A)) @ X
    return X

G = torch.randn(2048, 2048); G = G / G.norm() * 0.1     # a late-layer matrix grad
print("no clip :", newton_schulz(G).norm().item())       # ~40  (orthogonal)
print("clipped :", newton_schulz(G * c).norm().item())   # collapses toward 0

Training-level, single-variable (fixed 32-example overfit batch, optimizer=dist_muon, only clip_grad changed):

clip_grad=1.0 : loss 0.596 0.671 ... 0.583   stalls ~0.5, never overfits; grad_norm 7.5e7 -> 2.2e11
clip_grad=0   : loss 0.596 0.530 ... 0.019   clean overfit

(Keeping clip_grad=1.0 and instead only lowering Newton–Schulz eps 1e-7→1e-30 also unblocks learning, confirming the clip → NS-eps interaction.)

Expected behavior

Orthogonalizing / Muon param groups should not be subject to magnitude-based global gradient clipping (it is a no-op-at-best for scale-invariant updates). Training should proceed (the fixed-batch overfit should drive loss → 0), as it does with clip_grad=0.

Additional context

• Location: megatron/core/optimizer/optimizer.py, ChainedOptimizer.step() — grad_norm = self.get_grad_norm() then per sub-optimizer clip_grad_by_total_norm_fp32(..., total_norm=grad_norm).
• Suggested fix: skip global norm-clipping for orthogonalizing/Muon groups; or default clip_grad=0 for dist_muon; or clamp the clip coefficient with a floor so it cannot push grads under the optimizer's numerical eps.
• Related: the silent Newton–Schulz degeneration that makes this hard to diagnose is filed on NVIDIA-NeMo/Emerging-Optimizers — issue newton_schulz silently degenerates (non-orthogonal output) for small-norm inputs due to F.normalize eps floor NVIDIA-NeMo/Emerging-Optimizers#229 and PR fix(muon_utils): keep newton_schulz scale-invariant for small-norm inputs (#229) NVIDIA-NeMo/Emerging-Optimizers#230. The two interact: the clip coefficient pushes Muon's per-matrix gradients below Newton–Schulz's eps floor, so the orthogonalization silently degenerates and the forward/loss look completely normal.
• Environment: Megatron-LM 0.19, Emerging-Optimizers 0.4.0a0, PyTorch 2.11+cu128, B200.

[yuchenwang3]

Background: hit while running ms-swift + Megatron-Core SFT of Qwen3.5-35B-A3B (GatedDeltaNet hybrid, optimizer=dist_muon) on 16× B200 across 2 nodes (2×8).

Clarification on scope (it's not Muon-specific): the root cause is optimizer-agnostic — a benign-but-large global grad_norm makes the clip coefficient tiny (≈2e-8), which pushes per-matrix gradients below the optimizer's numerical floor and freezes updates. AdamW under the same clip also stalls (its eps floor), so the clip — not the optimizer — is the cause. (We used dist_muon here following the model's training recipe; GatedDeltaNet itself is not tied to Muon.)

The proposed fix #5395 targets the orthogonalizing-optimizer case because magnitude clipping is semantically meaningless there (Newton–Schulz discards magnitude) → a clean, side-effect-free skip. There's no equally clean "skip clip" for a magnitude-based optimizer; the optimizer-agnostic alternative is to clamp the clip coefficient so a benign-large grad_norm can't drive gradients under the eps floor (larger blast radius). Fix PR: #5395.

[yuchenwang3]

Update — optimizer-agnostic confirmation. Controlled follow-up on a fresh 8× B200 node (Qwen3.5-35B-A3B, fixed 64-example overfit, 200 steps):

optimizer	clip_grad	result
adam	1.0	stalls ~0.5
adam	0	overfits → 0.029
dist_muon	0	overfits → 0.019

grad_norm stayed ~1e8–3e12 for the whole run with no NaN — AdamW's per-parameter second-moment normalization handles the (architecturally large) raw gradients fine once they are not crushed by the clip. With clip_grad=1.0, c = clip_grad/grad_norm ≈ 1e-12 annihilates the update for any optimizer; with clip_grad=0 both AdamW and Muon learn.

So the clip is the sole cause and this is not Muon-specific — AdamW is viable without the clip; an orthogonalizing optimizer is simply naturally immune (magnitude clipping is a no-op for it). Proposed fix: #5395.

[Saurav-Gupta-9741]

Hi @yuchenwang3 Sir, I would like to tackle this bug.

My plan is to modify the ChainedOptimizer logic so that param groups explicitly designated for the Muon optimizer (or any orthogonalizing optimizer) bypass the standard global grad-norm clipping phase. I will ensure that the gradient norms for these specific groups are excluded from the total norm calculation to prevent the silent training stall, while still allowing standard parameters to clip correctly.

I will begin implementing this and aim to open a PR within the next 24-48 hours. Please let me know if there are any specific edge cases with the pipeline-parallel integration you'd like me to watch out for!

[yuchenwang3]

Thanks for the interest @Saurav-Gupta-9741! I've already opened #5395 for this — it adds a per-optimizer skip_grad_norm_clip attribute so the orthogonalizing (Muon) param groups bypass the global grad-norm clip in ChainedOptimizer.step. (A config-level fix doesn't work because ChainedOptimizer asserts all sub-optimizers share one config, hence the per-optimizer attribute.) Review/feedback there very welcome.

One nuance worth flagging so it isn't missed: the stall is actually optimizer-agnostic — plain AdamW hits the same wall once the global clip coefficient collapses (its update is then crushed below the adam_eps floor), it's not specific to Muon. The orthogonalizing optimizers just make it more visible. So the real fix is decoupling the clip for these groups, which #5395 does.