Add Gemma4 GGUF support with per-layer KV heads fix - #234
Merged
Conversation
Fix the GGUF architecture mapping so gemma4 GGUF files build using the text-only task (gemma4_text) instead of the multimodal task, matching the gemma3 pattern. GGUF files only contain the text backbone weights. Add Gemma4-specific config postprocessing to extract dual-regime fields from GGUF metadata: - Dual head_dim: key_length_swa (sliding) vs key_length (global) - Dual RoPE theta: freq_base_swa (sliding) vs freq_base (global) - Layer types from sliding_window_pattern bool array - Sliding window, softcapping, KV sharing, per-layer input gating - Per-layer KV head arrays collapsed to majority value Returns Gemma4Config (not plain ArchitectureConfig) so the model receives all architecture-specific fields. Tests: - 13 unit tests for Gemma4 config mapping (dual head_dim, dual RoPE, layer types, softcapping, KV heads, per-layer input, etc.) - 1 end-to-end test: build_from_gguf with synthetic Gemma4 GGUF file Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4 31B uses different KV head counts per layer type: 16 for
sliding-window layers and 4 for full-attention layers (HF config:
num_global_key_value_heads=4, attention_k_eq_v=True).
The model was previously using a uniform num_key_value_heads=16 for
all layers, causing k_proj/v_proj shape mismatches on full-attention
layers: expected [8192, 5376] but weights are [2048, 5376].
Changes:
- Add num_global_key_value_heads field to Gemma4Config
- Extract from HF config (attention_k_eq_v + num_global_key_value_heads)
- Extract from GGUF per-layer KV head array (minority value)
- Use per-layer-type KV heads in Gemma4TextAttention and KV cache
- Verified end-to-end: build_from_gguf('unsloth/gemma-4-31B-it-GGUF')
succeeds with 833 tensors loaded, 1962 graph nodes
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
…otary_factor When a GGUF file lacks the non-standard key_length_swa and freq_base_swa keys (unsloth extensions), emit a warning so users know the sliding-window head_dim and rope_theta may be incorrect. Also add a comment citing the HF Gemma4Config source for the hardcoded global_partial_rotary_factor = 0.25 default. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds GGUF import support for Gemma4 text-only models and fixes Gemma4’s per-layer KV head handling so full-attention layers use the correct (global) KV head count.
Changes:
- Remap GGUF
gemma4architecture togemma4_textand add Gemma4-specific GGUF → config postprocessing. - Fix per-layer KV head selection for full-attention layers across Gemma4 attention modules and KV-cache task inputs.
- Add synthetic GGUF-based unit/e2e tests for Gemma4 config extraction and
build_from_gguf().
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/tasks/_gemma4.py |
Use num_global_key_value_heads for full-attention KV-cache input shapes. |
src/mobius/models/gemma4.py |
Set per-layer num_key_value_heads based on layer type to match HF Gemma4 semantics. |
src/mobius/integrations/gguf/_config_mapping.py |
Map gemma4 → gemma4_text, collapse per-layer KV head arrays, and add Gemma4 postprocessor to build Gemma4Config. |
src/mobius/integrations/gguf/_reader_test.py |
Add synthetic Gemma4 GGUF writer + tests for Gemma4 config mapping and build_from_gguf(). |
src/mobius/_configs.py |
Add num_global_key_value_heads to Gemma4Config and populate it from HF configs when applicable. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Address PR review comments: - Validate sliding_window_pattern length matches num_hidden_layers - Validate head_count_kv array length matches sliding_window_pattern - Add layer_types length check in Gemma4TextDecoderLayer.__init__ Clear ValueError messages on malformed GGUF metadata instead of cryptic IndexError at runtime. Signed-off-by: Justin Chu <justinchu@microsoft.com>
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
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
Add GGUF support for Gemma4 (text-only) and fix a pre-existing bug where full-attention layers used wrong KV head counts.
Changes
1. GGUF model_type mapping fix (1-line)
gemma4→gemma4_textinGGUF_ARCH_TO_MODEL_TYPE— GGUF only has text weights, not multimodal. Follows the gemma3 → gemma3_text pattern.2. Gemma4 GGUF config extraction
_gemma4_postprocess()in_config_mapping.py— extracts dual head_dim, dual RoPE theta, layer_types, sliding_window, softcapping, KV sharing, and per-layer-input gating from GGUF metadata.key_length_swa,freq_base_swa) with warnings when missing.num_global_key_value_headsfrom per-layer KV head arrays.3. Per-layer KV heads bug fix
num_key_value_heads=16(sliding) instead ofnum_global_key_value_heads=4(full-attention).Gemma4TextAttention,Gemma4Config,Gemma4TextCausalLMTaskKV cache, and GGUF config extraction.4. Tests
build_from_gguf()with synthetic GGUFVerified
unsloth/gemma-4-31B-it-UD-IQ2_XXS.gguf→ 833 tensors, 1962 nodes, SUCCESSbuild_from_gguf()correctly produces single text-only model (not multimodal)