Add Gemma4 12b unified model support in Fairseq2

recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml

@@ -0,0 +1,97 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Gemma4 12B (Unified family, dense) GSM8K SFT Fine-tuning Config
+#
+# 12B: 12B params total, 48 layers, model_dim=3840, head_dim=256,
+# 16 attn heads, 8 KV heads, num_global_key_value_heads=1 (MQA global),
+# attention_k_eq_v=True, ffn_inner=15360, sliding_window=1024,
+# max_seq_len=262144, no PLE, no MoE.
+#
+# Recommended hardware: 8x H100/H200 (FSDP). Reduce max_seq_len/max_num_tokens
+# on smaller setups.
+#
+# Usage:
+#   torchrun --standalone --nproc_per_node=8 -m recipes.lm.sft \
+#     --config-file recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml \
+#     /path/to/output_dir
+
+model:
+  family: "gemma4"
+  # Fine-tune the -it variant. SFT on the instruction-tuned checkpoint is the
+  # standard pattern for domain adaptation. The base gemma4_12b card is also
+  # registered (points at /checkpoint/fairseq2/shared/models/gemma-4-12B) if
+  # the user wants to SFT from base — just swap the name below.
+  name: "gemma4_12b_it"
+  dtype: bfloat16
+
+tokenizer:
+  family: "gemma4"
+  name: "gemma4_12b_it"
+
+dataset:
+  max_seq_len: 4096
+  max_num_tokens: 8192
+  valid_split: "sft_test"
+  # NOTE: chat_mode=true requires the chat_template.jinja to use
+  # {% generation %} markers so apply_chat_template can return a non-empty
+  # assistant_masks. Google's Gemma 4 chat_template.jinja does NOT use
+  # {% generation %} (verified 2026-06-07), so target_mask is all-false and
+  # loss collapses to 0 with zero gradient. Until the chat template is
+  # patched upstream (or replaced with a fairseq2-side variant), use
+  # chat_mode=false: the SFT runs as LM continuation on the src+tgt pair,
+  # which still validates the recipe wiring end-to-end and produces a
+  # nonzero loss signal.
+  chat_mode: false
+  config_overrides:
+    sources:
+      train:
+      - path: "hg://facebook/fairseq2-lm-gsm8k"
+        split: "sft_train"
+        weight: 1.0
+      sft_test:
+      - path: "hg://facebook/fairseq2-lm-gsm8k"
+        split: "sft_test"
+        weight: 1.0
+
+trainer:
+  data_parallelism: fsdp
+  max_grad_norm: 1.0
+  mixed_precision:
+    mode: static
+    dtype: bfloat16
+
+optimizer:
+  name: adamw
+  config:
+    lr: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.1
+    impl: fused
+
+lr_scheduler:
+  name: cosine_annealing
+  config:
+    final_lr_scale: 0.1
+    num_warmup_steps: 100
+
+regime:
+  num_steps: 100000
+  checkpoint_every_n_steps: 100
+  validate_every_n_steps: 100
+  keep_last_n_checkpoints: 10
+  publish_metrics_every_n_steps: 1
+  save_model_only: false
+
+common:
+  seed: 0
+  metric_recorders:
+    wandb:
+      enabled: true
+      project: "gemma-4-fairseq2"
+      run_name: "sft_gemma4_12b_gsm8k"
+    tensorboard:
+      enabled: true

src/fairseq2/assets/cards/models/gemma4.yaml

@@ -73,3 +73,21 @@ model_arch: e2b_it
 checkpoint: "hg://google/gemma-4-E2B-it"
 tokenizer: "hg://google/gemma-4-E2B-it"
 tokenizer_family: gemma4
+
+---
+
+name: gemma4_12b
+model_family: gemma4
+model_arch: 12b
+checkpoint: "hg://google/gemma-4-12B"
+tokenizer: "hg://google/gemma-4-12B"
+tokenizer_family: gemma4
+
+---
+
+name: gemma4_12b_it
+model_family: gemma4
+model_arch: 12b_it
+checkpoint: "hg://google/gemma-4-12B-it"
+tokenizer: "hg://google/gemma-4-12B-it"
+tokenizer_family: gemma4

src/fairseq2/models/gemma4/__init__.py

