huggingface / huggingface/diffusers
cuDNN attention backward corrupts gradients: K/V transposed twice, Q once
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
# cuDNN attention backward pass corrupts gradients (K/V transposed twice, Q once)
Opening this as an issue first, per the *AI-assisted and agentic contributions*
section of CONTRIBUTING.md — I have a fix and a repro ready and would like a
maintainer's ack on the scope before opening the PR.
## What's wrong
`_cudnn_attention_forward_op` saves the attention inputs **after** transposing them
to `(B, H, S, D)`:
```python
# src/diffusers/models/attention_dispatch.py:934-937
query = query.transpose(1, 2).contiguous()
key = key.transpose(1, 2).contiguous()
value = value.transpose(1, 2).contiguous()
tensors_to_save += (query, key, value)
```
`_cudnn_attention_backward_op` then transposes `key` and `value` a **second** time,
while leaving `query` alone:
```python
# src/diffusers/models/attention_dispatch.py:971-974
query, key, value, out, lse, ... = ctx.saved_tensors # already (B, H, S, D)
grad_out = grad_out.transpose(1, 2).contiguous() # correct: grad_out arrives BSHD
key = key.transpose(1, 2).contiguous() # -> (B, S, H, D)
value = value.transpose(1, 2).contiguous() # -> (B, S, H, D)
```
so `torch.ops.aten._scaled_dot_product_cudnn_attention_backward` receives Q as
`(B, H, S, D)` and K/V as `(B, S, H, D)`. The `grad_out` transpose is correct and
should stay — the forward transposes `out` back to BSHD at :963.
## Impact
Only the context-parallel `_native_cudnn` path is affected (the simple path at
:3607 goes through `F.scaled_dot_product_attention`, which has its own autograd).
Forward is unaffected, so inference is fine and no forward-only test can catch it —
but any training/fine-tuning run on this backend optimises against garbage gradients.
## Repro
H100, bf16, `B=2 S=16 H=16 D=64`. `num_heads == seq_len` is deliberate, so a shape
check cannot mask the layout error. Errors are `max_abs` vs an eager fp32 reference:
| variant | out | dQ | dK | dV |
|---|---|---|---|---|
| all three passed as saved (BHSD) | 0.0078 | 0.0078 | 0.0078 | 0.0156 |
| current code | 0.0078 | **7.8984** | **7.9727** | **3.9785** |
Reference gradient magnitudes are 1.859 / 2.266 / 2.750 — the errors exceed the
values being approximated, so the gradients are unrelated, not merely imprecise.
## Proposed fix
Drop the two `transpose(1, 2)` calls on `key`/`value` in the backward op; the saved
tensors are already in the layout the ATen op expects.
**Split out:** the CP path also forwards a `torch.bool` mask straight into the ATen op's
additive `attn_bias` slot, where it is consumed as an additive `0`/`1`. That is tracked
separately in #14342 so this issue can be closed cleanly by #14341, which fixes the
K/V transpose only.
## Context
Found during an audit of cuDNN SDPA integrations across the ecosystem, run by the
NVIDIA cuDNN team. I have the fix, a gradient regression test, and the self-review
report ready to attach to the PR once you've had a chance to weigh in on scope.
Beitragsleitfaden
Rechercherichtung
Beginne in src/diffusers/models/attention_dispatch.py bei den zitierten forward-Zeilen 934-937 und backward-Zeilen 971-974 und vergleiche anschließend die gespeicherten Tensorlayouts mit dem backward-Aufruf von ATen cuDNN. Reproduziere den H100-bf16-Fall und füge den Gradient-Regressionstest hinzu oder führe ihn aus; fertig ist es, wenn die K/V-Gradienten mit der eager-fp32-Referenz übereinstimmen, ohne das separate Maskenproblem zu ändern.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- backend, machine-learning
- Issue-Typ
- Bug
- Schwierigkeit
- 3/5
- Geschätzter Aufwand
- 1-2 Tage
- Aktivitätsstatus
- Ruhig
- Klarheit
- Klar beschrieben
- Anfängerfreundlichkeit
- 35/100