huggingface/pytorch-image-models · error · RuntimeError

NAdamW does not support sparse gradients

Error message

NAdamW does not support sparse gradients

What it means

RuntimeError raised in NAdamW.step() when a parameter's gradient is a torch sparse tensor. The NAdamW update kernels only handle dense gradients.

Source

Thrown at timm/optim/nadamw.py:125

        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad = []
            grads = []
            exp_avgs = []
            exp_avg_sqs = []
            state_steps = []
            beta1, beta2 = group['betas']

            for p in group['params']:
                if p.grad is None:
                    continue
                params_with_grad.append(p)
                if p.grad.is_sparse:
                    raise RuntimeError('NAdamW does not support sparse gradients')
                grads.append(p.grad)

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = _init_scalar(device=p.device if group['capturable'] else 'cpu')
                    # Exponential moving average of gradient values
                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # Exponential moving average of squared gradient values
                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                exp_avgs.append(state['exp_avg'])
                exp_avg_sqs.append(state['exp_avg_sq'])
                state_steps.append(state['step'])

            nadamw(
                params_with_grad,

View on GitHub (pinned to 9a5261e31b)

Solutions

  1. Set sparse=False on nn.Embedding layers (memory cost rises but gradients become dense)
  2. Use torch.optim.SparseAdam or an optimizer that supports sparse grads for embedding params
  3. Keep sparse-embedding params in a separate param group handled by a supporting optimizer

Example fix

# before
emb = nn.Embedding(num, dim, sparse=True)
opt = NAdamW(model.parameters(), lr=1e-3)

# after
emb = nn.Embedding(num, dim, sparse=False)
opt = NAdamW(model.parameters(), lr=1e-3)
Defensive patterns

Strategy: validation

Validate before calling

sparse_params = [n for n, p in model.named_parameters() if p.grad is not None and p.grad.is_sparse]
assert not sparse_params, f'sparse grads on: {sparse_params}'

Prevention

When it happens

Trigger: Calling optimizer.step() after a backward pass where at least one param grad is sparse, typically from nn.Embedding(sparse=True) used with NAdamW.

Common situations: NLP/recommendation models with sparse=True embeddings trained with NAdamW; switching an existing training script from SparseAdam/optimizers that tolerate sparsity to timm's NAdamW.

Related errors


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