Skip to content

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

Open
xucqX wants to merge 2 commits into
huggingface:mainfrom
xucqX:main
Open

xucqX wants to merge 2 commits 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

@github-actions github-actions Bot added fixes-issue size/S PR with diff < 50 LOC labels Sep 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fixes-issue models size/S PR with diff < 50 LOC

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

2 participants