huggingface / huggingface/candle

candle-nn: ops::softmax_last_dim is apply_op1_no_bwd; severs autograd (same pattern as #2168 / PR #3526)

Open
#3,569 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

## Severity: HIGH — silently severs autograd through attention weights

\`candle_nn::ops::softmax_last_dim\` at \`candle-nn/src/ops.rs:430\` uses \`apply_op1_no_bwd\`. The typical use site is attention weights inside a model's forward pass, and the typical loss is computed downstream of those — so a \`loss.backward()\` call returns with no error, but \`grads.get(var)\` returns \`None\` for any \`Var\` upstream of the softmax.

The workaround \`softmax(t, D::Minus1)\` is pure-Candle differentiable and produces the same output. Same naming-footgun pattern as #2168 / PR #3526 (RmsNorm) and #3568 (rope): users reach for the obvious name and silently lose gradients.

**Verified present at candle main HEAD \`1d7e9274591ef7eb101d965b71243a41994214d2\`** (committed 2026-05-26).

## Reproducer

https://github.com/nxrobins/candle-bug-4-softmax-no-bwd

\`\`\`bash
git clone https://github.com/nxrobins/candle-bug-4-softmax-no-bwd && cd candle-bug-4-softmax-no-bwd && cargo run --release
# exit 1: BUG PRESENT (softmax_last_dim severs autograd; softmax(t, dim) propagates)
# exit 0: bug appears fixed
\`\`\`

\`cargo run\` and \`cargo test\` (2 tests) pass. Loss shape is \`(out * out).sum_all()\` rather than \`out.sum_all()\` because the latter is mathematically a constant for any softmax (sums to 1 per row) so its gradient is legitimately zero — the squared sum is non-uniform w.r.t. the input and gives a meaningful gradient signal for the positive control.

## Suggested fix

Two options:

**Option A (doc-only floor)** — \`#[doc]\` warning on \`softmax_last_dim\` pointing at \`softmax(t, dim)\`.

**Option B (supremum)** — implement \`bwd\` for the \`SoftmaxLastDim\` Op. Math:

\`\`\`
grad_in = softmax * (grad_out - sum(grad_out * softmax, dim=-1))
\`\`\`

This would mirror the shape of PR #3526 (\`bwd\` impl on the existing \`CustomOp\`). Watch the cuDNN-backed path if compiled with \`--features cudnn\`.

## Related

- #2168 / PR #3526 — same \`apply_op*_no_bwd\` pattern in \`RmsNorm\` (awaiting review).
- #3568 — same pattern in \`rotary_emb::rope\`.
- #3011 — broader LayerNorm gradient-flow.

## Context

Surfaced alongside three other \`_no_bwd\` bugs during Phase 0.5b validation: https://github.com/nxrobins/wave-agent/blob/main/docs/upstream-candle-bugs.md

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.