Fix BF16 vision encoder: keep OneHot values as float32 - #252
Merged
Conversation
ONNX OneHot op doesn't support bfloat16 for the values tensor. Keep OneHot values as float32 and cast the output to model dtype afterward for the subsequent MatMul. Previously, CastLike cast the values to bf16 before OneHot, causing the entire vision encoder to fail at model load time. Tested: bf16 build succeeds for google/gemma-4-e2b-it (all 4 models). Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
🏗️ Architecture Diff
gemma4 (gemma4) / vision_encoder — 2 change(s)Op summary: 202 → 202 nodes --- base
+++ head
@@ -181,8 +181,8 @@
Mul
Add
Constant
-CastLike
OneHot
+CastLike
Transpose
MatMul
ConstantAdded nodes:
Removed nodes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
titaiwangms
approved these changes
May 5, 2026
titaiwangms
left a comment
titaiwangms
left a comment
Contributor
There was a problem hiding this comment.
Did copilot miss op spec when it's designing model? Maybe we need to re-enforce into SKILL.md about op spec?
Member
Author
Onehot does support it. I think it is missing support in ort: https://onnx.ai/onnx/operators/onnx__OneHot.html |
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
2 participants
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.
Summary
ONNX OneHot op does not support bfloat16 for the values tensor. When building Gemma4 with
--dtype bf16, the vision encoder fails to load becauseCastLikecasts the OneHot values to bf16 before the op.Fix
Keep OneHot values as float32 and cast the output to model dtype afterward for the subsequent MatMul.
Testing
mobius build --model google/gemma-4-e2b-it --dtype bf16 --no-weights