Add text-backbone support for Gemma 4 12B and 12B-it

docs/models.md

@@ -11,6 +11,7 @@ Tunix supports the following models:
 | Gemma | 2B, 7B, 9B |
 | Gemma 2 | 2B, 9B |
 | Gemma 3 | 270M, 1B, 4B, 12B, 27B |
+| Gemma 4 | E2B, E4B, 12B, 26B-A4B, 31B |
 | Llama 3 | 70B, 405B |
 | Llama 3.1 | 8B, 70B, 405B |
 | Llama 3.2 | 1B, 3B |

tests/models/automodel_test.py

@@ -33,6 +33,12 @@ def _get_all_models_test_parameters():
       dict(testcase_name="gemma-3-12b-it", model_name="gemma-3-12b-it"),
       dict(testcase_name="gemma-3-27b-pt", model_name="gemma-3-27b-pt"),
       dict(testcase_name="gemma-3-27b-it", model_name="gemma-3-27b-it"),
+      dict(testcase_name="gemma-4-e2b", model_name="gemma-4-e2b"),
+      dict(testcase_name="gemma-4-e4b", model_name="gemma-4-e4b"),
+      dict(testcase_name="gemma-4-12b", model_name="gemma-4-12b"),
+      dict(testcase_name="gemma-4-12b-it", model_name="gemma-4-12b-it"),
+      dict(testcase_name="gemma-4-31b", model_name="gemma-4-31b"),
+      dict(testcase_name="gemma-4-26b-a4b", model_name="gemma-4-26b-a4b"),
       dict(testcase_name="llama-3-70b", model_name="llama-3-70b"),
       dict(testcase_name="llama-3.1-70b", model_name="llama-3.1-70b"),
       dict(testcase_name="llama-3.1-405b", model_name="llama-3.1-405b"),
@@ -260,7 +266,13 @@ def test_create_model_dynamically(
         mode="mode",
     )
 
-    if naming_info.model_family in ("gemma", "gemma1p1", "gemma2", "gemma3"):
+    if naming_info.model_family in (
+        "gemma",
+        "gemma1p1",
+        "gemma2",
+        "gemma3",
+        "gemma4",
+    ):
       expected_module_type = automodel.ModelModule.PARAMS_SAFETENSORS
     else:
       expected_module_type = automodel.ModelModule.PARAMS

tests/models/gemma4/model_test.py

@@ -25,6 +25,45 @@
 
 class ModelTest(absltest.TestCase):
 
+  def test_gemma4_12b_config(self):
+    config = model_lib.ModelConfig.gemma4_12b()
+
+    self.assertEqual(config.num_layers, 48)
+    self.assertEqual(config.num_embed, 262144)
+    self.assertEqual(config.embed_dim, 3840)
+    self.assertEqual(config.hidden_dim, 15360)
+    self.assertEqual(config.num_heads, 16)
+    self.assertEqual(config.head_dim, 256)
+    self.assertEqual(config.num_kv_heads, 8)
+    self.assertEqual(config.num_global_kv_heads, 1)
+    self.assertEqual(config.global_key_size, 512)
+    self.assertEqual(config.sliding_window_size, 1024)
+    self.assertTrue(config.k_eq_v_global)
+    self.assertEqual(config.per_layer_input_dim, 0)
+    self.assertEqual(
+        config.attention_pattern,
+        (
+            model_lib.AttentionType.LOCAL_SLIDING,
+            model_lib.AttentionType.LOCAL_SLIDING,
+            model_lib.AttentionType.LOCAL_SLIDING,
+            model_lib.AttentionType.LOCAL_SLIDING,
+            model_lib.AttentionType.LOCAL_SLIDING,
+            model_lib.AttentionType.GLOBAL,
+        ),
+    )
+
+  def test_gemma4_12b_it_config_matches_base(self):
+    config = model_lib.ModelConfig.gemma4_12b()
+    it_config = model_lib.ModelConfig.gemma4_12b_it()
+
+    self.assertEqual(it_config.num_layers, config.num_layers)
+    self.assertEqual(it_config.embed_dim, config.embed_dim)
+    self.assertEqual(it_config.hidden_dim, config.hidden_dim)
+    self.assertEqual(it_config.num_heads, config.num_heads)
+    self.assertEqual(it_config.num_kv_heads, config.num_kv_heads)
+    self.assertEqual(it_config.num_global_kv_heads, config.num_global_kv_heads)
+    self.assertEqual(it_config.attention_pattern, config.attention_pattern)
+
   def test_forward_pass_dense(self):
     config = model_lib.ModelConfig.gemma4_e2b()
     config.num_layers = 1