@@ -24,6 +24,7 @@
 )
 from fairseq2.models.gemma4.config import GEMMA4_FAMILY as GEMMA4_FAMILY
 from fairseq2.models.gemma4.config import Gemma4Config as Gemma4Config
+from fairseq2.models.gemma4.config import get_gemma4_12b_config as get_gemma4_12b_config
 from fairseq2.models.gemma4.config import (
     get_gemma4_26b_a4b_config as get_gemma4_26b_a4b_config,
 )
@@ -83,6 +84,7 @@
     "apply_fsdp_to_gemma4",
     "convert_gemma4_state_dict",
     "create_gemma4_model",
+    "get_gemma4_12b_config",
     "get_gemma4_26b_a4b_config",
     "get_gemma4_31b_config",
     "get_gemma4_e2b_config",

src/fairseq2/models/gemma4/audio/config.py

@@ -7,23 +7,43 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Literal
 
 
 @dataclass(kw_only=True)
 class Gemma4AudioConfig:
-    """Configuration for the Gemma 4 audio tower (USM Conformer).
+    """Configuration for the Gemma 4 audio pipeline.
 
-    Default values correspond to the E4B model.
+    Default values correspond to the E4B model (Conformer-based mel pipeline).
+
+    The ``audio_mode`` field selects between the two Gemma 4 audio pipelines:
+
+    * ``"conformer"`` (default; used by E4B / classic Gemma 4): mel-spectrogram
+      input goes through a subsampling Conv2d stack and 12 Conformer layers
+      (the audio tower) before the multimodal embedder projects to text space.
+
+    * ``"linear"`` (used by Gemma 4 Unified family — 12B+): raw waveform frames
+      of ``audio_samples_per_token`` (640) samples each are fed directly through
+      the multimodal embedder (RMSNorm + Linear) — no tower, no mel, no convs.
+      The ``hidden_size``, ``num_hidden_layers``, conv/attention etc. fields
+      are ignored in this mode; only ``output_proj_dims`` (= 640 for unified)
+      and ``rms_norm_eps`` are read.
     """
 
+    audio_mode: Literal["conformer", "linear"] = "conformer"
+    """Audio pipeline selector: ``"conformer"`` or ``"linear"``."""
+
     hidden_size: int = 1024
-    """Audio encoder hidden dimension."""
+    """Audio encoder hidden dimension. (conformer mode only)"""
 
     output_proj_dims: int = 1536
-    """Output projection dimension (before text embedder)."""
+    """Output projection dimension (before text embedder).
+
+    For the unified ``linear`` mode this equals the raw-waveform frame size
+    (typically 640 samples = 40 ms @ 16 kHz)."""
 
     num_hidden_layers: int = 12
-    """Number of conformer layers."""
+    """Number of conformer layers. (conformer mode only)"""
 
     num_attention_heads: int = 8
     """Number of attention heads. head_dim = hidden_size / num_attention_heads."""

src/fairseq2/models/gemma4/audio/embedder.py

@@ -27,25 +27,46 @@ class Gemma4MultimodalAudioEmbedder(Module):
     followed by a ``Linear`` projection from ``output_proj_dims`` to
     ``text_model_dim``.
 
+    This single class implements **two** matching HF classes:
+
+    * ``transformers.Gemma4MultimodalEmbedder`` (classic gemma4 family,
+      E4B/31B/26B-A4B): forward is ``RMSNorm -> Linear``, callers are
+      expected to feed inputs already in the embedder's dtype.
+
+    * ``transformers.Gemma4UnifiedMultimodalEmbedder`` (gemma4_unified
+      family, 12B+): identical math, but the forward additionally casts
+      ``inputs_embeds`` to ``self.embedding_projection.weight.dtype`` before
+      the norm. This matters when raw waveform features (typically fp32 from
+      the feature extractor) are fed into a bf16 embedder.
+
+    The ``cast_input_dtype`` ctor flag selects between the two: ``False``
+    (default) preserves the classic gemma4 behaviour bit-for-bit; ``True``
+    activates the Unified family's input cast. The factory sets it to
+    ``True`` when ``audio_config.audio_mode == "linear"``.
+
     Note: HF does NOT use ClippableLinear for the embedder projection --
     the checkpoint key is ``model.embed_audio.embedding_projection.weight``
     (plain ``nn.Linear``, no clipping buffers).
     """
 
     embedding_pre_projection_norm: Gemma4AudioRMSNorm
     embedding_projection: Linear
