huggingface / huggingface/diffusers

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

Aperta
#14,380 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug models
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### 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

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia con diffusers.models.normalization.RMSNorm e il percorso _maybe_modify_attn_mask_npu in diffusers.models.attention_dispatch, quindi esegui i due sketch di riproduzione Ascend NPU forniti. Il lavoro è completato quando RMSNorm funziona con elementwise_affine=False e _native_npu accetta maschere in stile LTX con forma [B, N, 1, Skv], gestendo correttamente la dimensione della lunghezza della query.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
52/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.