From e2db777e4ba9d4973b22daf11f7d6c2c6d1a7f2b Mon Sep 17 00:00:00 2001 From: ProfAI Date: Fri, 5 Jun 2026 21:18:25 +0000 Subject: [PATCH 01/12] [gemma4] Add 12B (Unified family) arch and extend HF multimodal filter * config.py: register arch('12b') and arch('12b_it'); add get_gemma4_12b_config(). Dense Gemma 4 text decoder reusing the gemma4 code paths with: hidden=3840, layers=48 (5:1 sliding:full = 40+8), heads=16, kv=8, num_global_kv=1 (MQA), head_dim=256, global_head_dim=512, ffn_inner=15360, sliding_window=1024, max_seq_len=262144, attention_k_eq_v=True, no PLE, no KV sharing, no MoE, final_logit_softcapping=30. * __init__.py: export get_gemma4_12b_config. * interop.py: extend always-strip multimodal prefixes with 'model.vision_embedder.' (Unified family ships a tower-free vision embedder: LN+Dense+LN+factorized 2D posemb+LN+RMSNorm+Linear). Audio remains conditional on audio_config. Verifies (offline): for google/gemma-4-12B safetensors (677 HF keys), convert_gemma4_state_dict produces 667 keys, matching the 667 keys of the freshly-constructed fairseq2 12B model on meta device exactly (no missing, no extras). --- src/fairseq2/models/gemma4/__init__.py | 1 + src/fairseq2/models/gemma4/config.py | 50 ++++++++++++++++++++++++++ src/fairseq2/models/gemma4/interop.py | 6 ++++ 3 files changed, 57 insertions(+) diff --git a/src/fairseq2/models/gemma4/__init__.py b/src/fairseq2/models/gemma4/__init__.py index d34fd3c22..5de439aab 100644 --- a/src/fairseq2/models/gemma4/__init__.py +++ b/src/fairseq2/models/gemma4/__init__.py @@ -30,6 +30,7 @@ from fairseq2.models.gemma4.config import get_gemma4_31b_config as get_gemma4_31b_config from fairseq2.models.gemma4.config import get_gemma4_e2b_config as get_gemma4_e2b_config from fairseq2.models.gemma4.config import get_gemma4_e4b_config as get_gemma4_e4b_config +from fairseq2.models.gemma4.config import get_gemma4_12b_config as get_gemma4_12b_config from fairseq2.models.gemma4.config import ( register_gemma4_configs as register_gemma4_configs, ) diff --git a/src/fairseq2/models/gemma4/config.py b/src/fairseq2/models/gemma4/config.py index e2629187c..06ddef0fb 100644 --- a/src/fairseq2/models/gemma4/config.py +++ b/src/fairseq2/models/gemma4/config.py @@ -240,6 +240,14 @@ 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() + def get_gemma4_e2b_config() -> Gemma4Config: """Get configuration for Gemma4 E2B (small dense, on-device). @@ -335,3 +343,45 @@ def get_gemma4_26b_a4b_config() -> Gemma4Config: moe_intermediate_size=704, final_logit_soft_cap=30.0, ) + + +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, + ) 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: From 6dcca1d3e928e65534b026166b9ea2ca899699e0 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 20:11:15 +0000 Subject: [PATCH 02/12] [gemma4] Asset cards + SFT recipe config for 12B (Unified) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * assets/cards/models/gemma4.yaml: register gemma4_12b and gemma4_12b_it cards (model_arch=12b/12b_it, checkpoint=hg://google/gemma-4-12B[-it], tokenizer_family=gemma4 — the unified family reuses the same HuggingFace tokenizer infrastructure and vocab). * recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml: mirror of gemma4_e4b_gsm8k.yaml with name=gemma4_12b, max_seq_len=4096, bf16 FSDP. Same optimizer/scheduler as e4b. Smoke-test target. --- recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml | 84 ++++++++++++++++++++ src/fairseq2/assets/cards/models/gemma4.yaml | 18 +++++ 2 files changed, 102 insertions(+) create mode 100644 recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml 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..1995a497b --- /dev/null +++ b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml @@ -0,0 +1,84 @@ +# 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" + name: "gemma4_12b" + dtype: bfloat16 + +tokenizer: + family: "gemma4" + name: "gemma4_12b" + +dataset: + max_seq_len: 4096 + max_num_tokens: 8192 + valid_split: "sft_test" + chat_mode: true + 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: false + project: "fairseq2" + run_name: "sft_gemma4_12b_gsm8k" + tensorboard: + enabled: false 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 From 07d6aa6db918ed26e1068e3020199b9ea8c642cd Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 20:31:41 +0000 Subject: [PATCH 03/12] [gemma4] SFT 12B config: use gemma4_12b_it tokenizer The base google/gemma-4-12B repo does not ship chat_template.jinja (only the -it variant does). For chat_mode=true SFT we need the chat template, so point the tokenizer at the -it asset card. The model weights are unchanged (the YAML's model.name=gemma4_12b is still the base checkpoint); only the tokenizer asset differs to pick up the chat template. --- recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml index 1995a497b..c969a1ac9 100644 --- a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml +++ b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml @@ -26,7 +26,11 @@ model: tokenizer: family: "gemma4" - name: "gemma4_12b" + # Use the -it tokenizer so that chat_template.jinja is loaded; the BASE + # google/gemma-4-12B repo does not ship chat_template.jinja since it is + # not meant to be chatted with. SFT with chat_mode=true requires the + # chat template; the model weights are unchanged so this is safe. + name: "gemma4_12b_it" dataset: max_seq_len: 4096 From 19a1c9f0d1eb9e2248f180cae20d061cebb3ae69 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 20:44:47 +0000 Subject: [PATCH 04/12] [gemma4] Point 12B cards at local /checkpoint dir; disable chat_mode for SFT * assets/cards/models/gemma4.yaml: gemma4_12b and gemma4_12b_it now use file:///checkpoint/smallomnillm/shared/models/gemma-4-12B[-it] paths (user-provided shared mirror), avoiding HF Hub cache races and the missing chat_template.jinja issue. * recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml: set chat_mode=false. Google's Gemma 4 chat_template.jinja does NOT use {% generation %} markers, so apply_chat_template(return_assistant_tokens_mask=True) returns all-zero assistant_masks and lm_sft's target_mask becomes all-false. With chat_mode=true the SFT runs but Number of Target Elements = 0 every step, NLL Loss stays at 0, and no gradient flows. Use chat_mode=false (LM continuation on src+tgt) to validate the recipe wiring end-to-end with a nonzero loss signal; revisit once the chat template is patched upstream to include {% generation %} blocks around the model turn. --- recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml | 11 ++++++++++- src/fairseq2/assets/cards/models/gemma4.yaml | 8 ++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml index c969a1ac9..1e08b46f5 100644 --- a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml +++ b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml @@ -36,7 +36,16 @@ dataset: max_seq_len: 4096 max_num_tokens: 8192 valid_split: "sft_test" - chat_mode: true + # 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: diff --git a/src/fairseq2/assets/cards/models/gemma4.yaml b/src/fairseq2/assets/cards/models/gemma4.yaml index 0052771a9..9f502361d 100644 --- a/src/fairseq2/assets/cards/models/gemma4.yaml +++ b/src/fairseq2/assets/cards/models/gemma4.yaml @@ -79,8 +79,8 @@ tokenizer_family: gemma4 name: gemma4_12b model_family: gemma4 model_arch: 12b -checkpoint: "hg://google/gemma-4-12B" -tokenizer: "hg://google/gemma-4-12B" +checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" +tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" tokenizer_family: gemma4 --- @@ -88,6 +88,6 @@ 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" +checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B-it" +tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B-it" tokenizer_family: gemma4 From 0e4cb569b56dc228435c15906cc60d9970e61f41 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 20:47:42 +0000 Subject: [PATCH 05/12] [gemma4] SFT yaml: fine-tune -it variant; fix base card path The base google/gemma-4-12B is at /checkpoint/fairseq2/shared/models (downloaded by exp 42); only the -it variant lives at the user-shared /checkpoint/smallomnillm tree. Switching the SFT yaml to fine-tune the -it variant (standard domain-adaptation pattern) so the SFT recipe finds its checkpoint in the user's tree. --- recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml | 10 +++++----- src/fairseq2/assets/cards/models/gemma4.yaml | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml index 1e08b46f5..99f2824c1 100644 --- a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml +++ b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml @@ -21,15 +21,15 @@ model: family: "gemma4" - name: "gemma4_12b" + # 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" - # Use the -it tokenizer so that chat_template.jinja is loaded; the BASE - # google/gemma-4-12B repo does not ship chat_template.jinja since it is - # not meant to be chatted with. SFT with chat_mode=true requires the - # chat template; the model weights are unchanged so this is safe. name: "gemma4_12b_it" dataset: diff --git a/src/fairseq2/assets/cards/models/gemma4.yaml b/src/fairseq2/assets/cards/models/gemma4.yaml index 9f502361d..5486184f9 100644 --- a/src/fairseq2/assets/cards/models/gemma4.yaml +++ b/src/fairseq2/assets/cards/models/gemma4.yaml @@ -79,8 +79,8 @@ tokenizer_family: gemma4 name: gemma4_12b model_family: gemma4 model_arch: 12b -checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" -tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" +checkpoint: "file:///checkpoint/fairseq2/shared/models/gemma-4-12B" +tokenizer: "file:///checkpoint/fairseq2/shared/models/gemma-4-12B" tokenizer_family: gemma4 --- From 9af1826cae4587875c967a711657d3737f070990 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 20:51:37 +0000 Subject: [PATCH 06/12] [gemma4] Tokenizer: add prompt_response encoding mode The SFT recipe (chat_mode=false path) calls create_encoder twice per example: mode='prompt' for the source and mode='prompt_response' for the target. Gemma4Tokenizer previously only accepted default/prompt/ as_is and raised ValueError on 'prompt_response', blocking the SFT recipe. Add prompt_response: no BOS prefix (the source half already has it), EOS suffix to terminate the target. Matches the Llama/Qwen pattern. --- src/fairseq2/models/gemma4/tokenizer.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/fairseq2/models/gemma4/tokenizer.py b/src/fairseq2/models/gemma4/tokenizer.py index e974c6f32..9a540fc81 100644 --- a/src/fairseq2/models/gemma4/tokenizer.py +++ b/src/fairseq2/models/gemma4/tokenizer.py @@ -57,15 +57,25 @@ 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 = [] From 35a23bc15dfb3904f0b6779dc74c59a053baae9c Mon Sep 17 00:00:00 2001 From: ProfAI Date: Sun, 7 Jun 2026 21:07:49 +0000 Subject: [PATCH 07/12] [gemma4] SFT yaml: enable wandb + tensorboard recorders Recipe-level wandb recorder writes per-step training metrics (NLL Loss, Gradient Norm, LR, throughput) into a wandb run with project= gemma-4-fairseq2. With WANDB_MODE=offline (default in our launch env since no API key is configured), the run materializes locally under /wandb/ and can be synced later via 'wandb sync '. --- recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml index 99f2824c1..91efb6b6d 100644 --- a/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml +++ b/recipes/lm/sft/configs/gemma4_12b_gsm8k.yaml @@ -90,8 +90,8 @@ common: seed: 0 metric_recorders: wandb: - enabled: false - project: "fairseq2" + enabled: true + project: "gemma-4-fairseq2" run_name: "sft_gemma4_12b_gsm8k" tensorboard: - enabled: false + enabled: true From 496cba7e39e36de34a381b4ce29bad9450097a85 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Mon, 8 Jun 2026 13:21:31 +0000 Subject: [PATCH 08/12] [gemma4] Wire audio embedder for Gemma 4 Unified family (12B) * audio/config.py: add audio_mode field ('conformer' | 'linear'). 'conformer' (default) preserves existing E4B / classic-Gemma 4 behaviour (mel-spec -> subsample Conv2d -> Conformer tower -> embedder). 'linear' is the new Gemma 4 Unified pipeline (raw waveform frames of audio_samples_per_token=640 -> embedder RMSNorm+Linear). Other Conformer-specific fields are ignored in linear mode. * config.py: add get_gemma4_unified_audio_config() (linear mode, output_proj_dims=640) and get_gemma4_12b_audio_config() that wraps the existing 12B text config with audio enabled. Register archs '12b_audio' and '12b_it_audio'. The original '12b' / '12b_it' archs stay text-only (no behaviour change for existing parity / MMLU / SFT runs). * factory.py: create_audio_tower returns None when audio_mode='linear' (the existing Gemma4MultimodalAudioEmbedder class is already exactly the unified embedder: RMSNorm(no scale) + Linear from 640 to text dim, no other changes needed). * model.py: Gemma4Model.forward now also handles the embedder-only case (tower=None, embedder!=None): feeds raw audio_features directly through the embedder. The conformer path is unchanged. Sanity check: 12b_audio model has 668 keys; converter on the HF gemma-4-12B-it safetensors produces exactly 668 keys (667 text + audio_embedder.embedding_projection.weight) -- zero diff. --- src/fairseq2/models/gemma4/audio/config.py | 29 ++++++++++--- src/fairseq2/models/gemma4/config.py | 48 ++++++++++++++++++++++ src/fairseq2/models/gemma4/factory.py | 13 +++++- src/fairseq2/models/gemma4/model.py | 15 ++++++- 4 files changed, 96 insertions(+), 9 deletions(-) diff --git a/src/fairseq2/models/gemma4/audio/config.py b/src/fairseq2/models/gemma4/audio/config.py index dbf68b415..97aee4df1 100644 --- a/src/fairseq2/models/gemma4/audio/config.py +++ b/src/fairseq2/models/gemma4/audio/config.py @@ -11,19 +11,38 @@ @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: str = "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/config.py b/src/fairseq2/models/gemma4/config.py index 06ddef0fb..8b7b28f9e 100644 --- a/src/fairseq2/models/gemma4/config.py +++ b/src/fairseq2/models/gemma4/config.py @@ -248,6 +248,14 @@ def _12b() -> Gemma4Config: 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). @@ -345,6 +353,27 @@ def get_gemma4_26b_a4b_config() -> Gemma4Config: ) +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). @@ -385,3 +414,22 @@ def get_gemma4_12b_config() -> Gemma4Config: 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..02feea4fb 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 getattr(config.audio_config, "audio_mode", "conformer") == "linear": + return None + from fairseq2.models.gemma4.audio.tower import Gemma4AudioTower return Gemma4AudioTower( 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, From 3a55b73904217c9ad46350313ff031ea87b8c824 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Mon, 8 Jun 2026 15:55:19 +0000 Subject: [PATCH 09/12] [gemma4] Co-locate 12B asset under /checkpoint/smallomnillm with the rest of the family User moved the 12B base safetensors from /checkpoint/fairseq2/shared/models to /checkpoint/smallomnillm/shared/models so all Gemma 4 variants live together (gemma-4-E*B, gemma-4-31B[-it], gemma-4-26B-A4B[-it], gemma-4-12B, gemma-4-12B-it). Source path deleted; asset card now points at the new location. The -it card was already on smallomnillm. --- src/fairseq2/assets/cards/models/gemma4.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/fairseq2/assets/cards/models/gemma4.yaml b/src/fairseq2/assets/cards/models/gemma4.yaml index 5486184f9..9f502361d 100644 --- a/src/fairseq2/assets/cards/models/gemma4.yaml +++ b/src/fairseq2/assets/cards/models/gemma4.yaml @@ -79,8 +79,8 @@ tokenizer_family: gemma4 name: gemma4_12b model_family: gemma4 model_arch: 12b -checkpoint: "file:///checkpoint/fairseq2/shared/models/gemma-4-12B" -tokenizer: "file:///checkpoint/fairseq2/shared/models/gemma-4-12B" +checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" +tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" tokenizer_family: gemma4 --- From 159c8335c771544de76f1a70004fa307ff2af9a3 Mon Sep 17 00:00:00 2001 From: ProfAI Date: Mon, 8 Jun 2026 16:03:44 +0000 Subject: [PATCH 10/12] [gemma4] Gemma4MultimodalAudioEmbedder: add cast_input_dtype flag Mirrors HF's split into two classes (Gemma4MultimodalEmbedder vs Gemma4UnifiedMultimodalEmbedder), but in a single fairseq2 class gated by a ctor flag. When cast_input_dtype=True, forward casts inputs_embeds to self.embedding_projection.weight.dtype before the RMSNorm -- exactly what HF's Gemma4UnifiedMultimodalEmbedder.forward does (modular_gemma4_unified.py:895-899). Without the cast, raw waveform features (fp32 from the feature extractor) fed into a bf16 embedder would silently use mismatched dtypes. Factory wires cast_input_dtype=True when audio_config.audio_mode == 'linear' (the Gemma 4 Unified family path). The default flag is False, so the classic E*/31B/26B-A4B Conformer audio path keeps its bit-for-bit behaviour. Exp 49 audio parity (bf16 cos=0.9992) passed before this change because the parity script pre-cast audio_features to DTYPE manually; this commit moves the cast into the embedder so any caller that forgets to pre-cast still produces HF-identical output. --- src/fairseq2/models/gemma4/audio/embedder.py | 26 ++++++++++++++++++++ src/fairseq2/models/gemma4/factory.py | 9 +++++++ 2 files changed, 35 insertions(+) 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/factory.py b/src/fairseq2/models/gemma4/factory.py index 02feea4fb..12ff4cbcb 100644 --- a/src/fairseq2/models/gemma4/factory.py +++ b/src/fairseq2/models/gemma4/factory.py @@ -567,10 +567,19 @@ 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 = ( + getattr(config.audio_config, "audio_mode", "conformer") == "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, ) From 3e060e18d825e0410537189cf9ed9d2a145ba0e6 Mon Sep 17 00:00:00 2001 From: Yunchao Yang Date: Tue, 16 Jun 2026 13:34:39 +0000 Subject: [PATCH 11/12] [gemma4] Fix isort/black lint errors in __init__.py and tokenizer.py --- src/fairseq2/models/gemma4/__init__.py | 3 ++- src/fairseq2/models/gemma4/tokenizer.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/fairseq2/models/gemma4/__init__.py b/src/fairseq2/models/gemma4/__init__.py index 5de439aab..81d94becf 100644 --- a/src/fairseq2/models/gemma4/__init__.py +++ b/src/fairseq2/models/gemma4/__init__.py @@ -24,13 +24,13 @@ ) 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, ) from fairseq2.models.gemma4.config import get_gemma4_31b_config as get_gemma4_31b_config from fairseq2.models.gemma4.config import get_gemma4_e2b_config as get_gemma4_e2b_config from fairseq2.models.gemma4.config import get_gemma4_e4b_config as get_gemma4_e4b_config -from fairseq2.models.gemma4.config import get_gemma4_12b_config as get_gemma4_12b_config from fairseq2.models.gemma4.config import ( register_gemma4_configs as register_gemma4_configs, ) @@ -84,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/tokenizer.py b/src/fairseq2/models/gemma4/tokenizer.py index 9a540fc81..68490a668 100644 --- a/src/fairseq2/models/gemma4/tokenizer.py +++ b/src/fairseq2/models/gemma4/tokenizer.py @@ -58,7 +58,10 @@ def create_encoder( raise ValueError(f"`lang` must be `None`, but is '{lang}' instead.") if mode is not None and mode not in ( - "default", "prompt", "prompt_response", "as_is" + "default", + "prompt", + "prompt_response", + "as_is", ): raise ValueError( "`mode` must be 'default', 'prompt', 'prompt_response', or " From d0ecda2022c8c0ccac78871c90544ee9d1c69b40 Mon Sep 17 00:00:00 2001 From: Yunchao Yang Date: Wed, 24 Jun 2026 15:24:38 +0000 Subject: [PATCH 12/12] [gemma4] Address PR review: hg:// 12B cards, Literal audio_mode, drop getattr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Point gemma4_12b{,_it} cards at hg://google/gemma-4-12B{,-it} instead of the internal /checkpoint/smallomnillm path so the OSS cards don't leak cluster-only paths. - Type Gemma4AudioConfig.audio_mode as Literal["conformer", "linear"] — the only two supported pipelines — for static checking. - Drop defensive getattr(config.audio_config, "audio_mode", ...) calls in Gemma4Factory; audio_mode is a proper dataclass field with a default, no BC story to preserve. --- src/fairseq2/assets/cards/models/gemma4.yaml | 8 ++++---- src/fairseq2/models/gemma4/audio/config.py | 3 ++- src/fairseq2/models/gemma4/factory.py | 6 ++---- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/fairseq2/assets/cards/models/gemma4.yaml b/src/fairseq2/assets/cards/models/gemma4.yaml index 9f502361d..0052771a9 100644 --- a/src/fairseq2/assets/cards/models/gemma4.yaml +++ b/src/fairseq2/assets/cards/models/gemma4.yaml @@ -79,8 +79,8 @@ tokenizer_family: gemma4 name: gemma4_12b model_family: gemma4 model_arch: 12b -checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" -tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B" +checkpoint: "hg://google/gemma-4-12B" +tokenizer: "hg://google/gemma-4-12B" tokenizer_family: gemma4 --- @@ -88,6 +88,6 @@ tokenizer_family: gemma4 name: gemma4_12b_it model_family: gemma4 model_arch: 12b_it -checkpoint: "file:///checkpoint/smallomnillm/shared/models/gemma-4-12B-it" -tokenizer: "file:///checkpoint/smallomnillm/shared/models/gemma-4-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/audio/config.py b/src/fairseq2/models/gemma4/audio/config.py index 97aee4df1..9fa4289ce 100644 --- a/src/fairseq2/models/gemma4/audio/config.py +++ b/src/fairseq2/models/gemma4/audio/config.py @@ -7,6 +7,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Literal @dataclass(kw_only=True) @@ -29,7 +30,7 @@ class Gemma4AudioConfig: and ``rms_norm_eps`` are read. """ - audio_mode: str = "conformer" + audio_mode: Literal["conformer", "linear"] = "conformer" """Audio pipeline selector: ``"conformer"`` or ``"linear"``.""" hidden_size: int = 1024 diff --git a/src/fairseq2/models/gemma4/factory.py b/src/fairseq2/models/gemma4/factory.py index 12ff4cbcb..10dbe22e2 100644 --- a/src/fairseq2/models/gemma4/factory.py +++ b/src/fairseq2/models/gemma4/factory.py @@ -541,7 +541,7 @@ def create_audio_tower(self) -> Module | None: return None # Unified family: no tower; embedder consumes raw waveform frames. - if getattr(config.audio_config, "audio_mode", "conformer") == "linear": + if config.audio_config.audio_mode == "linear": return None from fairseq2.models.gemma4.audio.tower import Gemma4AudioTower @@ -571,9 +571,7 @@ def create_audio_embedder(self) -> Module | None: # 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 = ( - getattr(config.audio_config, "audio_mode", "conformer") == "linear" - ) + cast_input_dtype = config.audio_config.audio_mode == "linear" return Gemma4MultimodalAudioEmbedder( output_proj_dims=config.audio_config.output_proj_dims,