Skip to content

fix: npu_rms_norm does not support gamma=None, use all-ones tensor in… - #14585

Open
xucqX wants to merge 1 commit into
huggingface:mainfrom
xucqX:main
Open

fix: npu_rms_norm does not support gamma=None, use all-ones tensor in…#14585
xucqX wants to merge 1 commit into
huggingface:mainfrom
xucqX:main

Conversation

@xucqX

@xucqX xucqX commented Aug 24, 2026

Copy link
Copy Markdown

fixes: #14590

What does this PR do?

Fixes a crash in RMSNorm.forward on Ascend NPU when the layer is built with
elementwise_affine=False (i.e. self.weight is None).

Context

RMSNorm supports elementwise_affine=False, which sets self.weight = None
(no learnable gamma, mathematically gamma == 1). The non-NPU else branch
handles this fine, but the NPU branch explicitly calls the fused CANN op

torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)

and passes self.weight straight through. npu_rms_norm does not accept
gamma=None (it requires a real tensor), so any model that uses weightless
RMSNorm blocks crashes on Ascend the moment forward hits one of them.

Concrete in-repo trigger: LTX2VideoTransformerBlock
(src/diffusers/models/transformers/transformer_ltx2.py) constructs its block
norms (norm1/norm2/norm3, the audio variants, and the a2v/v2a cross-attn
norms) with elementwise_affine=False, so LTX-2 is unusable on NPU without
this fix.

The fix

When self.weight is None, substitute an all-ones tensor of the same shape,
dtype and device:

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]

Multiplying by an all-ones tensor is identical to "no affine" (gamma=1), so
this preserves the original semantics while giving the CANN op a valid tensor.
The CPU/GPU else branch and bias handling are unchanged.

No new dependencies.

Self-review notes (AI-assisted)

  • Math equivalence: ones ≡ gamma=1 ≡ the "no affine" semantics the
    else branch already implements for the weight=None case. No behaviour
    change on CPU/GPU.
  • Alternative considered: fall back to the else (decomposed) branch on
    NPU when weight is None. Rejected — the explicit NPU branch exists
    precisely to get the fused CANN kernel in eager mode; the ones-shim keeps
    the fast path rather than degrading to a multi-launch decomposition.
  • dtype/device correctness: the substitute weight is created with
    hidden_states.device / hidden_states.dtype, so the subsequent
    .to(weight.dtype) is a no-op for the None case and does not silently
    promote/demote precision.
  • Tests not added: the NPU path only runs when is_torch_npu_available()
    is true, which is not the case in HF CI. Mocking the guard would not
    exercise the real CANN op and would be fragile. Happy to add a
    device-agnostic test (assert the weight=None path matches a reference
    RMSNorm with gamma=1 on CPU) if reviewers want one.
  • Intentionally did not touch the else branch or bias handling.

Before submitting

  • Did you use an AI agent (Claude Code) to help with this PR?
  • Read the contributor guideline.
  • Read the philosophy doc (change is small and localized, bugfix only).
  • Discussed via issue/forum: none yet — opening this PR as the discussion. Can file a tracking issue if preferred.
  • Documentation: no doc changes (internal forward path, no public API change).
  • Tests: not added, see self-review notes.

Who can review?

NPU doesn't have a dedicated owner in the list; this is a model/normalization
change, so tagging:

@yiyixuxu @dg845

@github-actions github-actions Bot added size/S PR with diff < 50 LOC models and removed size/S PR with diff < 50 LOC labels Aug 24, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Hi @xucqX, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

@xucqX

xucqX commented Aug 24, 2026

Copy link
Copy Markdown
Author

fixes: #14590

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant