huggingface / huggingface/diffusers
RMSNorm crashes on NPU when elementwise_affine=False (weight=None): npu_rms_norm requires a real gamma tensor
- 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
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