Skip to content

Fix Gemma4 scale-free V norm FP16 overflow on CUDA - #253

Merged
justinchuby merged 4 commits into
mainfrom
fix-gemma4-vnorm-fp16-overflow
May 5, 2026
Merged

Fix Gemma4 scale-free V norm FP16 overflow on CUDA#253
justinchuby merged 4 commits into
mainfrom
fix-gemma4-vnorm-fp16-overflow

Conversation

@justinchuby

@justinchuby justinchuby commented May 5, 2026

Copy link
Copy Markdown
Member

Gemma4's parameterless V normalization (v / sqrt(mean(v²) + ε)) squared FP16 values directly. V projection outputs reach ~888, and 888² overflows FP16 max (65504) → inf → 0. This caused all-zero V outputs on CUDA (CPU uses FP32 internally).

Fix: Cast to FP32 before squaring, compute full RMSNorm in FP32, CastLike back.

Result: F16 CUDA 151.5 tok/s on H200 (was NaN).

@justinchuby
Fix Gemma4 scale-free V norm FP16 overflow on CUDA
e228380 
Gemma4's parameterless V normalization (v / sqrt(mean(v²) + ε))
squared FP16 values directly via op.Mul(v, v). V projection outputs
can reach ~888, and 888² = 788,544 overflows FP16 max (65504),
producing inf → mean(inf) → sqrt(inf) → v/inf = 0.

Fix: Cast to FP32 before squaring, compute RMSNorm in FP32, CastLike
back. Applied to _Gemma4ScaleFreeRMSNorm and both inline V norms.

Result: F16 CUDA 151.5 tok/s on H200 (was NaN).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby requested review from a team and Copilot May 5, 2026 18:54
This was referenced May 5, 2026
Closed
Open
@github-actions

github-actions Bot commented May 5, 2026 •
edited
Loading

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing c1daeaa → 19ed4d0

Model Sub-model Changes Status
gemma4 (gemma4) decoder 9 🟡
gemma4 (gemma4) embedding 0 ⚪
gemma4 (gemma4) vision_encoder 74 🟡
gemma4_text model 9 🟡
gemma4 (gemma4) / decoder — 9 change(s)

Op summary: 125 → 127 nodes

--- base
+++ head
@@ -47,13 +47,14 @@
 MatMul
 Constant
 Reshape
+Cast
 Mul
 ReduceMean
 Constant
-CastLike
 Add
 Sqrt
 Div
+CastLike
 Reshape
 Attention
 Transpose
@@ -90,13 +91,14 @@
 MatMul
 Constant
 Reshape
+Cast
 Mul
 ReduceMean
 Constant
-CastLike
 Add
 Sqrt
 Div
+CastLike
 Reshape
 Attention
 Transpose

Added nodes:

  • + Cast
  • + CastLike
  • + Cast
  • + CastLike

Removed nodes:

  • - CastLike
  • - CastLike

Connectivity changes:

  • node[53] Add: input_ids [101, 103] → [102, 103]
  • node[55] Div: input_ids [99, 105] → [100, 105]
  • node[105] RMSNormalization: input_ids [159, 43] → [159, 42]
gemma4 (gemma4) / vision_encoder — 74 change(s)

Op summary: 202 → 190 nodes

--- base
+++ head
@@ -50,14 +50,8 @@
 Reshape
 RMSNormalization
 RMSNormalization
-Mul
-Constant
-ReduceMean
-Constant
 CastLike
-Add
-Sqrt
-Div
+RMSNormalization
 Constant
 Gather
 Constant
@@ -181,21 +175,15 @@
 Mul
 Add
 Constant
+CastLike
 OneHot
-CastLike
 Transpose
 MatMul
 Constant
 CastLike
 Mul
-Mul
-Constant
-ReduceMean
-Constant
 CastLike
-Add
-Sqrt
-Div
+RMSNormalization
 Transpose
 MatMul
 Constant

Added nodes:

  • + RMSNormalization
  • + CastLike
  • + RMSNormalization

