Fix vision encoder for 26b/31b: respect use_clipped_linears config - #277
Merged
Conversation
Gemma4 26b-A4B and 31b models set use_clipped_linears=False in their vision config, meaning they don't use quantization-aware clipping bounds. Previously, the vision encoder always used ClippableLinear (which creates input_min/max/output_min/max parameters), causing 756 missing weight errors when loading these models. Now Gemma4VisionSelfAttention and Gemma4VisionEncoderLayer accept a use_clipped_linears parameter that selects between ClippableLinear (e2b/e4b models with QAT) and standard Linear (26b/31b without QAT). The flag is read from VisionConfig.use_clipped_linears which is already populated from the HF config. Tested with google/gemma-4-26b-a4b: all 3 submodels (decoder, vision_encoder, embedding) now export and save successfully. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
gemma4 (gemma4) / vision_encoder — 52 change(s)Op summary: 190 → 176 nodes --- base
+++ head
@@ -33,18 +33,12 @@
Where
Unsqueeze
RMSNormalization
-Clip
Transpose
MatMul
-Clip
-Clip
Transpose
MatMul
-Clip
-Clip
Transpose
MatMul
-Clip
Reshape
Reshape
Reshape
@@ -112,27 +106,19 @@
MatMul
Transpose
Reshape
-Clip
Transpose
MatMul
-Clip
RMSNormalization
Add
RMSNormalization
-Clip
Transpose
MatMul
-Clip
Gelu
-Clip
Transpose
MatMul
-Clip
Mul
-Clip
Transpose
MatMul
-Clip
RMSNormalization
Add
UnsqueezeRemoved nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
37 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Gemma4 26b-A4B and 31b models fail to export with 756 missing weight errors:
Root Cause
The vision encoder always used
ClippableLinearfor all projections, which createsinput_min/input_max/output_min/output_maxparameters. These QAT (quantization-aware training) clipping bounds exist in the HF checkpoint for e2b/e4b models but not for 26b/31b.The HF config has
use_clipped_linears:True(224 vision clipping weights in checkpoint)False(0 clipping weights)False(0 clipping weights)Our
VisionConfigalready reads this field, but the model code ignored it.Fix
Gemma4VisionSelfAttentionandGemma4VisionEncoderLayernow acceptuse_clipped_linearsand selectClippableLinearvsLinearaccordingly. The flag is passed fromVisionConfigthrough the encoder core.Testing
mobius build --model google/gemma-4-26b-a4bexports all 3 submodels successfullyTrue)