Fix Gemma4 weight mapping for 26b MoE and 31b alternative attention - #237
Closed
justinchuby wants to merge 5 commits into
Closed
Fix Gemma4 weight mapping for 26b MoE and 31b alternative attention#237justinchuby wants to merge 5 commits into
justinchuby wants to merge 5 commits into
Conversation
When attention_k_eq_v=True (26b-a4b, 31b models), full_attention layers derive V from K (no separate v_proj), matching HF Gemma4TextAttention. Changes: - Add attention_k_eq_v and num_global_key_value_heads to Gemma4Config - Skip v_proj creation for alternative-attention layers - In forward: V = raw k_proj output (before k_norm/RoPE), then v_norm - Per-layer KV head count: full_attention uses num_global_key_value_heads - Update task KV cache to use per-layer head counts Tests: 2680 passed (12 gemma4), lint clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4Model (multimodal) has its own preprocess_weights that was missing the expert weight rename (gate_up_proj → fc1_experts_weights) and router scale folding. These were only in Gemma4CausalLMModel.preprocess_weights. Added both transformations to the multimodal model's preprocess_weights, fixing 26b-a4b and 26b-a4b-it exports. 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! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes Gemma4 HuggingFace weight/graph mapping for variants that enable alternative attention (attention_k_eq_v=True) and for multimodal Gemma4 MoE checkpoints, aligning Mobius’ ONNX module structure with HF’s parameterization.
Changes:
- Implement alternative-attention behavior for full-attention layers (V derived from raw K projection, no
v_proj, and per-layer KV head counts). - Add missing MoE expert weight renames and router scale folding to
Gemma4Model.preprocess_weights(multimodal path). - Extend
Gemma4Configwithattention_k_eq_vand plumb it throughfrom_transformers.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/mobius/tasks/_gemma4.py |
Adjusts KV-cache input shapes to account for alternative-attention KV head counts. |
src/mobius/models/gemma4.py |
Implements V=K alternative attention (no v_proj) and adds MoE weight mapping + router scale folding for multimodal Gemma4. |
src/mobius/_configs.py |
Adds attention_k_eq_v to Gemma4Config and extracts it from HF configs. |
The need_fallback check only considered KV-shared layers but not layers where head_dim exceeds CUDA GQA's max (256). Full-attention layers with global_head_dim=512 fell through to GQA fallback with an empty bias dict, causing KeyError: 'full_attention'. Fix: include head_dim > 256 in the need_fallback condition. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Latest ORT supports head_dim=512 in GroupQueryAttention. Remove the fallback that routed full-attention layers (global_head_dim=512) to standard Attention instead of GQA. All non-shared layers now use GQA on CUDA EP regardless of head_dim. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Decouple two independent Gemma4 features: - num_global_key_value_heads: controls KV head count for full-attention layers. Now used whenever set (not None), regardless of attention_k_eq_v. - attention_k_eq_v: controls V=K sharing (no v_proj). Gated separately. Previously both were coupled: num_global_key_value_heads only took effect when attention_k_eq_v was True. This was incorrect — a model could have different KV head counts for global/local layers without sharing V=K. Changes: - gemma4.py: Use num_global_key_value_heads for full-attention layers whenever it's not None - _gemma4.py: Same decoupling in _make_gemma4_kv_cache_inputs - build_graph_test.py: Add test with k_eq_v=True and num_global_key_value_heads != num_key_value_heads, verifying no v_proj and correct KV cache shapes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Member
Author
|
Replaced by #239 |
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
Two weight mapping fixes for Gemma4 models that use
attention_k_eq_v=True(26b-a4b, 31b):Fix 1: Alternative attention (V=K) for full_attention layers
When
attention_k_eq_v=True, full-attention layers derive V from K (no separatev_proj), matching HFGemma4TextAttention. Changes:v_projcreation for alternative-attention layersk_projoutput (before k_norm/RoPE), then v_norm appliednum_global_key_value_headsFix 2: MoE expert weight mapping in Gemma4Model
Gemma4Model(multimodal) had its ownpreprocess_weightsthat was missing:experts.gate_up_proj→fc1_experts_weights)hidden_size^-0.5)These were only in
Gemma4CausalLMModel.preprocess_weights.Testing
Models affected