Affineless RMSNorm Issue Related to Ascend
An issue related to invoking affine-less RMSNorm operator on Ascend NPUs. The issue can be reproduced on diffusers 0.40.0 but has already been fixed in 0.41.0
Table of Contents
1. Problem Description
LTX2.3 transformers set norm_elementwise_affine=False, so every block norm (norm1 or audio_norm1 or …) carries weight=None. diffusers’ RMSNorm.forward NPU branch feeds that None unguarded into torch_npu.npu_rms_norm – "Expected a proper Tensor but got None ... (argument 'gamma')" – crashing the first trainer-side forward right after the first rollout.
2. Solution/Patch
Note, this problem has already been fixed in diffusers 0.41.0 (Issue #14380, related fix PR #14288). The following patch applies to verl-omni 0.2.0rc1 (commit)
Patch diffusers.models.normalization.RMSNorm.forward so the NPU + weight=None case computes the same plain RMS normalization the upstream CPU/CUDA branch already uses, and delegate every other path (NPU with weights, all non-NPU backends) to the upstream implementation untouched. Idempotent, installed from verl_omni/__init__.py before any engine forward; no behavior change on CUDA/CPU.