From 8bf0267cd145726888b0528d7a7b274b770c7d67 Mon Sep 17 00:00:00 2001 From: direwolf_xcq Date: Mon, 24 Aug 2026 11:36:56 +0800 Subject: [PATCH] fix: npu_rms_norm does not support gamma=None, use all-ones tensor instead --- src/diffusers/models/normalization.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/normalization.py b/src/diffusers/models/normalization.py index 84ffb67bfd6a..b7e9be8dcf8c 100644 --- a/src/diffusers/models/normalization.py +++ b/src/diffusers/models/normalization.py @@ -542,11 +542,13 @@ def forward(self, hidden_states): if is_torch_npu_available(): import torch_npu - if self.weight is not None: - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)[0] + # npu_rms_norm does not support gamma=None, use all-ones tensor instead + weight = self.weight + if weight is None: + weight = torch.ones(self.dim, device=hidden_states.device, dtype=hidden_states.dtype) + if weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(weight.dtype) + hidden_states = torch_npu.npu_rms_norm(hidden_states, weight, epsilon=self.eps)[0] if self.bias is not None: hidden_states = hidden_states + self.bias else: