From 46ebbe7685faae8acc7f0b73f4d99962281b1065 Mon Sep 17 00:00:00 2001 From: XianlongLi <34086799+XianlongLi@users.noreply.github.com> Date: Fri, 10 Jul 2026 08:24:26 +0000 Subject: [PATCH 1/4] Fix NpuRMSNorm accuracy: false residual-param detection on trained Qwen3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The _twinkle_residual_param() heuristic used abs(weight.mean()) < 0.3 to detect residual parameterization (Qwen3.5: scale = 1 + weight). This was unreliable: trained Qwen3-8B weights can have small means (e.g. 0.011), falsely triggering residual mode on 18/145 modules and inflating the scale by ~90x (using 1+weight instead of weight). This corrupted the forward output and inflated the training loss by ~1.9x (8.83 vs correct 4.69). Fix: detect by attribute name instead of weight values — Qwen3.5 uses 'eps' (residual), Qwen3/Qwen2/Llama use 'variance_epsilon' (standard). Also add TWINKLE_NPU_GATED_RMSNorm_FP32 env var support to NpuRMSNorm.forward (matching npu_gated_rms_norm_forward), upcasting hidden_states/weight to fp32 when enabled. Verified on Qwen3-8B (4x NPU): loss 4.68 vs HF reference 4.64 (was 8.83). --- src/twinkle/kernel/npu_impls/rms_norm.py | 32 ++++++++++++++++++------ 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/src/twinkle/kernel/npu_impls/rms_norm.py b/src/twinkle/kernel/npu_impls/rms_norm.py index 7fd7a5f7..10e11ae5 100644 --- a/src/twinkle/kernel/npu_impls/rms_norm.py +++ b/src/twinkle/kernel/npu_impls/rms_norm.py @@ -25,10 +25,20 @@ class NpuRMSNorm(nn.Module): """ def _twinkle_residual_param(self) -> bool: - """Lazily detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight).""" + """Detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight). + + Detection is by attribute name, not weight values: Qwen3.5 RMSNorm uses + ``eps`` and residual parameterization (scale = 1 + weight); Qwen3 / + Qwen2 / Llama etc. use ``variance_epsilon`` and standard parameterization + (scale = weight). The previous weight-value heuristic + (``abs(mean) < 0.3``) was unreliable: trained Qwen3 weights can have + small means, falsely triggering residual mode and inflating the scale + by ~90×, which corrupted the forward output and inflated the training + loss by ~1.9×. + """ 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)') @@ -39,12 +49,20 @@ def _twinkle_eps(self) -> float: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: import torch_npu - target_dtype = hidden_states.dtype - if self._twinkle_residual_param(): - scale = (1.0 + self.weight).to(target_dtype) + input_dtype = hidden_states.dtype + if _FORCE_FP32: + hidden_states = hidden_states.to(torch.float32) + if self._twinkle_residual_param(): + scale = (1.0 + self.weight).float() + else: + scale = self.weight.float() else: - scale = self.weight.to(target_dtype) - return torch_npu.npu_rms_norm(hidden_states, scale, epsilon=self._twinkle_eps())[0] + if self._twinkle_residual_param(): + scale = (1.0 + self.weight).to(input_dtype) + else: + scale = self.weight.to(input_dtype) + out = torch_npu.npu_rms_norm(hidden_states, scale, epsilon=self._twinkle_eps())[0] + return out.to(input_dtype) if _FORCE_FP32 else out # Resolved once at import: matches the legacy "patch-time, process-wide" invariant. From 643f26026cd833f9d5606ddff644f4a38e22e79b Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Fri, 10 Jul 2026 18:06:46 +0800 Subject: [PATCH 2/4] Update rms_norm.py --- src/twinkle/kernel/npu_impls/rms_norm.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/src/twinkle/kernel/npu_impls/rms_norm.py b/src/twinkle/kernel/npu_impls/rms_norm.py index 10e11ae5..963082f8 100644 --- a/src/twinkle/kernel/npu_impls/rms_norm.py +++ b/src/twinkle/kernel/npu_impls/rms_norm.py @@ -25,17 +25,6 @@ class NpuRMSNorm(nn.Module): """ def _twinkle_residual_param(self) -> bool: - """Detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight). - - Detection is by attribute name, not weight values: Qwen3.5 RMSNorm uses - ``eps`` and residual parameterization (scale = 1 + weight); Qwen3 / - Qwen2 / Llama etc. use ``variance_epsilon`` and standard parameterization - (scale = weight). The previous weight-value heuristic - (``abs(mean) < 0.3``) was unreliable: trained Qwen3 weights can have - small means, falsely triggering residual mode and inflating the scale - by ~90×, which corrupted the forward output and inflated the training - loss by ~1.9×. - """ cached = getattr(self, '_twinkle_residual_cached', None) if cached is None: cached = not hasattr(self, 'variance_epsilon') From 8a2d466cdf0059f0a65ae5e1b313370438d6ce7f Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Fri, 10 Jul 2026 18:11:39 +0800 Subject: [PATCH 3/4] Update rms_norm.py --- src/twinkle/kernel/npu_impls/rms_norm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/twinkle/kernel/npu_impls/rms_norm.py b/src/twinkle/kernel/npu_impls/rms_norm.py index 963082f8..198585ca 100644 --- a/src/twinkle/kernel/npu_impls/rms_norm.py +++ b/src/twinkle/kernel/npu_impls/rms_norm.py @@ -25,6 +25,7 @@ class NpuRMSNorm(nn.Module): """ 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 = not hasattr(self, 'variance_epsilon') From 1ea0832a4ebbe17ef845f56e9334a9db12ff4d53 Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Fri, 10 Jul 2026 18:49:47 +0800 Subject: [PATCH 4/4] Update rms_norm.py --- src/twinkle/kernel/npu_impls/rms_norm.py | 27 ++++++++++-------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/src/twinkle/kernel/npu_impls/rms_norm.py b/src/twinkle/kernel/npu_impls/rms_norm.py index 198585ca..7996db95 100644 --- a/src/twinkle/kernel/npu_impls/rms_norm.py +++ b/src/twinkle/kernel/npu_impls/rms_norm.py @@ -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. @@ -25,7 +29,7 @@ class NpuRMSNorm(nn.Module): """ def _twinkle_residual_param(self) -> bool: - """Lazily detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight).""" + """Lazily detect residual parameterization (e.g. Qwen3.5: scale = 1 + weight).""" cached = getattr(self, '_twinkle_residual_cached', None) if cached is None: cached = not hasattr(self, 'variance_epsilon') @@ -40,26 +44,17 @@ def _twinkle_eps(self) -> float: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: import torch_npu 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).float() - else: - scale = self.weight.float() + if self._twinkle_residual_param(): + scale = 1.0 + self.weight.to(target_dtype) else: - if self._twinkle_residual_param(): - scale = (1.0 + self.weight).to(input_dtype) - else: - scale = self.weight.to(input_dtype) + scale = self.weight.to(target_dtype) out = torch_npu.npu_rms_norm(hidden_states, scale, epsilon=self._twinkle_eps())[0] return out.to(input_dtype) if _FORCE_FP32 else out -# 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') - - def npu_gated_rms_norm_forward(self, hidden_states, gate=None): """Forward replacement for Gated RMSNorm variants (e.g. Qwen3.5-MoE).""" import torch_npu @@ -69,7 +64,7 @@ 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 @@ -77,4 +72,4 @@ def npu_gated_rms_norm_forward(self, hidden_states, gate=None): 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