+    cast_input_dtype: bool
 
     def __init__(
         self,
         output_proj_dims: int,
         text_model_dim: int,
         rms_norm_eps: float = 1e-6,
         *,
+        cast_input_dtype: bool = False,
         device: Device | None = None,
         dtype: DataType | None = None,
     ) -> None:
         super().__init__()
 
+        self.cast_input_dtype = cast_input_dtype
+
         # RMSNorm without learnable scale (elementwise_affine=False)
         self.embedding_pre_projection_norm = Gemma4AudioRMSNorm(
             output_proj_dims,
@@ -70,6 +91,11 @@ def forward(self, features: Tensor) -> Tensor:
         :param features: Audio tower output. *Shape:* :math:`(N,T,D)`.
         :returns: Text-space embeddings. *Shape:* :math:`(N,T,H_{text})`.
         """
+        if self.cast_input_dtype:
+            # Match HF Gemma4UnifiedMultimodalEmbedder: cast raw inputs
+            # (often fp32 from the feature extractor) to the embedder weight
+            # dtype (typically bf16) before the norm.
+            features = features.to(self.embedding_projection.weight.dtype)
         features = self.embedding_pre_projection_norm(features)
         features = self.embedding_projection(features)
         return features

src/fairseq2/models/gemma4/config.py

@@ -240,6 +240,22 @@ def _e2b() -> Gemma4Config:
     def _e2b_it() -> Gemma4Config:
         return get_gemma4_e2b_config()
 
+    @arch("12b")
+    def _12b() -> Gemma4Config:
+        return get_gemma4_12b_config()
+
+    @arch("12b_it")
+    def _12b_it() -> Gemma4Config:
+        return get_gemma4_12b_config()
+
+    @arch("12b_audio")
+    def _12b_audio() -> Gemma4Config:
+        return get_gemma4_12b_audio_config()
+
+    @arch("12b_it_audio")
+    def _12b_it_audio() -> Gemma4Config:
+        return get_gemma4_12b_audio_config()
+
 
 def get_gemma4_e2b_config() -> Gemma4Config:
     """Get configuration for Gemma4 E2B (small dense, on-device).
@@ -335,3 +351,85 @@ def get_gemma4_26b_a4b_config() -> Gemma4Config:
         moe_intermediate_size=704,
         final_logit_soft_cap=30.0,
     )
