diff --git a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml new file mode 100644 index 000000000..91efb6b6d --- /dev/null +++ b/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 diff --git a/src/fairseq2/assets/cards/models/gemma4.yaml b/src/fairseq2/assets/cards/models/gemma4.yaml index e1b38693f..0052771a9 100644 --- a/src/fairseq2/assets/cards/models/gemma4.yaml +++ b/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 diff --git a/src/fairseq2/models/gemma4/__init__.py b/src/fairseq2/models/gemma4/__init__.py index d34fd3c22..81d94becf 100644 --- a/src/fairseq2/models/gemma4/__init__.py +++ b/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", diff --git a/src/fairseq2/models/gemma4/audio/config.py b/src/fairseq2/models/gemma4/audio/config.py index dbf68b415..9fa4289ce 100644 --- a/src/fairseq2/models/gemma4/audio/config.py +++ b/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.""" diff --git a/src/fairseq2/models/gemma4/audio/embedder.py b/src/fairseq2/models/gemma4/audio/embedder.py index 0238f67d8..f9acf06f2 100644 --- a/src/fairseq2/models/gemma4/audio/embedder.py +++ b/src/fairseq2/models/gemma4/audio/embedder.py @@ -27,6 +27,23 @@ 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). @@ -34,6 +51,7 @@ class Gemma4MultimodalAudioEmbedder(Module): embedding_pre_projection_norm: Gemma4AudioRMSNorm embedding_projection: Linear + cast_input_dtype: bool def __init__( self, @@ -41,11 +59,14 @@ def __init__( 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 diff --git a/src/fairseq2/models/gemma4/config.py b/src/fairseq2/models/gemma4/config.py index e2629187c..8b7b28f9e 100644 --- a/src/fairseq2/models/gemma4/config.py +++ b/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 diff --git a/src/fairseq2/models/gemma4/factory.py b/src/fairseq2/models/gemma4/factory.py index d095d3073..10dbe22e2 100644 --- a/src/fairseq2/models/gemma4/factory.py +++ b/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, ) diff --git a/src/fairseq2/models/gemma4/interop.py b/src/fairseq2/models/gemma4/interop.py index 8c148d433..0261304fa 100644 --- a/src/fairseq2/models/gemma4/interop.py +++ b/src/fairseq2/models/gemma4/interop.py @@ -204,9 +204,15 @@ def convert_gemma4_state_dict( """ # Determine which multimodal prefixes to filter out. # Always filter vision; only filter audio when not configured. + # + # Vision prefixes cover both the existing ``gemma4`` family (vision_tower + # + multi_modal_projector) and the newer ``gemma4_unified`` family which + # ships a ``vision_embedder`` (LN + Dense + LN + factorized 2D positional + # embedding + RMSNorm + Linear) in place of a SigLIP-style tower. multimodal_prefixes = [ "model.vision_tower.", "model.embed_vision.", + "model.vision_embedder.", "model.multi_modal_projector.", ] if config.audio_config is None: diff --git a/src/fairseq2/models/gemma4/model.py b/src/fairseq2/models/gemma4/model.py index 1abf242ec..e4eb32546 100644 --- a/src/fairseq2/models/gemma4/model.py +++ b/src/fairseq2/models/gemma4/model.py @@ -169,11 +169,22 @@ def forward( :returns: Logits or loss (or both if return_logits=True). """ # Encode audio through tower + embedder before frontend. + # Two supported audio pipelines: + # (a) Conformer family (E*, classic Gemma 4): mel-spec -> + # audio_tower -> audio_embedder -> text-space embeds. + # (b) Unified family (gemma4_unified / 12B+): raw waveform frames + # of audio_samples_per_token=640 samples each, fed directly + # through the embedder (RMSNorm + Linear). No tower. + # Selection: presence of self.audio_tower. audio_embeds: Tensor | None = None - if audio_features is not None: - if self.audio_tower is not None and self.audio_embedder is not None: + if audio_features is not None and self.audio_embedder is not None: + if self.audio_tower is not None: tower_output = self.audio_tower(audio_features) audio_embeds = self.audio_embedder(tower_output) + else: + # Unified family: embedder consumes raw waveform frames + # (shape (B, T, audio_samples_per_token)) directly. + audio_embeds = self.audio_embedder(audio_features) seqs, seqs_layout, per_layer_embeds = self.decoder_frontend( seqs, diff --git a/src/fairseq2/models/gemma4/tokenizer.py b/src/fairseq2/models/gemma4/tokenizer.py index e974c6f32..68490a668 100644 --- a/src/fairseq2/models/gemma4/tokenizer.py +++ b/src/fairseq2/models/gemma4/tokenizer.py @@ -57,15 +57,28 @@ def create_encoder( if lang is not None: raise ValueError(f"`lang` must be `None`, but is '{lang}' instead.") - if mode is not None and mode not in ("default", "prompt", "as_is"): + if mode is not None and mode not in ( + "default", + "prompt", + "prompt_response", + "as_is", + ): raise ValueError( - f"`mode` must be 'default', 'prompt', or 'as_is', but is '{mode}' instead." + "`mode` must be 'default', 'prompt', 'prompt_response', or " + f"'as_is', but is '{mode}' instead." ) - # Gemma 4 uses BOS token (ID 2) as prefix, no EOS suffix. + # Gemma 4 uses BOS token (ID 2) as prefix, no EOS suffix by default. + # The SFT recipe (chat_mode=false) calls create_encoder twice per + # example: once with mode='prompt' for the source (BOS-prefixed, no + # EOS) and once with mode='prompt_response' for the target (no BOS, + # EOS-suffixed). The two encodings are concatenated. if mode == "as_is": prefix_tokens: list[str] = [] suffix_tokens: list[str] = [] + elif mode == "prompt_response": + prefix_tokens = [] + suffix_tokens = [""] else: prefix_tokens = [""] suffix_tokens = []