Small Gemma draft

Author: gagika

pytest.ini

@@ -14,6 +14,7 @@ addopts =
     --ignore=tests/unit/dequantize_pack_quantized_int4_test.py
     --ignore=tests/unit/gemma3_layers_test.py
     --ignore=tests/unit/gemma4_layers_test.py
+    --ignore=tests/unit/gemma4_small_layers_test.py
     --ignore=tests/unit/gpt_vs_reference_test.py
     --ignore=tests/unit/llama4_layers_test.py
     --ignore=tests/unit/hf_checkpoint_conversion_test.py

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -144,14 +144,137 @@
 )
 
 
+gemma4_e2b_dict = {
+    "architectures": ["Gemma4ForConditionalGeneration"],
+    "audio_config": None,
+    "audio_token_id": 258881,
+    "boa_token_id": 256000,
+    "boi_token_id": 255999,
+    "dtype": "bfloat16",
+    "eoa_token_id": 258883,
+    "eoa_token_index": 258883,
+    "eoi_token_id": 258882,
+    "eos_token_id": [1, 106],
+    "image_token_id": 258880,
+    "initializer_range": 0.02,
+    "model_type": "gemma4",
+    "text_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "attention_k_eq_v": False,
+        "bos_token_id": 2,
+        "dtype": "bfloat16",
+        "enable_moe_block": False,
+        "eos_token_id": 1,
+        "expert_intermediate_size": None,
+        "final_logit_softcapping": 30.0,
+        "global_head_dim": 512,
+        "head_dim": 256,
+        "hidden_activation": "gelu_pytorch_tanh",
+        "hidden_size": 1536,
+        "hidden_size_per_layer_input": 256,
+        "initializer_range": 0.02,
+        "intermediate_size": 6144,
+        "layer_types": [
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+        ]
+        * 7,
+        "max_position_embeddings": 131072,
+        "model_type": "gemma4_text",
+        "num_attention_heads": 8,
+        "num_experts": None,
+        "num_global_key_value_heads": None,
+        "num_hidden_layers": 35,
+        "num_key_value_heads": 1,
+        "num_kv_shared_layers": 20,
+        "pad_token_id": 0,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "full_attention": {
+                "partial_rotary_factor": 0.25,
+                "rope_theta": 1_000_000.0,
+                "rope_type": "proportional",
+            },
+            "sliding_attention": {"rope_theta": 10_000.0, "rope_type": "default"},
+        },
+        "sliding_window": 512,
+        "tie_word_embeddings": True,
+        "top_k_experts": None,
+        "use_bidirectional_attention": None,
+        "use_cache": True,
+        "use_double_wide_mlp": True,
+        "vocab_size": 262144,
+        "vocab_size_per_layer_input": 262144,
+    },
+    "tie_word_embeddings": True,
+    "transformers_version": "5.5.0.dev0",
+    "video_token_id": 258884,
+    "vision_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "default_output_length": 280,
+        "dtype": "bfloat16",
+        "global_head_dim": 64,
+        "head_dim": 64,
+        "hidden_activation": "gelu_pytorch_tanh",
+        "hidden_size": 768,
+        "intermediate_size": 3072,
+        "max_position_embeddings": 131072,
+        "model_type": "gemma4_vision",
+        "num_attention_heads": 12,
+        "num_hidden_layers": 16,
+        "num_key_value_heads": 12,
+        "patch_size": 16,
+        "pooling_kernel_size": 3,
+        "position_embedding_size": 10240,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {"rope_theta": 100.0, "rope_type": "default"},
+        "standardize": False,
+        "use_clipped_linears": True,
+    },
+    "vision_soft_tokens_per_image": 280,
+}
+
+
+gemma4_e4b_dict = gemma4_e2b_dict.copy()
+gemma4_e4b_dict["text_config"] = gemma4_e2b_dict["text_config"].copy()
+gemma4_e4b_dict["text_config"].update(
+    {
+        "hidden_size": 2560,
+        "intermediate_size": 10240,
+        "layer_types": [
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+        ]
+        * 7,
+        "num_hidden_layers": 42,
+        "num_key_value_heads": 2,
+        "num_kv_shared_layers": 18,
+        "use_double_wide_mlp": False,
+    }
+)
+
+
 try:
   # Will execute successfully if Transformers is updated with Gemma 4 support
   gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict)
   gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict)
