huggingface / huggingface/candle
candle-nn: ops::softmax_last_dim is apply_op1_no_bwd; severs autograd (same pattern as #2168 / PR #3526)
- 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.