huggingface / huggingface/diffusers

cuDNN attention backward corrupts gradients: K/V transposed twice, Q once

Offen
#14,338 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
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

Beitragsleitfaden öffnen

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

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.