[Bugfix] Fix NpuRMSNorm accuracy on Qwen3#245
Conversation
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).
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the residual parameterization detection in NpuRMSNorm to check for the absence of variance_epsilon instead of using a mean weight heuristic. It also adds support for forcing FP32 precision during the forward pass. The reviewer suggests simplifying the forward method to avoid duplicating the residual parameterization check and casting self.weight to the target dtype before performing the addition to prevent potential precision loss or overflow.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
2 participants
PR type
PR information
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.
Experiment results
Verified on Qwen3-8B (4x NPU): loss 4.68 vs HF reference 4.64 (was 8.83).