huggingface / huggingface/diffusers
Kolors text encoder carries an unreachable torch<2 attention path built on the MPS-unsafe baddbmm(empty, beta=0) idiom
- 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
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