From 490a1b6f8410692e97260124e0d90a3756a7a57f Mon Sep 17 00:00:00 2001 From: Chansung Park Date: Tue, 9 Jun 2026 09:37:00 +0900 Subject: [PATCH] Add Gemma 4 12B model config support --- docs/models.md | 1 + tests/models/automodel_test.py | 14 +++- tests/models/gemma4/model_test.py | 66 ++++++++++++++++ .../models/gemma4/params_safetensors_test.py | 78 +++++++++++++++++++ tunix/models/gemma4/model.py | 33 +++++++- tunix/models/registry.py | 16 ++++ 6 files changed, 203 insertions(+), 5 deletions(-) create mode 100644 tests/models/gemma4/params_safetensors_test.py diff --git a/docs/models.md b/docs/models.md index bf6bba259..ff3a3a8ff 100644 --- a/docs/models.md +++ b/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 | diff --git a/tests/models/automodel_test.py b/tests/models/automodel_test.py index fecff2ab5..c7feb57dc 100644 --- a/tests/models/automodel_test.py +++ b/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 diff --git a/tests/models/gemma4/model_test.py b/tests/models/gemma4/model_test.py index fdf8b457e..84219aa2a 100644 --- a/tests/models/gemma4/model_test.py +++ b/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 diff --git a/tests/models/gemma4/params_safetensors_test.py b/tests/models/gemma4/params_safetensors_test.py new file mode 100644 index 000000000..3ee87723a --- /dev/null +++ b/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() diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py index ab8cd940b..907fffa27 100644 --- a/tunix/models/gemma4/model.py +++ b/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 ), diff --git a/tunix/models/registry.py b/tunix/models/registry.py index 989adeb13..68251dcc3 100644 --- a/tunix/models/registry.py +++ b/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',