+  gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict)
+  gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict)
 except AttributeError:
   # Graceful fallback to raw dict-based PTConfig if Gemma 4 natively is missing
   gemma4_26b_config = PTConfig(**gemma4_26b_dict)  # pytype: disable=wrong-arg-types
   gemma4_31b_config = PTConfig(**gemma4_31b_dict)  # pytype: disable=wrong-arg-types
+  gemma4_e2b_config = PTConfig(**gemma4_e2b_dict)  # pytype: disable=wrong-arg-types
+  gemma4_e4b_config = PTConfig(**gemma4_e4b_dict)  # pytype: disable=wrong-arg-types
 
 
 gemma3_4b_config = transformers.Gemma3Config(
@@ -1185,6 +1308,8 @@ def __init__(self, **kwargs):
     "gemma3-27b": gemma3_27b_config,
     "gemma4-26b": gemma4_26b_config,
     "gemma4-31b": gemma4_31b_config,
+    "gemma4-e2b": gemma4_e2b_config,
+    "gemma4-e4b": gemma4_e4b_config,
     "qwen2.5-1.5b": qwen25_1_5b_config,
     "qwen2.5-7b": qwen25_7b_config,
     "qwen2.5-14b": qwen25_14b_config,

src/maxtext/checkpoint_conversion/utils/hf_shape.py

@@ -284,6 +284,92 @@ def GEMMA4_HF_WEIGHTS_TO_SHAPE(config):
   return shapes
 
 
+def GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE(config):
+  """Generates HF parameter shapes for Gemma 4 small (E2B / E4B).
+
+  Differs from GEMMA4_HF_WEIGHTS_TO_SHAPE in that it:
+    * derives global-vs-sliding from the per-model ``layer_types`` list
+      (E2B has period-5, E4B has period-6),
+    * emits the Per-Layer-Embedding parameters when ``hidden_size_per_layer_input`` > 0,
+    * omits k_proj/v_proj/k_norm/v_norm shapes on KV-shared layers, and
+    * doubles ``intermediate_size`` on shared layers when ``use_double_wide_mlp``
+      is set (E2B).
+  """
+  shapes = {}
+
+  text_cfg = config.get("text_config", config)
+  vision_cfg = config.get("vision_config", {})
+  text_base = "model.language_model" if vision_cfg else "model"
+
+  hidden_size = text_cfg["hidden_size"]
+  intermediate_size = text_cfg["intermediate_size"]
+  num_hidden_layers = text_cfg["num_hidden_layers"]
+  num_attention_heads = text_cfg["num_attention_heads"]
+  num_key_value_heads = text_cfg["num_key_value_heads"]
+  num_global_key_value_heads = text_cfg.get("num_global_key_value_heads") or num_key_value_heads
+  head_dim = text_cfg["head_dim"]
+  global_head_dim = text_cfg.get("global_head_dim", head_dim)
+  vocab_size = text_cfg["vocab_size"]
+  layer_types = text_cfg.get("layer_types", [])
+
+  ple_dim = text_cfg.get("hidden_size_per_layer_input", 0) or 0
+  vocab_ple = text_cfg.get("vocab_size_per_layer_input", 0) or 0
+  num_kv_shared = text_cfg.get("num_kv_shared_layers", 0) or 0
+  first_shared = max(0, num_hidden_layers - num_kv_shared) if num_kv_shared > 0 else num_hidden_layers
+  use_double_wide_mlp = bool(text_cfg.get("use_double_wide_mlp", False))
+
+  shapes[f"{text_base}.embed_tokens.weight"] = [vocab_size, hidden_size]
+  shapes[f"{text_base}.norm.weight"] = [hidden_size]
+
+  if ple_dim > 0:
+    shapes[f"{text_base}.embed_tokens_per_layer.weight"] = [vocab_ple, num_hidden_layers * ple_dim]
+    shapes[f"{text_base}.per_layer_model_projection.weight"] = [num_hidden_layers * ple_dim, hidden_size]
+    shapes[f"{text_base}.per_layer_projection_norm.weight"] = [ple_dim]
+
+  for i in range(num_hidden_layers):
+    hf_prefix = f"{text_base}.layers.{i}"
+    is_global = i < len(layer_types) and layer_types[i] == "full_attention"
+    is_shared = num_kv_shared > 0 and i >= first_shared
+
+    if is_global:
+      q_dim = num_attention_heads * global_head_dim
+      kv_dim = num_global_key_value_heads * global_head_dim
+      norm_dim = global_head_dim
+    else:
+      q_dim = num_attention_heads * head_dim
+      kv_dim = num_key_value_heads * head_dim
+      norm_dim = head_dim
+
+    shapes[f"{hf_prefix}.self_attn.q_proj.weight"] = [q_dim, hidden_size]
+    shapes[f"{hf_prefix}.self_attn.o_proj.weight"] = [hidden_size, q_dim]
+    shapes[f"{hf_prefix}.self_attn.q_norm.weight"] = [norm_dim]
+    if not is_shared:
+      shapes[f"{hf_prefix}.self_attn.k_proj.weight"] = [kv_dim, hidden_size]
+      shapes[f"{hf_prefix}.self_attn.v_proj.weight"] = [kv_dim, hidden_size]
+      shapes[f"{hf_prefix}.self_attn.k_norm.weight"] = [norm_dim]
+      # v_norm only when scale is enabled in MaxText; param_mapping suppresses
+      # this key otherwise, so emit the shape unconditionally — extras are ignored.
+      shapes[f"{hf_prefix}.self_attn.v_norm.weight"] = [norm_dim]
+
+    shapes[f"{hf_prefix}.input_layernorm.weight"] = [hidden_size]
+    shapes[f"{hf_prefix}.post_attention_layernorm.weight"] = [hidden_size]
+    shapes[f"{hf_prefix}.pre_feedforward_layernorm.weight"] = [hidden_size]
+    shapes[f"{hf_prefix}.post_feedforward_layernorm.weight"] = [hidden_size]
+    shapes[f"{hf_prefix}.layer_scalar"] = [1]
+
+    mlp_dim = intermediate_size * 2 if (is_shared and use_double_wide_mlp) else intermediate_size
+    shapes[f"{hf_prefix}.mlp.gate_proj.weight"] = [mlp_dim, hidden_size]
+    shapes[f"{hf_prefix}.mlp.up_proj.weight"] = [mlp_dim, hidden_size]
+    shapes[f"{hf_prefix}.mlp.down_proj.weight"] = [hidden_size, mlp_dim]
+
+    if ple_dim > 0:
+      shapes[f"{hf_prefix}.per_layer_input_gate.weight"] = [ple_dim, hidden_size]
+      shapes[f"{hf_prefix}.per_layer_projection.weight"] = [hidden_size, ple_dim]
+      shapes[f"{hf_prefix}.post_per_layer_input_norm.weight"] = [hidden_size]
+
+  return shapes
+
+
 def GEMMA2_HF_WEIGHTS_TO_SHAPE(config):
   """Returns mapping between HuggingFace weights path and weights shape.
 
@@ -920,6 +1006,8 @@ def MIXTRAL_HF_WEIGHTS_TO_SHAPE(config):
     "gemma3-27b": GEMMA3_HF_WEIGHTS_TO_SHAPE,
     "gemma4-26b": GEMMA4_HF_WEIGHTS_TO_SHAPE,
     "gemma4-31b": GEMMA4_HF_WEIGHTS_TO_SHAPE,
+    "gemma4-e2b": GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE,
+    "gemma4-e4b": GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE,
     "qwen2.5-1.5b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen2.5-7b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen2.5-14b": QWEN_HF_WEIGHTS_TO_SHAPE,