huggingface / huggingface/diffusers

[Bug] Ascend NPU: RMSNorm crashes with elementwise_affine=False; _native_npu FA rejects [B, N, 1, Skv] masks (LTX-2)

Offen
#14,380 3 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug models
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### Describe the bug

When running LTX-2 with `attn_backend=_native_npu` on Ascend NPU, two incompatibilities show up:

1. **RMSNorm** – layers with `elementwise_affine=False` leave `weight=None`, but `torch_npu.npu_rms_norm` requires a gamma tensor → crash (`gamma is None`).

2. **Fused attention mask** – LTX cross-attn uses masks shaped `[B, N, 1, Skv]` (e.g. `[1, 32, 1, 1024]`). Ascend FA does not broadcast the singleton query-length dim the way SDPA does, so `_maybe_modify_attn_mask_npu` must expand it to `[B, N, Sq, Skv]`. Today only `[B, 1, 1, Skv]` is expanded.

These bugs were discovered while running LTX-2.3 text-to-audio-video FlowGRPO training on Ascend NPU using [verl-omni](https://github.com/verl-project/verl-omni/blob/main/examples/flowgrpo_trainer/ltx2/run_ltx2_3_t2av_lora_npu.sh).

### Reproduction

Requires Ascend NPU + `torch_npu`. Minimal sketches of both failure modes:

```python
import torch
from diffusers.models.normalization import RMSNorm
from diffusers.models.attention_dispatch import (
AttentionBackendName,
attention_backend,
dispatch_attention_fn,
)

# --- Bug 1: RMSNorm with elementwise_affine=False ---
norm = RMSNorm(dim=64, eps=1e-6, elementwise_affine=False).to("npu")
x = torch.randn(2, 16, 64, device="npu", dtype=torch.float16)
# Crashes: npu_rms_norm called with weight=None
y = norm(x)

# --- Bug 2: FA mask [B, N, 1, Skv] under _native_npu ---
B, Sq, Skv, N, D = 1, 384, 1024, 32, 64
q = torch.randn(B, Sq, N, D, device="npu", dtype=torch.float16)
k = torch.randn(B, Skv, N, D, device="npu", dtype=torch.float16)
v = torch.randn(B, Skv, N, D, device="npu", dtype=torch.float16)
attn_mask = torch.zeros(B, N, 1, Skv, device="npu", dtype=torch.float16) # LTX-style

with attention_backend(AttentionBackendName._NATIVE_NPU):
# Fails: Ascend FA expects Sq on dim=-2, not a singleton 1
out = dispatch_attention_fn(q, k, v, attn_mask=attn_mask)
```

### Logs

# Error 1

Image

# Error 2

Image

### System Info

- OS: Linux aarch64
- Hardware: Ascend NPU
- Python: 3.11
- torch: 2.10.0
- torch_npu: 2.10.0
- CANN: 9.0.0
- diffusers: main

### Who can help?

@yiyixuxu @sayakpaul

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginne mit diffusers.models.normalization.RMSNorm und dem _maybe_modify_attn_mask_npu-Pfad in diffusers.models.attention_dispatch und führe anschließend die beiden bereitgestellten Ascend NPU-Reproduktionsskizzen aus. Als abgeschlossen gilt die Aufgabe, wenn RMSNorm mit elementwise_affine=False funktioniert und _native_npu LTX-artige Masken im Format [B, N, 1, Skv] akzeptiert, indem die Query-Längen-Dimension korrekt behandelt wird.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Aktivitätsstatus
Ruhig
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
52/100

Neue Issues direkt in Ihr Postfach

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