fix: GPT-OSS MXFP4 dequant + Gemma3 VL auto-routing - #194
Merged
Conversation
The openai/gpt-oss-20b checkpoint stores expert weights in MXFP4 format (_blocks + _scales tensors) instead of full weight tensors. Add a dequantization phase using HF's _convert_moe_packed_tensors before the existing split/transpose logic. Also make gate_up_proj and down_proj matching more precise with endswith() to avoid matching _blocks/_scales suffixes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR addresses two build-time correctness issues: (1) loading GPT-OSS checkpoints that store MoE expert weights in MXFP4 packed form, and (2) automatically routing Gemma3 configs with vision/audio sub-configs to the multimodal registry entry so the correct multi-model package is produced.
Changes:
- Add an MXFP4 dequantization phase in
GPTOSSCausalLMModel.preprocess_weights()that converts_blocks/_scalestensors into full expert weight tensors before existing split/transpose logic runs. - Tighten GPT-OSS expert weight key matching using
endswith()to avoid accidentally matching MXFP4 suffix keys. - Add multimodal auto-detection in
build()to promote{model_type}→{model_type}_multimodalwhen the HF config indicates vision/audio encoders and a corresponding registry key exists.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
src/mobius/models/gptoss.py |
Dequantizes MXFP4-packed MoE expert tensors and refines key matching to prevent suffix collisions. |
src/mobius/_builder.py |
Promotes HF model_type to a registered multimodal variant based on vision_config/audio_config presence. |
justinchuby
force-pushed
the
justinchu/gpt-oss
branch
2 times, most recently
from
April 23, 2026 17:57
2191d22 to
086710b
Compare
All model_type=gemma3 HF models (4B/12B/27B) are multimodal with vision_config. Only gemma3_text (1B) is text-only. Update the registry to map gemma3 → Gemma3MultiModalModel with vision-language task, matching HF's type assignment. Remove the gemma3_multimodal alias — gemma3 now IS the multimodal entry. Update test configs and references accordingly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
justinchu/gpt-oss
branch
from
April 23, 2026 17:58
086710b to
60ecea2
Compare
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
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
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
Two fixes on this branch:
1. GPT-OSS MXFP4 dequantization
The
openai/gpt-oss-20bcheckpoint stores expert weights in MXFP4 format (_blocks+_scalesuint8 tensors) instead of full weight tensors. This caused 2304 missing expert weights during build.Fix: Added a dequantization phase in
preprocess_weightsusing HF's_convert_moe_packed_tensors(4-bit nibble-packed with shared exponent). The dequantized tensors are then handled by the existing split/transpose logic.Also made
gate_up_projanddown_projmatching more precise withendswith()to avoid matching_blocks/_scalessuffixes.2. Gemma3 VL auto-routing
google/gemma-3-4b-ithasmodel_type="gemma3"but is actually multimodal (Gemma3ForConditionalGenerationwithvision_config). The registry mapsgemma3→ text-only task, producing only 1 model instead of 3.Fix: Added multimodal auto-detection in
_builder.py. When a HF config hasvision_configoraudio_configand a{model_type}_multimodalregistry key exists,model_typeis promoted to the multimodal variant. This routes gemma3 → gemma3_multimodal → 3-model VL output.Tests
All 2555 tests pass.