fix: npu_rms_norm does not support gamma=None, use all-ones tensor in… - #14585
Open
xucqX wants to merge 1 commit into
Open
fix: npu_rms_norm does not support gamma=None, use all-ones tensor in…#14585xucqX wants to merge 1 commit into
xucqX wants to merge 1 commit into
Conversation
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. 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 |
Author
|
fixes: #14590 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fixes: #14590
What does this PR do?
Fixes a crash in
RMSNorm.forwardon Ascend NPU when the layer is built withelementwise_affine=False(i.e.self.weight is None).Context
RMSNormsupportselementwise_affine=False, which setsself.weight = None(no learnable gamma, mathematically gamma == 1). The non-NPU
elsebranchhandles this fine, but the NPU branch explicitly calls the fused CANN op
and passes
self.weightstraight through.npu_rms_normdoes not acceptgamma=None(it requires a real tensor), so any model that uses weightlessRMSNorm 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 blocknorms (
norm1/norm2/norm3, the audio variants, and the a2v/v2a cross-attnnorms) with
elementwise_affine=False, so LTX-2 is unusable on NPU withoutthis fix.
The fix
When
self.weight is None, substitute an all-ones tensor of the same shape,dtype and device:
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
elsebranch and bias handling are unchanged.No new dependencies.
Self-review notes (AI-assisted)
ones≡ gamma=1 ≡ the "no affine" semantics theelsebranch already implements for theweight=Nonecase. No behaviourchange on CPU/GPU.
else(decomposed) branch onNPU when
weight is None. Rejected — the explicit NPU branch existsprecisely to get the fused CANN kernel in eager mode; the ones-shim keeps
the fast path rather than degrading to a multi-launch decomposition.
weightis created withhidden_states.device/hidden_states.dtype, so the subsequent.to(weight.dtype)is a no-op for theNonecase and does not silentlypromote/demote precision.
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=Nonepath matches a referenceRMSNorm with gamma=1 on CPU) if reviewers want one.
elsebranch or bias handling.Before submitting
self-reviewon the diff (notes above).Who can review?
NPU doesn't have a dedicated owner in the list; this is a model/normalization
change, so tagging:
@yiyixuxu @dg845