huggingface / huggingface/diffusers

Kolors text encoder carries an unreachable torch<2 attention path built on the MPS-unsafe baddbmm(empty, beta=0) idiom

Offen
#14,624 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

### Describe the bug

`CoreAttention.forward` in `pipelines/kolors/text_encoder.py` computes raw attention scores with:

```python
matmul_input_buffer = torch.empty(
output_size[0] * output_size[1], output_size[2], output_size[3],
dtype=query_layer.dtype, device=query_layer.device,
)
matmul_result = torch.baddbmm(matmul_input_buffer, ..., beta=0.0, alpha=(1.0 / self.norm_factor))
```

This relies on the documented `baddbmm` contract that `beta=0` causes `input` to be ignored, so NaN/Inf in the uninitialised buffer must not propagate. The MPS backend violates that contract: pytorch/pytorch#187521 (labeled `module: correctness (silent)`, fixed in pytorch `main` but **not in any released torch**, including 2.13.0 — the fix missed the release branch). #14438 was this exact mechanism biting SDXL through `Attention.get_attention_scores`, fixed by #14459. After #14459, this Kolors call site is the last remaining instance of the idiom in the repository.

### Reproduction

The primitive fails at exactly the shapes this code requests. On Apple Silicon (torch 2.13.0, M-series):

```python
import torch

for b_np, sq in [(16, 1024), (16, 2048)]: # b*np, seq — Kolors text-encoder score shapes
shape = (b_np, sq, sq)
junk = torch.full(shape, float("nan"), device="mps", dtype=torch.float16)
del junk # freed NaN pages go back to the allocator pool
buf = torch.empty(shape, device="mps", dtype=torch.float16) # recycles them
q = torch.randn(b_np, sq, 64, device="mps", dtype=torch.float16)
k = torch.randn(b_np, sq, 64, device="mps", dtype=torch.float16)
out = torch.baddbmm(buf, q, k.transpose(1, 2), beta=0.0, alpha=0.125)
print(shape, "buffer had NaN:", True, "-> output has NaN:", bool(out.isnan().any()))
```

prints `output has NaN: True` for both shapes on my machine (M5 Pro).

**Scope caveat, stated honestly:** I have *not* reproduced NaN in `CoreAttention`'s output end-to-end through the Kolors pipeline — in my attempts the allocator handed the buffer clean pages even after aggressive dirtying. In the #14438 case it is component offloading that reliably leaves large freed dirty regions behind, and the Kolors text encoder has no equivalent. So this is a latent correctness hazard (unsafe idiom + shapes confirmed vulnerable at the primitive level), not a bug I can show corrupting Kolors outputs today.

Fix in #14620: route the MPS case through a buffer-free scaled `bmm`, identical to the approach taken in #14459; all other devices keep the current `baddbmm` path unchanged.

### System Info

- diffusers @ `main`
- torch 2.13.0 (any released torch is affected; fixed only in pytorch `main`)
- Apple Silicon (M5 Pro), macOS / MPS backend

### Who can help?

@yiyixuxu @asomoza

---

**Correction (same day):** on re-verification I found the vulnerable call is *unreachable on any supported torch*. `CoreAttention.forward` branches on `int(torch.__version__.split(".")[0]) >= 2` and uses `scaled_dot_product_attention` for all of torch 2.x; the `baddbmm` code above only executes on torch 1.x, and diffusers requires torch ≥ 2.6. So this is not a live bug — it is a latent hazard inside dead code (which is also why end-to-end corruption was never reproducible). #14620 has been reworked accordingly: it now deletes the entire torch < 2 branch, verified bit-identical to `main` across devices, dtypes, and mask branches.

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Read pipelines/kolors/text_encoder.py and inspect CoreAttention.forward, especially the torch-version branch and its attention paths. Compare the proposed change in #14620 and run the relevant Kolors tests or verification across the mentioned devices, dtypes, and mask branches; done means the obsolete branch is removed without changing supported behavior.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Refactoring
Schwierigkeit
2/5
Geschätzter Aufwand
1-3 Stunden
Aktivitätsstatus
Veraltet
Klarheit
Klar beschrieben
Anfängerfreundlichkeit
25/100

Neue Issues direkt in Ihr Postfach

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