huggingface / huggingface/diffusers
cuDNN attention backward corrupts gradients: K/V transposed twice, Q once
- Lingua principale
- Python
- Stelle
- 34.5k
- Fork
- 7.3k
- Merge medio
- 3g 3h
- PR unite (30g)
- 91
Descrizione
# 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.
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Inizia in src/diffusers/models/attention_dispatch.py, alle righe forward citate 934-937 e alle righe backward 971-974, quindi confronta i layout dei tensori salvati con la chiamata backward di ATen cuDNN. Riproduci il caso H100 bf16 e aggiungi o esegui il test di regressione dei gradienti; il lavoro è completato quando i gradienti K/V corrispondono al riferimento eager fp32 senza modificare il problema separato della mask.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- python, pytorch
- Ambito
- backend, machine-learning
- Tipo di issue
- Bug
- Difficoltà
- 3/5
- Tempo stimato
- 1-2 giorni
- Stato di attività
- Tranquilla
- Chiarezza
- Specificata chiaramente
- Idoneità per principianti
- 35/100