@@ -81,6 +120,33 @@ def test_forward_pass_moe(self):
 
     self.assertEqual(logits.shape, (2, 32, config.num_embed))
 
+  def test_forward_pass_gemma4_12b(self):
+    config = model_lib.ModelConfig.gemma4_12b()
+    config.num_layers = 6
+    config.num_embed = 128
+    config.embed_dim = 256
+    config.hidden_dim = 512
+    config.num_heads = 4
+    config.head_dim = 64
+    config.num_kv_heads = 2
+    config.num_global_kv_heads = 1
+
+    rngs = nnx.Rngs(0)
+    model = model_lib.Gemma4(config, rngs=rngs)
+
+    tokens = jax.random.randint(
+        jax.random.PRNGKey(0), (1, 8), 0, config.num_embed
+    )
+    positions = jnp.tile(
+        jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1)
+    )
+    attn_mask = jnp.tril(
+        jnp.ones((tokens.shape[1], tokens.shape[1]), dtype=jnp.bool_)
+    )[None, ...]
+    logits, _ = model(tokens, positions=positions, attention_mask=attn_mask)
+
+    self.assertEqual(logits.shape, (1, 8, config.num_embed))
+
   def test_remat_block(self):
     config = model_lib.ModelConfig.gemma4_e2b()
     config.num_layers = 1

tests/models/gemma4/params_safetensors_test.py

