Skip to content

[Bugfix] Fix NpuRMSNorm accuracy on Qwen3 #245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ys2025-AI wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[Bugfix] Fix NpuRMSNorm accuracy on Qwen3 #245

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
46ebbe7
Fix NpuRMSNorm accuracy: false residual-param detection on trained Qwen3
XianlongLi Jul 10, 2026
643f260
Update rms_norm.py
ys2025-AI Jul 10, 2026
8a2d466
Update rms_norm.py
ys2025-AI Jul 10, 2026
1ea0832
Update rms_norm.py
ys2025-AI Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions src/twinkle/kernel/npu_impls/rms_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@

logger = get_logger()

# Resolved once at import: matches the legacy "patch-time, process-wide" invariant.
# Mid-process env mutation will not retroactively change behavior.
_FORCE_FP32 = os.environ.get('TWINKLE_NPU_GATED_RMSNorm_FP32', '0').lower() in ('1', 'true', 'on', 'yes')


class NpuRMSNorm(nn.Module):
"""Class-replacement impl for HF RMSNorm variants.
Expand All @@ -28,7 +32,7 @@ def _twinkle_residual_param(self) -> bool:
"""Lazily detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight)."""
cached = getattr(self, '_twinkle_residual_cached', None)
if cached is None:
cached = abs(self.weight.data.mean().item()) < 0.3
cached = not hasattr(self, 'variance_epsilon')
self._twinkle_residual_cached = cached
if cached:
logger.debug('[NPU] NpuRMSNorm using residual parameterization (1.0 + weight)')
Expand All @@ -39,17 +43,16 @@ def _twinkle_eps(self) -> float:

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
import torch_npu
target_dtype = hidden_states.dtype
input_dtype = hidden_states.dtype
target_dtype = torch.float32 if _FORCE_FP32 else input_dtype
if _FORCE_FP32:
hidden_states = hidden_states.to(torch.float32)
if self._twinkle_residual_param():
scale = (1.0 + self.weight).to(target_dtype)
scale = 1.0 + self.weight.to(target_dtype)
else:
scale = self.weight.to(target_dtype)
return torch_npu.npu_rms_norm(hidden_states, scale, epsilon=self._twinkle_eps())[0]


# Resolved once at import: matches the legacy "patch-time, process-wide" invariant.
# Mid-process env mutation will not retroactively change behavior.
_FORCE_FP32 = os.environ.get('TWINKLE_NPU_GATED_RMSNorm_FP32', '0').lower() in ('1', 'true', 'on', 'yes')
out = torch_npu.npu_rms_norm(hidden_states, scale, epsilon=self._twinkle_eps())[0]
return out.to(input_dtype) if _FORCE_FP32 else out
Comment thread
ys2025-AI marked this conversation as resolved.


def npu_gated_rms_norm_forward(self, hidden_states, gate=None):
Expand All @@ -61,12 +64,12 @@ def npu_gated_rms_norm_forward(self, hidden_states, gate=None):

if _FORCE_FP32:
hidden_states = hidden_states.to(torch.float32)
weight = self.weight.float()
weight = self.weight.to(torch.float32)
gate = gate.to(torch.float32) if gate is not None else None
else:
weight = self.weight

hidden_states = torch_npu.npu_rms_norm(hidden_states, weight, epsilon=_eps)[0]
if gate is not None:
hidden_states = hidden_states * F.silu(gate)
return hidden_states.to(input_dtype)
return hidden_states.to(input_dtype) if _FORCE_FP32 else hidden_states
Loading