huggingface / huggingface/diffusers

RMSNorm crashes on NPU when elementwise_affine=False (weight=None): npu_rms_norm requires a real gamma tensor

Open
#14,590 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

`diffusers.models.normalization.RMSNorm` crashes on Ascend NPU when constructed
with `elementwise_affine=False` (i.e. `self.weight is None`).

The NPU branch of `RMSNorm.forward` passes `self.weight` straight into the
fused CANN op `torch_npu.npu_rms_norm`, but that op's schema declares `gamma`
as a required `Tensor` (not `Optional`), so a `None` weight is invalid and the
forward errors out. The non-NPU `else` branch handles `weight=None` correctly.

### Reproduction

```python
import torch
from diffusers.models.normalization import RMSNorm

# LTX-2 style block norm: no learnable affine => weight is None
norm = RMSNorm(dim=4096, eps=1e-6, elementwise_affine=False).npu().to(torch.bfloat16)
x = torch.randn(1, 256, 4096, device="npu", dtype=torch.bfloat16)
out = norm(x) # raises on NPU
```

(Note: this needs an actual Ascend NPU device; on a non-NPU build it fails earlier
with a backend-availability error rather than the gamma=None error.)

### Logs

```shell
On Ascend NPU the call to `torch_npu.npu_rms_norm(hidden_states, None, eps)` violates
the op schema and the forward aborts. The op's schema (from `torch_npu`) is:

npu::npu_rms_norm(Tensor input, Tensor gamma, float epsilon=1e-6) -> (Tensor, Tensor)

`gamma` is a required `Tensor`, while `RMSNorm` can legitimately have
`self.weight = None` (when `elementwise_affine=False`, i.e. gamma == 1, no affine).
```

### System Info

- 🤗 Diffusers version: 0.40.0.dev0
- Platform: Linux-5.15.0-119-generic-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.11.15
- PyTorch version (GPU?): 2.9.0+cpu (False)
- Huggingface_hub version: 1.27.0
- Transformers version: 5.3.0
- Accelerate version: 1.10.1
- PEFT version: 0.18.0
- bitsandbytes version: 0.49.2
- optimum-quanto version: 0.2.7
- Safetensors version: 0.8.0
- xFormers version: not installed
- Accelerator: Ascend NPU (torch_npu 2.9.0.dev20260207)
- Using GPU in script?: Yes (Ascend NPU via `torch_npu`, `is_torch_npu_available()==True`)
- Using distributed or parallel set-up in script?: Yes (accelerate multi-NPU)

### Who can help?

_No response_

Contributor guide

Open the contributing guide

Research direction

Start in diffusers.models.normalization.RMSNorm, focusing on the NPU branch of forward and how it passes self.weight to torch_npu.npu_rms_norm. Reproduce on an Ascend NPU with elementwise_affine=False, then verify that forward no longer crashes for weight=None while the existing non-NPU behavior remains correct.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.