Removed nodes:

  • - Mul
  • - Constant
  • - ReduceMean
  • - Constant
  • - Add
  • - Sqrt
  • - Div
  • - CastLike
  • - Mul
  • - Constant
  • - ReduceMean
  • - Constant
  • - Add
  • - Sqrt
  • - Div

Modified attributes:

  • node[65] Gather: axis: 2 → 0
  • node[144] Constant: value_float: 0.0 → None, value_int: None → 1
  • node[160] Constant: value_int: 0 → 1
  • node[162] Constant: value_int: 0 → None, value_ints: None → [1]
  • node[164] Constant: value_int: None → 1, value_ints: [1] → None
  • node[166] Constant: value_int: 1 → None, value_ints: None → [3]
  • node[170] Constant: value_int: 1 → None, value_ints: None → [0, 1]
  • node[176] Constant: value_floats: None → [0.0, 0.1111111111111111], value_ints: [0, 1] → None

Connectivity changes:

  • node[1] CastLike: input_ids [59, 0] → [61, 0]
  • node[3] CastLike: input_ids [61, 0] → [63, 0]
  • node[4] Mul: input_ids [0, 60] → [0, 62]
  • node[5] Sub: input_ids [63, 62] → [65, 64]
  • node[7] MatMul: input_ids [64, 65] → [66, 67]
  • node[9] Gather: input_ids [1, 67] → [1, 69]
  • node[11] Equal: input_ids [68, 69] → [70, 71]
  • node[13] Clip: input_ids [1, 71] → [1, 73]
  • node[15] Gather: input_ids [72, 73] → [74, 75]
  • node[17] Gather: input_ids [72, 75] → [74, 77]
  • node[19] Gather: input_ids [2, 77] → [2, 79]
  • node[21] Gather: input_ids [2, 79] → [2, 81]
  • node[22] Gather: input_ids [78, 74] → [80, 76]
  • node[23] Gather: input_ids [80, 76] → [82, 78]
  • node[24] Add: input_ids [81, 82] → [83, 84]
  • node[25] CastLike: input_ids [4, 83] → [4, 85]
  • node[26] Not: input_ids [70] → [72]
  • node[27] Unsqueeze: input_ids [85, 5] → [87, 5]
  • node[28] Where: input_ids [86, 83, 84] → [88, 85, 86]
  • node[29] Add: input_ids [66, 87] → [68, 89]
  • node[30] CastLike: input_ids [6, 88] → [6, 90]
  • node[31] CastLike: input_ids [4, 88] → [4, 90]
  • node[32] Where: input_ids [70, 89, 90] → [72, 91, 92]
  • node[33] Unsqueeze: input_ids [91, 7] → [93, 7]
  • node[34] RMSNormalization: input_ids [88, 8] → [90, 8]
  • node[35] Clip: input_ids [93, 12, 13] → [95, 12, 13]
  • node[37] MatMul: input_ids [94, 95] → [96, 97]
  • node[38] Clip: input_ids [96, 14, 15] → [98, 14, 15]
  • node[39] Clip: input_ids [93, 17, 18] → [95, 17, 18]
  • node[41] MatMul: input_ids [98, 99] → [100, 101]
  • node[42] Clip: input_ids [100, 19, 20] → [102, 19, 20]
  • node[43] Clip: input_ids [93, 22, 23] → [95, 22, 23]
  • node[45] MatMul: input_ids [102, 103] → [104, 105]
  • node[46] Clip: input_ids [104, 24, 25] → [106, 24, 25]
  • node[47] Reshape: input_ids [97, 26] → [99, 26]
  • node[48] Reshape: input_ids [101, 26] → [103, 26]
  • node[49] Reshape: input_ids [105, 26] → [107, 26]
  • node[50] RMSNormalization: input_ids [106, 27] → [108, 27]
  • node[51] RMSNormalization: input_ids [107, 28] → [109, 28]
  • node[65] Gather: input_ids [1, 123] → [10, 122]
  • node[81] Slice: input_ids [135, 33, 30, 31] → [132, 30, 34, 32]
  • node[85] Mul: input_ids [142, 132] → [132, 129]
  • node[96] Slice: input_ids [137, 33, 30, 31] → [134, 30, 34, 32]
  • node[100] Mul: input_ids [157, 132] → [134, 129]
  • node[112] Transpose: input_ids [118] → [172]
  • node[173] Div: input_ids [226, 231] → [218, 227]
  • node[184] CastLike: input_ids [242, 218] → [59, 244]