+
+
+def get_gemma4_unified_audio_config() -> Gemma4AudioConfig:
+    """Audio config for the Gemma 4 Unified family (12B+).
+
+    Linear (tower-free) pipeline: raw 16 kHz waveform is chunked into frames
+    of ``audio_samples_per_token`` = 640 samples (40 ms each), then projected
+    to the text model dim through RMSNorm + Linear. No mel-spectrogram, no
+    Conformer.
+
+    Matches HF's ``Gemma4UnifiedAudioConfig`` (model_type=gemma4_unified_audio)
+    where ``audio_embed_dim = audio_samples_per_token = output_proj_dims = 640``.
+    """
+    return Gemma4AudioConfig(
+        audio_mode="linear",
+        # In linear mode, only output_proj_dims (= 640 raw samples per token)
+        # and rms_norm_eps are read. Other Conformer fields default values
+        # are unused.
+        output_proj_dims=640,
+        rms_norm_eps=1e-6,
+    )
+
+
+def get_gemma4_12b_config() -> Gemma4Config:
+    """Get configuration for Gemma4 12B (Unified family, dense).
+
+    First member of HF ``gemma4_unified`` model_type (released 2026-05-23).
+    The text decoder is a dense Gemma 4 model that reuses the same attention,
+    decoder, and frontend code paths as the existing 31B dense variant.
+    Distinguishing values vs the existing dense archs:
+
+    * ``num_global_key_value_heads = 1`` — multi-query (MQA) global attention.
+      Previously the dense archs only used 2 (26B-A4B) and 4 (31B).
+    * ``attention_k_eq_v = True`` — keys reused as values in global layers.
+    * ``hidden_size_per_layer_input = 0`` — no PLE.
+    * ``num_kv_shared_layers = 0`` — no KV sharing.
+    * 48 layers, 5:1 sliding:full pattern (40 sliding + 8 full).
+
+    The 12B Unified checkpoint also ships an ``embed_audio`` projection
+    (``[3840, 640]``) and a ``vision_embedder`` pipeline (LN + Dense + LN +
+    factorized 2D positional embedding + RMSNorm + Linear), but those are
+    not required for text-only logit parity or downstream text evaluation.
+    Multimodal embedders are intentionally not registered here — when
+    ``audio_config`` is ``None`` (the default), :func:`convert_gemma4_state_dict`
+    filters multimodal keys (audio_tower, embed_audio, vision_tower,
+    embed_vision, vision_embedder, multi_modal_projector).
+    """
+    return Gemma4Config(
+        model_dim=3840,
+        max_seq_len=262_144,
+        num_layers=48,
+        num_attn_heads=16,
+        num_key_value_heads=8,
+        head_dim=256,
+        global_head_dim=512,
+        num_global_key_value_heads=1,
+        ffn_inner_dim=15_360,
+        sliding_window=1024,
+        attention_k_eq_v=True,
+        num_kv_shared_layers=0,
+        hidden_size_per_layer_input=0,  # PLE disabled (Unified family has no PLE)
+        final_logit_soft_cap=30.0,
+    )
+
+
+def get_gemma4_12b_audio_config() -> Gemma4Config:
+    """Get configuration for Gemma4 12B (Unified family) WITH the audio
+    embedder enabled.
+
+    Identical to :func:`get_gemma4_12b_config` except that ``audio_config`` is
+    set to :func:`get_gemma4_unified_audio_config` (linear mode, no tower).
+    Use this arch when you want to consume audio inputs through the
+    fairseq2 inference path (audio+text -> text).
+
+    The text-only ``12b`` / ``12b_it`` archs are unchanged and remain the
+    canonical entry point for logit parity, MMLU, SFT — keeping the audio
+    embedder out of those configs avoids loading unused parameters and
+    preserves the converter's multimodal filter.
+    """
+    cfg = get_gemma4_12b_config()
+    cfg.audio_config = get_gemma4_unified_audio_config()
+    return cfg

src/fairseq2/models/gemma4/factory.py

@@ -527,14 +527,23 @@ def create_final_projection(self, embed: Embedding) -> Projection:
     def create_audio_tower(self) -> Module | None:
         """Create the audio tower for mel-spectrogram encoding.
 
-        :returns: A :class:`Gemma4AudioTower` if audio is configured,
-            ``None`` otherwise.
+        :returns: A :class:`Gemma4AudioTower` if audio is configured AND the
+            audio_mode is ``"conformer"``; ``None`` otherwise.
+
+        The Gemma 4 Unified family (``audio_mode="linear"``) has no audio
+        tower — raw waveform frames are fed directly through the
+        multimodal embedder (see :meth:`create_audio_embedder` and
+        ``Gemma4Model.forward``).
         """
         config = self._config
 
         if config.audio_config is None:
             return None
 
+        # Unified family: no tower; embedder consumes raw waveform frames.
+        if config.audio_config.audio_mode == "linear":
+            return None
+
         from fairseq2.models.gemma4.audio.tower import Gemma4AudioTower
 
         return Gemma4AudioTower(
@@ -558,10 +567,17 @@ def create_audio_embedder(self) -> Module | None:
             Gemma4MultimodalAudioEmbedder,
         )
 
+        # Linear mode (Gemma 4 Unified family) requires the HF Unified
+        # behaviour where the embedder casts raw inputs to its weight dtype
+        # before the norm. The classic conformer path does not need this
+        # (its inputs come from the audio tower already in the right dtype).
+        cast_input_dtype = config.audio_config.audio_mode == "linear"
+
         return Gemma4MultimodalAudioEmbedder(
             output_proj_dims=config.audio_config.output_proj_dims,
             text_model_dim=config.model_dim,
             rms_norm_eps=config.audio_config.rms_norm_eps,
+            cast_input_dtype=cast_input_dtype,
             device=self._device,
             dtype=self._dtype,
         )