Support for gemma4 12b dense model (#2060)

Author: jordimas

docs/guides/transformers.md

@@ -227,19 +227,19 @@ print(tok.convert_tokens_to_string(res[0].sequences[0]))
 [Gemma 4](https://ai.google.dev/gemma/docs/gemma4) is Google's next generation of lightweight open-weight models, featuring a hybrid attention architecture with interleaved global and sliding-window attention layers.
 
 ```{note}
-Only the 31B dense model (`gemma-4-31B`) is currently supported. The MoE variants (E2B, E4B) are not supported.
+Only the 31B and 12B dense models are currently supported. The MoE variants (E2B, E4B) are not supported.
 ```
 
-Gemma 4 models come in two flavors: instruction tuned (it) models and base models.
+Gemma 4 models come in two flavors: instruction tuned (it) models and pre-trained models.
 
 Instruction tuned models use the same [prompt template format](https://ai.google.dev/gemma/docs/core/prompt-structure) as Gemma 3.
 
 When converting an instruction-tuned model, CTranslate2 sets `<end_of_turn>` as the default end-of-sequence token.
 
-To convert the 31B instruction-tuned model:
+To convert the 12B instruction-tuned model:
 
 ```bash
-ct2-transformers-converter --model google/gemma-4-31B-it --quantization float16 --output_dir gemma-4-31b-it
+ct2-transformers-converter --model google/gemma-4-12B-it --quantization float16 --output_dir gemma-4-12b-it
 ```
 
 Usage sample:
@@ -248,8 +248,8 @@ Usage sample:
 from transformers import AutoTokenizer
 import ctranslate2
 
-tok = AutoTokenizer.from_pretrained("google/gemma-4-31B-it")
-gen = ctranslate2.Generator("gemma-4-31b-it", device="auto")
+tok = AutoTokenizer.from_pretrained("google/gemma-4-12B-it")
+gen = ctranslate2.Generator("gemma-4-12b-it", device="auto")
 
 prompt = "<start_of_turn>user\nGenerate a 200 word text talking about George Orwell.<end_of_turn>\n<start_of_turn>model\n"
 tokens = tok.convert_ids_to_tokens(tok.encode(prompt))

python/ctranslate2/converters/transformers.py

@@ -1886,9 +1886,11 @@ def get_model_spec(self, model):
             )
             if sliding_window_pattern is not None:
                 layer_types = [
-                    "full_attention"
-                    if (i + 1) % sliding_window_pattern == 0
-                    else "sliding_attention"
+                    (
+                        "full_attention"
+                        if (i + 1) % sliding_window_pattern == 0
+                        else "sliding_attention"
+                    )
                     for i in range(num_layers)
                 ]
 
@@ -2058,6 +2060,8 @@ def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
             gc.collect()
 
 
+@register_loader("Gemma4UnifiedTextConfig")
+@register_loader("Gemma4UnifiedConfig")
 @register_loader("Gemma4TextConfig")
 @register_loader("Gemma4Config")
 class Gemma4Loader(ModelLoader):
@@ -2068,6 +2072,8 @@ def architecture_name(self):
     def get_model_class(self, config, default_class):
         if config.__class__.__name__ == "Gemma4Config":
             return transformers.Gemma4ForConditionalGeneration
+        if config.__class__.__name__ == "Gemma4UnifiedConfig":
+            return transformers.Gemma4UnifiedForConditionalGeneration
         return default_class
 
     def get_model_spec(self, model):
@@ -2109,18 +2115,28 @@ def get_model_spec(self, model):
         rope_local_base_freq = float(sliding_rope.get("rope_theta", 10_000))
         rope_theta = float(global_rope.get("rope_theta", 1_000_000))
 
-        # Proportional RoPE: only a fraction of global_head_dim uses RoPE
+        # Proportional RoPE (HF `rope_type="proportional"`, currently Gemma4-only):
+        # halves on full head_dim and zero-pads trailing freqs, unlike GPT-NeoX-style
+        # partial RoPE which halves on rotary_dim. HF: `1 / rope_theta^(2i/head_dim)`;
+        # CT2's RotaryEmbeddings: `1 / base^(2i/rotary_dim)`. Rescale base to match.
         global_partial_factor = float(global_rope.get("partial_rotary_factor", 1.0))
         global_rotary_dim = int(global_head_dim * global_partial_factor)
+        global_rope_base = (
+            rope_theta ** (global_rotary_dim / global_head_dim)
+            if 0 < global_rotary_dim < global_head_dim
+            else rope_theta
+        )
 
         sliding_window = getattr(text_config, "sliding_window", 512)
         layer_types = getattr(text_config, "layer_types", None)
         if layer_types is None:
             sliding_window_pattern = 6
             layer_types = [
-                "sliding_attention"
-                if bool((i + 1) % sliding_window_pattern)
-                else "full_attention"
+                (
+                    "sliding_attention"
+                    if bool((i + 1) % sliding_window_pattern)
+                    else "full_attention"
+                )
                 for i in range(num_layers)
             ]
 
@@ -2140,7 +2156,7 @@ def get_model_spec(self, model):
             quant_group_size = None
             quant_bits = None
 
-        # Build spec with sliding-attention defaults; global layers overridden per-layer below
+        # Build spec with sliding-attention defaults; global layers overridden per-layer below.
         spec = transformer_spec.TransformerDecoderModelSpec.from_config(
             num_layers,
             num_heads,
@@ -2166,8 +2182,14 @@ def get_model_spec(self, model):
             v_norm=True,
         )
 
+        # Set it to 0 so the decoder processes all tokens at once; per-layer sliding_window
+        # set below handles KV-cache trimming for sliding-attention layers.
+        spec.decoder.sliding_window = np.dtype("int32").type(0)
+
         self._layer_types = layer_types
         self._attention_k_eq_v = attention_k_eq_v
+        self._global_head_dim = global_head_dim
+        self._global_rotary_dim = global_rotary_dim
 
         # Per-layer overrides for full-attention layers
         for i, layer_type in enumerate(layer_types):
@@ -2178,7 +2200,9 @@ def get_model_spec(self, model):
                 layer.self_attention.rotary_dim = np.dtype("int32").type(
                     global_rotary_dim
                 )
-                layer.self_attention.rotary_base = np.dtype("float32").type(rope_theta)
+                layer.self_attention.rotary_base = np.dtype("float32").type(
+                    global_rope_base
+                )
                 layer.self_attention.sliding_window = np.dtype("int32").type(0)
                 layer.self_attention.head_dim = np.dtype("int32").type(global_head_dim)
                 if num_global_kv_heads is not None:
@@ -2226,7 +2250,10 @@ def set_config(self, config, model, tokenizer):
             and isinstance(tokenizer.chat_template, str)
             and tokenizer.chat_template.strip()
         ):
-            config.eos_token = "<end_of_turn>"
+            if "<turn|>" in tokenizer.chat_template:
+                config.eos_token = "<turn|>"
+            else:
+                config.eos_token = "<end_of_turn>"
         else:
             config.eos_token = tokenizer.eos_token
 
@@ -2241,6 +2268,18 @@ def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
         self.set_layer_norm(spec.layer_norm, module.norm)
 
         attention_k_eq_v = getattr(self, "_attention_k_eq_v", False)
+        ghd = getattr(self, "_global_head_dim", None)
+        grd = getattr(self, "_global_rotary_dim", None)
+        # HF's proportional partial-RoPE pairs channels [0:R/2]↔[HD/2:HD/2+R/2];
+        # CT2's RotaryEmbeddings pairs [0:R/2]↔[R/2:R]. Permute Q/K accordingly.
+        partial_perm = None
+        if ghd and grd and 0 < grd < ghd:
+            partial_perm = (
+                list(range(0, grd // 2))
+                + list(range(ghd // 2, ghd // 2 + grd // 2))
+                + list(range(grd // 2, ghd // 2))
+                + list(range(ghd // 2 + grd // 2, ghd))
+            )
 
         for layer_spec, layer in zip(spec.layer, module.layers):
             self.set_layer_norm(layer_spec.input_layer_norm, layer.input_layernorm)
@@ -2294,6 +2333,22 @@ def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2):
                     layer_spec.self_attention.linear[0], split_layers, cc_dim
                 )
 
+            # Apply the partial-RoPE permutation to Q/K rows of the fused QKV (and
+            # the norm gammas); V rows are left untouched (V is not RoPE-rotated).
+            if is_full_attn and partial_perm is not None:
+                fused = layer_spec.self_attention.linear[0].weight
+                q_rows = split_layers[0].weight.shape[0]
+                k_rows = split_layers[1].weight.shape[0]
+                qk = fused[: q_rows + k_rows].view(-1, ghd, fused.shape[1])
+                fused[: q_rows + k_rows] = qk[:, partial_perm, :].reshape(
+                    q_rows + k_rows, fused.shape[1]
+                )
+                for norm_spec in (
+                    layer_spec.self_attention.q_norm,
+                    layer_spec.self_attention.k_norm,
+                ):
+                    norm_spec.gamma = norm_spec.gamma[partial_perm]
+
             self.set_linear(
                 layer_spec.self_attention.linear[1],
                 layer.self_attn.o_proj,