Initializer changes:

  • initializer count 57 → 59
gemma4_text / model — 9 change(s)

Op summary: 127 → 129 nodes

--- base
+++ head
@@ -49,13 +49,14 @@
 MatMul
 Constant
 Reshape
+Cast
 Mul
 ReduceMean
 Constant
-CastLike
 Add
 Sqrt
 Div
+CastLike
 Reshape
 Attention
 Transpose
@@ -92,13 +93,14 @@
 MatMul
 Constant
 Reshape
+Cast
 Mul
 ReduceMean
 Constant
-CastLike
 Add
 Sqrt
 Div
+CastLike
 Reshape
 Attention
 Transpose

Added nodes:

  • + Cast
  • + CastLike
  • + Cast
  • + CastLike

Removed nodes:

  • - CastLike
  • - CastLike

Connectivity changes:

  • node[55] Add: input_ids [104, 106] → [105, 106]
  • node[57] Div: input_ids [102, 108] → [103, 108]
  • node[107] RMSNormalization: input_ids [162, 44] → [162, 43]

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

@github-actions

github-actions Bot commented May 5, 2026 •
edited
Loading

Copy link
Copy Markdown

Performance Comparison

Comparing c1daeaa → 19ed4d0

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0% ⚪
bert (feature-extraction) num_nodes 60 60 +0.0% ⚪
falcon model_size_bytes 364 KB 364 KB +0.0% ⚪
falcon num_nodes 66 66 +0.0% ⚪
gemma2 model_size_bytes 428 KB 428 KB +0.0% ⚪
gemma2 num_nodes 107 107 +0.0% ⚪
gpt2 model_size_bytes 388 KB 388 KB +0.0% ⚪
gpt2 num_nodes 53 53 +0.0% ⚪
llama model_size_bytes 425 KB 425 KB +0.0% ⚪
llama num_nodes 61 61 +0.0% ⚪
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0% ⚪
llama (static-cache) num_nodes 58 58 +0.0% ⚪
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0% ⚪
mamba (ssm-text-generation) num_nodes 98 98 +0.0% ⚪
phi3 model_size_bytes 421 KB 421 KB +0.0% ⚪
phi3 num_nodes 59 59 +0.0% ⚪
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0% ⚪
phi3 (static-cache) num_nodes 56 56 +0.0% ⚪
qwen2 model_size_bytes 425 KB 425 KB +0.0% ⚪
qwen2 num_nodes 61 61 +0.0% ⚪
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0% ⚪
qwen2 (static-cache) num_nodes 58 58 +0.0% ⚪
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0% ⚪
qwen3_5_moe (hybrid-text-generation) num_nodes 275 275 +0.0% ⚪
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0% ⚪
qwen3_5_text (hybrid-text-generation) num_nodes 129 129 +0.0% ⚪
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0% ⚪
qwen3_5_vl (hybrid-qwen-vl) num_nodes 408 408 +0.0% ⚪
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0% ⚪
t5 (seq2seq) num_nodes 166 166 +0.0% ⚪
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0% ⚪
whisper (speech-to-text) num_nodes 128 128 +0.0% ⚪

No performance regressions.

@codecov

codecov Bot commented May 5, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.65517% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/mobius/models/gemma4.py 57.14% 3 Missing ⚠️

📢 Thoughts on this report? Let us know!

titaiwangms
titaiwangms approved these changes May 5, 2026
View reviewed changes

@titaiwangms titaiwangms left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way to avoid copilot doing this again?

