huggingface/pytorch-image-models · error · RuntimeError

Muon does not support sparse gradients

Error message

Muon does not support sparse gradients

What it means

Muon's update relies on dense matrix operations (Newton–Schulz orthogonalization of the gradient), so step() raises immediately if any parameter's gradient is a sparse tensor.

Source

Thrown at timm/optim/muon.py:825

            muon_params = []
            muon_grads = []
            muon_momentum_bufs = []
            # Additional state for adamuon mode
            muon_exp_avg_sqs = []
            muon_state_steps = []

            adamw_params = []
            adamw_grads = []
            adamw_exp_avgs = []
            adamw_exp_avg_sqs = []
            adamw_state_steps = []

            for p in group["params"]:
                if p.grad is None:
                    continue

                if p.grad.is_sparse:
                    raise RuntimeError("Muon does not support sparse gradients")

                state = self.state[p]

                # Determine routing on first encounter (cache in state)
                if "use_muon" not in state:
                    # Check explicit flags first (support both 'use_fallback' and 'use_muon' for compatibility)
                    reason = None
                    if group.get("use_fallback", False):
                        # use_fallback=True means use AdamW (use_muon=False)
                        state["use_muon"] = False
                        if verbose:
                            reason = "use_fallback_flag"
                    elif "use_muon" in group:
                        # Explicit use_muon flag for compatibility with other Muon implementations
                        state["use_muon"] = group["use_muon"]
                        if verbose:
                            reason = "use_muon_flag"
                    else:

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Remove sparse=True from embeddings so gradients are dense
2. Put sparse-gradient params in a separate param group optimized by torch.optim.SparseAdam or SGD
3. Exclude embeddings from the Muon optimizer entirely

Example fix

# before
emb = nn.Embedding(vocab, dim, sparse=True)
opt = Muon(model.parameters())
# after
emb = nn.Embedding(vocab, dim)
opt = Muon(model.parameters())

assert all(p.grad is None or not p.grad.is_sparse for p in params), 'Muon cannot step on sparse gradients'

def has_sparse_grads(params) -> bool:
    return any(p.grad is not None and p.grad.is_sparse for p in params)

When it happens

Trigger: optimizer.step() with a parameter whose .grad is sparse — typically nn.Embedding(sparse=True) — while using the Muon optimizer.

Common situations: Applying Muon to a whole model that includes sparse embeddings (language models, recsys); reusing a sparse training pipeline with a new Muon config.

Related errors

• AdamW does not support sparse gradients
• ADOPT does not support sparse gradients
• momentum != 0 is not compatible with sparse gradients
• weight_decay option is not compatible with sparse gradients
• Adam does not support sparse gradients, please consider Spar

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/291c7434ba061d13. Report an issue: GitHub.