You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
RMSNorm crashes on NPU when elementwise_affine=False (weight=None): npu_rms_norm requires a real gamma tensor #14590
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
importtorchfromdiffusers.models.normalizationimportRMSNorm# LTX-2 style block norm: no learnable affine => weight is Nonenorm=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
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).
Activity
DN6 commentedon Sep 22, 2026
@xucqX I believe this should be fixed after #14288
DN6 commentedon Oct 1, 2026
@xucqX Closing this since we merged #14288. Feel free to open if the issue is unresolved.