@@ -0,0 +1,78 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Gemma 4 safetensors parameter loading helpers."""
+
+from absl.testing import absltest
+import jax.numpy as jnp
+from tunix.models.gemma4 import model as model_lib
+from tunix.models.gemma4 import params_safetensors
+from tunix.utils import torch_utils
+
+
+class ParamsSafetensorsTest(absltest.TestCase):
+
+  def test_gemma4_12b_key_mapping_uses_local_and_global_shapes(self):
+    config = model_lib.ModelConfig.gemma4_12b()
+    config.num_layers = 6
+    config.embed_dim = 16
+    config.num_heads = 2
+    config.head_dim = 4
+    config.num_kv_heads = 2
+    config.num_global_kv_heads = 1
+    config.global_key_size = 8
+
+    key_mapping = params_safetensors._get_key_and_transform_mapping(config)
+
+    local_key, local_transform = torch_utils.torch_key_to_jax_key(
+        key_mapping, "model.language_model.layers.0.self_attn.k_proj.weight"
+    )
+    self.assertEqual(local_key, "tmp.layers.0.attn.k")
+    self.assertEqual(local_transform[1], (16, 2, 4))
+
+    global_key, global_transform = torch_utils.torch_key_to_jax_key(
+        key_mapping, "model.language_model.layers.5.self_attn.k_proj.weight"
+    )
+    self.assertEqual(global_key, "tmp.layers.5.attn.k")
+    self.assertEqual(global_transform[1], (16, 1, 8))
+
+  def test_gemma4_12b_preprocess_handles_local_and_global_attention(self):
+    config = model_lib.ModelConfig.gemma4_12b()
+    config.num_layers = 6
+    config.embed_dim = 16
+    config.num_heads = 2
+    config.head_dim = 4
+    config.num_kv_heads = 2
+    config.num_global_kv_heads = 1
+    config.global_key_size = 8
+
+    tensors = {
+        "tmp.layers.0.attn.q": jnp.zeros((16, 2, 4)),
+        "tmp.layers.0.attn.k": jnp.zeros((16, 2, 4)),
+        "tmp.layers.0.attn.v": jnp.ones((16, 2, 4)),
+        "tmp.layers.5.attn.q": jnp.zeros((16, 2, 8)),
+        "tmp.layers.5.attn.k": jnp.ones((16, 1, 8)),
+    }
+
+    out = params_safetensors._make_preprocess_fn(config)(tensors)
+
+    self.assertEqual(out["layers.0.attn.q_einsum.w"].shape, (2, 16, 4))
+    self.assertEqual(out["layers.0.attn.kv_einsum.w"].shape, (2, 2, 16, 4))
+    self.assertEqual(out["layers.5.attn.q_einsum.w"].shape, (2, 16, 8))
+    self.assertEqual(out["layers.5.attn.k_einsum.w"].shape, (1, 16, 8))
+    self.assertNotIn("layers.5.attn.kv_einsum.w", out)
+
+
+if __name__ == "__main__":
+  absltest.main()

tunix/models/gemma4/model.py

@@ -208,6 +208,33 @@ def gemma4_e4b(
         ),
     )
 
+  @classmethod
+  def gemma4_12b(
+      cls,
+      sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
+  ) -> 'ModelConfig':
+    return cls(
+        num_layers=48,
+        num_embed=262144,
+        embed_dim=3840,
+        hidden_dim=3840 * 4,
+        num_heads=16,
+        head_dim=256,
+        num_kv_heads=8,
+        num_global_kv_heads=1,
+        sliding_window_size=1024,
+        shd_config=sharding_config,
+        k_eq_v_global=True,
+        attention_pattern=GEMMA4_ATTENTION_PATTERN,
+    )
+
+  @classmethod
+  def gemma4_12b_it(
+      cls,
+      sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
+  ) -> 'ModelConfig':
+    return cls.gemma4_12b(sharding_config=sharding_config)
+
   @classmethod
   def gemma4_31b(
       cls,
@@ -922,7 +949,7 @@ def init_cache(self, batch_size, max_seq_len, dtype):
     k = shard(
         np.zeros(cache_shape, dtype),
         self.config.shd_config.act_btnh,
-        eager=True
+        eager=True,
     )
     v = shard(
         np.zeros(cache_shape, dtype),
@@ -1321,9 +1348,7 @@ def get_model_input(self):
     dummy_batch_size = 2
     dummy_seq_len = 2
     return {
-        'tokens': jnp.ones(
-            (dummy_batch_size, dummy_seq_len), dtype=jnp.int32
-        ),
+        'tokens': jnp.ones((dummy_batch_size, dummy_seq_len), dtype=jnp.int32),
         'positions': jnp.ones(
             (dummy_batch_size, dummy_seq_len), dtype=jnp.int32
         ),

tunix/models/registry.py

@@ -195,6 +195,22 @@
         model_config_id='gemma4_e4b',
         model_config_category='gemma4',
     ),
+    naming.ModelNaming(
+        model_id='google/gemma-4-12b',
+        model_name='gemma-4-12b',
+        model_family='gemma4',
+        model_version='12b',
+        model_config_id='gemma4_12b',
+        model_config_category='gemma4',
+    ),
+    naming.ModelNaming(
+        model_id='google/gemma-4-12b-it',
+        model_name='gemma-4-12b-it',
+        model_family='gemma4',
+        model_version='12b_it',
+        model_config_id='gemma4_12b_it',
+        model_config_category='gemma4',
+    ),
     naming.ModelNaming(
         model_id='google/gemma-4-31b',
         model_name='gemma-4-31b',