@titaiwangms

titaiwangms commented May 5, 2026

Copy link
Copy Markdown
Contributor

It's concerning especially mobius would be heavily relying on ai if we don't have a good steering.

@justinchuby

justinchuby commented May 5, 2026

Copy link
Copy Markdown
Member Author

I think we can have unit tests to test the overflow behavior to harden this. What do you think?

@titaiwangms

titaiwangms commented May 5, 2026

Copy link
Copy Markdown
Contributor

I think we can have unit tests to test the overflow behavior to harden this. What do you think?

Sure. But also I think it would be good to have some ort knowledge in mobius as well. ORT is the major backend.

@justinchuby
Use RMSNormalization op for scale-free V norm with stash_type=1
70440c4 
Replace manual Cast+Mul+ReduceMean+Sqrt+Div decomposition with the
proper RMSNormalization op using stash_type=1 (float32 accumulation).
This handles FP16 overflow natively without manual F32 casts.

The constant all-ones weight is initialized with const_value so no
external weight file is needed. CastLike ensures type compatibility.

Add unit test verifying V norm handles input=888 (888²=788K > FP16
max 65504) without producing NaN/Inf.

6 gemma4 unit tests pass, 13 build tests pass.

Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby

justinchuby commented May 5, 2026

Copy link
Copy Markdown
Member Author

I think we can have unit tests to test the overflow behavior to harden this. What do you think?

Sure. But also I think it would be good to have some ort knowledge in mobius as well. ORT is the major backend.

Maybe have it as a skill? Or we can use onnx-doctor to detect unsupported dtypes? What do you think?

@justinchuby
Simplify: use nn.Parameter(data=...) for scale-free weight
532acfe 
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@titaiwangms

titaiwangms commented May 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

I think we can have unit tests to test the overflow behavior to harden this. What do you think?

Sure. But also I think it would be good to have some ort knowledge in mobius as well. ORT is the major backend.

Maybe have it as a skill? Or we can use onnx-doctor to detect unsupported dtypes? What do you think?

Maybe it's sufficient that we point it to https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/cuda_execution_provider.cc for registered kernels (as well as contrib op and cpu ones). Just for copilot to understand what options it has when it's building models.

@justinchuby
Fix lint errors
19ed4d0 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby justinchuby mentioned this pull request May 5, 2026
Open
@justinchuby
justinchuby force-pushed the fix-gemma4-vnorm-fp16-overflow branch from 532acfe to e228380 Compare May 5, 2026 21:28
@justinchuby
justinchuby merged commit a0b7ce1 into main May 5, 2026
19 of 23 checks passed
@justinchuby
justinchuby deleted the fix-gemma4-vnorm-fp16-overflow branch May 5, 2026 23:05
@justinchuby justinchuby mentioned this pull request May 5, 2026
Closed
justinchuby added a commit that referenced this pull request May 6, 2026
@justinchuby
Fix audio encoder SkipSimplifiedLayerNormalization CUDA crash
8700af6 
ORT fuses Add(output_proj.bias) + RMSNormalization into
SkipSimplifiedLayerNormalization, placing the 1D bias as the 'skip'
input. ORT's CUDA kernel rejects 1D skip (requires 2D+), while the
CPU kernel accepts it.

This fusion was enabled by PR #253 which changed _Gemma4ScaleFreeRMSNorm
from manual primitive ops to op.RMSNormalization(stash_type=1). The
RMSNormalization op is recognized by ORT's SkipLayerNorm fusion pattern.

Fix: inline manual RMSNorm ops in _Gemma4AudioEncoderModel.forward()
for the pre_projection_norm, preventing ORT from recognizing the fusion
pattern. This preserves the FP32 accumulation for numerical stability.

Tested: audio encoder runs on CUDA with correct output (standalone ORT
and GenAI at 117 tok/s).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@titaiwangms titaiwangms titaiwangms approved these changes
Copilot code review Copilot Awaiting requested review from Copilot

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@justinchuby @titaiwangms