From 5f6d26d7e1bc0d64831d480ba3d587a5a809f39f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 4 May 2026 18:05:37 +0000 Subject: [PATCH 1/2] Address PR #239 follow-up review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Update GQA rewrite rule MAX_HEAD_DIM from 256 to 512 — latest ORT supports head_dim=512 for Gemma4 full-attention layers 2. Set attention_k_eq_v in GGUF postprocessor when num_global_key_value_heads differs from num_key_value_heads, fixing build_from_gguf for 26b/31b 3. Extract _remap_moe_expert_weights() shared helper to avoid duplicating expert rename + router scale folding in CausalLM and multimodal paths 4. Add targeted preprocess_weights tests for both Gemma4CausalLMModel and Gemma4Model (multimodal) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../integrations/gguf/_config_mapping.py | 3 + src/mobius/models/gemma4.py | 75 ++++++----- src/mobius/models/gemma4_test.py | 119 ++++++++++++++++++ .../rewrite_rules/_group_query_attention.py | 11 +- 4 files changed, 164 insertions(+), 44 deletions(-) create mode 100644 src/mobius/models/gemma4_test.py diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index 22babb5d..a18adc69 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -465,6 +465,9 @@ def _gemma4_postprocess( else 1_000_000.0, global_partial_rotary_factor=global_partial_rotary_factor, num_global_key_value_heads=num_global_key_value_heads, + # attention_k_eq_v: derive from per-layer KV head counts. When + # full-attention layers use fewer KV heads, V = K (no v_proj). + attention_k_eq_v=num_global_key_value_heads is not None, final_logit_softcapping=float(final_logit_softcapping or 0.0), attn_logit_softcapping=float(attn_logit_softcapping or 0.0), num_kv_shared_layers=int(num_kv_shared_layers) diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 3a341e27..fb1baf97 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -55,6 +55,39 @@ if TYPE_CHECKING: from mobius.components._attention import GQAContext + +# --------------------------------------------------------------------------- +# Shared weight preprocessing helpers +# --------------------------------------------------------------------------- + + +def _remap_moe_expert_weights( + state_dict: dict[str, torch.Tensor], + config: Gemma4Config, +) -> None: + """Rename HF MoE expert weights and fold router scale in-place. + + Shared by ``Gemma4CausalLMModel`` and ``Gemma4Model`` to avoid + duplicating the rename/fold logic. + """ + # experts.gate_up_proj → fc1_experts_weights + # experts.down_proj → fc2_experts_weights + for key in list(state_dict.keys()): + if ".experts.gate_up_proj" in key: + new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") + state_dict[new_key] = state_dict.pop(key) + elif ".experts.down_proj" in key: + new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") + state_dict[new_key] = state_dict.pop(key) + + # Fold hidden_size^-0.5 into router.scale + if config.enable_moe_block: + scale_factor = float(config.hidden_size**-0.5) + for key in list(state_dict.keys()): + if ".router.scale" in key and ".per_expert_scale" not in key: + state_dict[key] = state_dict[key] * scale_factor + + # --------------------------------------------------------------------------- # Scale-free RMSNorm (Gemma4RMSNorm with with_scale=False) # --------------------------------------------------------------------------- @@ -1653,27 +1686,8 @@ def preprocess_weights( for i in range(num_layers): shard = value[:, i * per_layer_dim : (i + 1) * per_layer_dim] state_dict[f"model.embed_tokens_per_layer.{i}.weight"] = shard - # Map HF expert weight names to our 3D stacked parameter names. - # HF stores: layers.N.experts.gate_up_proj [E, 2*inter, H] - # layers.N.experts.down_proj [E, H, inter] - # We store: layers.N.fc1_experts_weights [E, 2*inter, H] - # layers.N.fc2_experts_weights [E, H, inter] - for key in list(state_dict.keys()): - if ".experts.gate_up_proj" in key: - new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") - state_dict[new_key] = state_dict.pop(key) - elif ".experts.down_proj" in key: - new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") - state_dict[new_key] = state_dict.pop(key) - # Fold hidden_size^-0.5 into router.scale. - # The router computes: x_normed * scale * hidden_size^-0.5. - # We pre-multiply scale by hidden_size^-0.5 here so the forward only needs - # x_normed * self.scale, avoiding float-constant name collisions across layers. - if self.config.enable_moe_block: - scale_factor = float(self.config.hidden_size**-0.5) - for key in list(state_dict.keys()): - if ".router.scale" in key: - state_dict[key] = state_dict[key] * scale_factor + # Map HF expert weight names and fold router scale + _remap_moe_expert_weights(state_dict, self.config) return super().preprocess_weights(state_dict) @@ -2150,22 +2164,7 @@ def preprocess_weights( else: renamed[key] = value - # Map HF expert weight names to our 3D stacked parameter names. - # HF: decoder.model.layers.N.experts.gate_up_proj [E, 2*inter, H] - # Us: decoder.model.layers.N.fc1_experts_weights [E, 2*inter, H] - for key in list(renamed.keys()): - if ".experts.gate_up_proj" in key: - new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") - renamed[new_key] = renamed.pop(key) - elif ".experts.down_proj" in key: - new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") - renamed[new_key] = renamed.pop(key) - - # Fold hidden_size^-0.5 into router.scale - if self.config.enable_moe_block: - scale_factor = float(self.config.hidden_size**-0.5) - for key in list(renamed.keys()): - if ".router.scale" in key and ".per_expert_scale" not in key: - renamed[key] = renamed[key] * scale_factor + # Map HF expert weight names and fold router scale + _remap_moe_expert_weights(renamed, self.config) return renamed diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py new file mode 100644 index 00000000..4e47422b --- /dev/null +++ b/src/mobius/models/gemma4_test.py @@ -0,0 +1,119 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for Gemma4 preprocess_weights — MoE expert rename and router scale.""" + +from __future__ import annotations + +import torch + +from mobius._configs import Gemma4Config +from mobius.models.gemma4 import Gemma4CausalLMModel, Gemma4Model + + +def _tiny_gemma4_config(**overrides) -> Gemma4Config: + """Create a minimal Gemma4Config for preprocess_weights tests.""" + from mobius._configs import VisionConfig + + defaults = dict( + model_type="gemma4", + vocab_size=256, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + intermediate_size=128, + head_dim=16, + global_head_dim=32, + hidden_act="gelu", + enable_moe_block=True, + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=32, + layer_types=["sliding_attention", "full_attention"], + attention_k_eq_v=True, + num_global_key_value_heads=1, + vision=VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + image_size=16, + patch_size=4, + ), + ) + defaults.update(overrides) + return Gemma4Config(**defaults) + + +class TestGemma4CausalLMPreprocessWeights: + """Test Gemma4CausalLMModel.preprocess_weights.""" + + def test_expert_weight_rename(self): + config = _tiny_gemma4_config() + model = Gemma4CausalLMModel(config) + + fake_sd = { + "model.layers.0.experts.gate_up_proj": torch.zeros(4, 64, 64), + "model.layers.0.experts.down_proj": torch.zeros(4, 64, 32), + } + result = model.preprocess_weights(fake_sd) + + assert "model.layers.0.fc1_experts_weights" in result + assert "model.layers.0.fc2_experts_weights" in result + assert "model.layers.0.experts.gate_up_proj" not in result + + def test_router_scale_folding(self): + config = _tiny_gemma4_config() + model = Gemma4CausalLMModel(config) + + scale_val = torch.tensor([2.0]) + fake_sd = {"model.layers.0.router.scale": scale_val.clone()} + result = model.preprocess_weights(fake_sd) + + expected = 2.0 * (64**-0.5) # hidden_size=64 + assert abs(result["model.layers.0.router.scale"].item() - expected) < 1e-6 + + +class TestGemma4ModelPreprocessWeights: + """Test Gemma4Model.preprocess_weights (multimodal path).""" + + def test_expert_weight_rename(self): + config = _tiny_gemma4_config() + model = Gemma4Model(config) + + fake_sd = { + "model.language_model.layers.0.experts.gate_up_proj": torch.zeros(4, 64, 64), + "model.language_model.layers.0.experts.down_proj": torch.zeros(4, 64, 32), + } + result = model.preprocess_weights(fake_sd) + + assert "decoder.model.layers.0.fc1_experts_weights" in result + assert "decoder.model.layers.0.fc2_experts_weights" in result + + def test_router_scale_folding(self): + config = _tiny_gemma4_config() + model = Gemma4Model(config) + + scale_val = torch.tensor([2.0]) + fake_sd = {"model.language_model.layers.0.router.scale": scale_val.clone()} + result = model.preprocess_weights(fake_sd) + + expected = 2.0 * (64**-0.5) + key = "decoder.model.layers.0.router.scale" + assert key in result + assert abs(result[key].item() - expected) < 1e-6 + + def test_per_expert_scale_not_folded(self): + """router.per_expert_scale should NOT be multiplied by scale_factor.""" + config = _tiny_gemma4_config() + model = Gemma4Model(config) + + fake_sd = { + "model.language_model.layers.0.router.per_expert_scale": torch.ones(4), + } + result = model.preprocess_weights(fake_sd) + + key = "decoder.model.layers.0.router.per_expert_scale" + assert key in result + assert torch.allclose(result[key], torch.ones(4)) diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index 4bdcf08c..2416b71b 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/src/mobius/rewrite_rules/_group_query_attention.py @@ -49,12 +49,11 @@ RewriteRuleSet, ) -# CUDA EP's GroupQueryAttention kernel enforces MAX_HEAD_SIZE = 256. -# Rewriting Attention → GQA for layers with head_dim > 256 would produce -# a node that fails at runtime. We skip those nodes so they stay as -# standard Attention ops (which can use the MHA code-path when -# q_num_heads == kv_num_heads after KV expansion). -_MAX_GQA_HEAD_DIM = 256 +# CUDA EP's GroupQueryAttention kernel historically enforced MAX_HEAD_SIZE = 256. +# Latest ORT (unreleased) supports head_dim=512 for Gemma4 full-attention layers. +# The rewrite rule (Attention → GQA) still skips head_dim > 512 to avoid +# generating nodes that would fail on older runtimes. +_MAX_GQA_HEAD_DIM = 512 def _head_dim_exceeds_gqa_limit(past_key) -> int | None: From 8cee3d6db545881246443b362e4742325cc68ad7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 4 May 2026 11:11:42 -0700 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/rewrite_rules/_group_query_attention.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index 2416b71b..476cd8c0 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/src/mobius/rewrite_rules/_group_query_attention.py @@ -50,10 +50,10 @@ ) # CUDA EP's GroupQueryAttention kernel historically enforced MAX_HEAD_SIZE = 256. -# Latest ORT (unreleased) supports head_dim=512 for Gemma4 full-attention layers. -# The rewrite rule (Attention → GQA) still skips head_dim > 512 to avoid -# generating nodes that would fail on older runtimes. -_MAX_GQA_HEAD_DIM = 512 +# Keep the rewrite-rule limit conservative until GQA fusion is gated on a +# runtime/EP capability check. This avoids emitting GroupQueryAttention nodes +# with head_dim=512 that can still fail on released ORT builds. +_MAX_GQA_HEAD_DIM = 256 def _head_dim_exceeds_gqa_limit(past_key) -> int | None: