Add 2-bit (uint2) Gemma 4 MoE support to the model builder - #2482
Open
Thiago Pereira Rocha (thpereir) wants to merge 5 commits into
Open
Add 2-bit (uint2) Gemma 4 MoE support to the model builder#2482Thiago Pereira Rocha (thpereir) wants to merge 5 commits into
Thiago Pereira Rocha (thpereir) wants to merge 5 commits into
Conversation
Adds Gemma4Model builder for the text component of Gemma4Unified (Gemma4UnifiedForConditionalGeneration). Verified against HuggingFace: full 48-layer fp32 parity gives logit correlation 1.0 and exact top-1 match. Handles the architecture's deviations from Gemma3: - Per-layer head dims (sliding 256 / full 512) and KV heads (8 / 1), emitted as concrete per-layer KV cache shapes so the runtime's DefaultKeyValueCache auto-detects the heterogeneous layout. - attention_k_eq_v: full-attention layers share the K projection as V (no v_proj), plus a scaleless value RMSNorm (v_norm) on all layers. - Proportional RoPE on full layers (partial rotary on the global head dim with a zero-padded NoPE tail), default RoPE on sliding layers. - Per-layer layer_scalar residual multiplier. - No-offset RMSNorm (weight used directly, unlike Gemma1/2/3). - Text-only weight loader that remaps the model.language_model.* prefix and avoids loading the vision/audio towers. Emits model_type "gemma4_text" to match the runtime's model type list. Co-Authored-By: Claude <noreply@anthropic.com>
Add Gemma4MoEModel subclassing the dense Gemma4Model. Every layer runs a parallel dense MLP and a fused-QMoE expert block, combined by Add. Uses the fused QMoE op path (gelu + swiglu_fusion=1 + normalize_routing_weights) with an explicit router pre-projection subgraph feeding raw logits, and folds per_expert_scale into the expert down_proj weights offline. Co-Authored-By: Claude <noreply@anthropic.com>
…nature) Co-Authored-By: Claude <noreply@anthropic.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
3 tasks
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
…tub test - make_moe_init: set both moe_attrs["op_type"] and ["moe_op_type"] to the resolved op. The shared emitters read "op_type" while the gemma parallel-FFN path reads "moe_op_type"; setting only one left quantized experts on the float "MoE" op, which rejects uint8 weights. - Gemma4MoEModel.__init__: hide config.moe_intermediate_size across super().__init__ so the base sets self.intermediate_size to the dense size directly, instead of reassigning it afterward (CodeQL "overwriting attribute"). - test_qmoe_weights.py: add Gemma4Model/Gemma4MoEModel to the stubbed `builders` class list so the builder-CLI import test matches builder.py's imports. Co-Authored-By: Claude <noreply@anthropic.com>
Build the fused Quark uint2 gemma-4-26B-A4B-it checkpoint to ONNX with a pre-quantized expert path: consume group-wise uint2 weights/scales/float zero-points directly, fold the router per-expert scale into the down-proj scales, and emit the shared gate/up prescale+rotation once before QMoE. - builder.py: add int2 precision plumbing - quantized_model.py: extend the Quark loader/repack to uint2 (group_size 64, MSB-first packing) and to the split gemma4 experts with bundled dense LoRA - gemma.py: pre-quantized uint2 expert branch in make_moe; emit block_size=64; omit zero_points on CUDA int2 (op reconstructs the symmetric -1.5*scale bias); keep float zero-points on CPU and interleave fc1 gate/up rows for the CPU kernel's swiglu_fusion=1 layout
Thiago Pereira Rocha (thpereir)
force-pushed
the
gemma4-2bit-upstream
branch
from
August 27, 2026 21:08
4fbcb5c to
c0e2355
Compare
Thiago Pereira Rocha (thpereir)
marked this pull request as ready for review
August 28, 2026 14:28
Thiago Pereira Rocha (thpereir)
requested a review
from a team
as a code owner
August 28, 2026 14:28
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Copilot started reviewing on behalf of
Thiago Pereira Rocha (thpereir)
August 28, 2026 14:29
View session
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Gemma 4 text-model builder support with a new 2-bit (uint2) pre-quantized Quark MoE path, including loader-side repacking/requantization utilities, builder-side QMoE graph emission (with factored input rotation + optional zero-point omission on CUDA), and new unit tests validating structural/parity expectations.
Changes:
- Add
--precision int2plumbing and quant-config mapping to drive uint2 weight/MoE quantization selection. - Extend Quark loader + base loader to support uint2 packing/unpacking, split-expert refusion, shared input rotations, and bundled dense LoRA adapter baking.
- Add Gemma4/Gemma4MoE builders and tests validating RoPE cache parity, per-layer KV geometry, and QMoE uint2 behaviors (e.g.,
expert_weight_bits=2, CPU vs CUDA zero-point handling).
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| test/python/models/test_gemma4moe_builder.py | Adds MoE-focused unit tests including the new Quark uint2 QMoE path expectations. |
| test/python/models/test_gemma4_builder.py | Adds dense Gemma4 structural/parity unit tests (RoPE caches, KV-cache shapes, attention scale). |
| test/python/builder/test_qmoe_weights.py | Registers Gemma4 builders in CLI-module loading for existing QMoE-related tests. |
| src/python/py/models/README.md | Documents -p int2 usage for pre-quantized 2-bit MoE. |
| src/python/py/models/quantization/quant_config.py | Adds uint2/int2 dtype descriptors and maps --precision int2 to weights.type=uint2 / MoE quant type. |
| src/python/py/models/loaders/quark.py | Extends Quark loader for uint2 expert handling, unpack/repack, and 2-bit-aware packing utilities. |
| src/python/py/models/loaders/base.py | Adds Gemma4-MoE layer/router structures, shared-input-rotation loading, LoRA adapter loading, and packed-splitting fixes for uint2. |
| src/python/py/models/builders/gemma.py | Implements Gemma4 + Gemma4MoE builders, including QMoE emission for pre-quantized uint2 experts and per-layer geometry/RoPE behavior. |
| src/python/py/models/builders/base.py | Threads true precision for int2, fixes MoE op-type wiring, adds factored-rotation + LoRA delta injection for quantized projections. |
| src/python/py/models/builders/init.py | Exports Gemma4Model/Gemma4MoEModel. |
| src/python/py/models/builder.py | Adds int2 precision option, dtype handling, supported MoE quant types, and Gemma4 architecture dispatch. |
Suppressed comments (2)
src/python/py/models/loaders/base.py:1064
_packed_out_factoris declared as a@staticmethod, but it is only used viaself._packed_out_factor(...). Making it a normal instance method avoids@staticmethod(which this codebase generally avoids outside of specific constructors) and keeps call sites unchanged.
@staticmethod
def _packed_out_factor(tensor, bits):
src/python/py/models/loaders/quark.py:212
_unpack_uint2_msb_firstis introduced as a@staticmethod, but it’s only called asself._unpack_uint2_msb_first(...). Converting it to an instance method avoids@staticmethodusage while keeping call sites unchanged.
@staticmethod
def _unpack_uint2_msb_first(packed):
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+273
to
+274
| # Factored online rotation (Quark rotation algo): x_rot = (x / input_prescale) @ shared_input_rotation_<in>. | ||
| # `shared_input_rotations` maps in_features -> [in, in] rotation matrix (shared across all layers). |
Comment on lines
+555
to
+560
| def make_layer(self, layer_id, layer): | ||
| # Stash the full layer so make_mlp can reach the router/experts/extra norms | ||
| # (the parent only passes layer.mlp to make_mlp). | ||
| self._current_layer = layer | ||
| super().make_layer(layer_id, layer) | ||
| self._current_layer = None |
Comment on lines
+18
to
+20
| # uint2/int2 pack 4 codes per byte (MSB-first); uint4/int4 pack 2 codes per byte. | ||
| _DTYPE_BITS = {"uint4": 4, "int4": 4, "uint2": 2, "int2": 2} | ||
|
|
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
Adds a pre-quantized 2-bit (uint2) expert path for the Gemma 4 MoE builder (
gemma-4-26B-A4B-it), consuming a fused Quark uint2 checkpoint directly into aQMoEgraph.int2precision plumbing (--precision int2,set_io_dtype/set_onnx_dtypeint2 as a FLOAT carrier,supported_moe_quant_types).make_moe_quark; emitblock_size, foldrouter.per_expert_scaleinto down-proj scales, emit the shared gate/up prescale+rotation once beforeQMoE; omit zero-points on CUDA int2 (op reconstructs the symmetric -1.5*scale bias), keep float zero-points on CPU and interleave fc1 gate/up rows for the CPU kernel'sswiglu_fusion=1layout.Stacking
This PR is stacked on top of the base Gemma 4 PR (#2473). Until #2473 merges, the diff here also shows the base-arch commits; the 2-bit-specific change is the top commit. Will rebase onto
mainonce #2473 lands.Test plan
test/python/models/test_gemma4moe_builder.py- addedTestGemma4MoEQuarkPath: QMoE emitsexpert_weight_bits=2, CPU emits float per-group zero-points as QMoE inputs, CUDA omits zero-points. (10 passed)ruff checkclean on the test file.Generated with Claude Code