Gemma4 text model support

[ruilin-gif]

🚀 The feature, motivation and pitch

Add apply_liger_kernel_to_gemma4_text patching support for the Gemma 4 text model family.

Gemma 4 was released by Google DeepMind on April 2, 2026, as their most capable open model family to date. The series includes 31B (Dense), 26B-A4B (MoE), E4B, and E2B variants, all under Apache 2.0 license. These models are seeing rapid adoption for fine-tuning and downstream tasks.

Gemma 4's text backbone (gemma4_text in HuggingFace transformers) introduces several architectural changes from Gemma 3:

• Hybrid Sliding/Full Attention with GQA, alternating between sliding window (512 tokens) and full attention layers
• GeGLU activation (gelu_pytorch_tanh) in the MLP, with conditional double-wide MLP for KV-shared layers
• Standard RMSNorm — unlike Gemma 3 which uses (1 + weight) * normed (zero-initialized weights with offset), Gemma 4 uses the standard weight * normed (one-initialized weights, no offset)
• Per-layer v_norm with with_scale=False (no learnable weight) in some attention heads
• Large vocabulary (262,144 tokens), making FLCE particularly beneficial for memory savings
• RoPE with dual configurations: theta=10000 for sliding attention, theta=1000000 with proportional scaling + partial rotary for full attention

Liger Kernel already supports the Gemma family:

• Gemma — via apply_liger_kernel_to_gemma
• Gemma2 — via apply_liger_kernel_to_gemma2
• Gemma3 (Text) — via apply_liger_kernel_to_gemma3_text
• Gemma3 (Multimodal) — via apply_liger_kernel_to_gemma3

Gemma 4 registers as a separate model type (gemma4_text) in HuggingFace transformers (v5.5.0+), so a dedicated patch function is needed for auto-patching to work (e.g., via AutoLigerKernelForCausalLM or the HF Trainer's use_liger_kernel=True).

The following kernels are applicable:

Kernel	Applicable?	Notes
RMSNorm	✅	Standard RMSNorm (offset=0.0), not the Gemma 3 offset variant. 6-7 instances per layer.
GeGLU	✅	Requires a custom wrapper (LigerGemma4TextMLP) due to (config, layer_idx) init signature and conditional double-wide MLP logic.
CrossEntropyLoss / FLCE	✅	Vocab size is 262,144 — FLCE memory savings are significant.
RoPE	❌	Gemma 4 uses apply_rotary_pos_emb(x, cos, sin) (single-tensor, called separately for q and k), incompatible with Liger's fused (q, k, cos, sin) variant.

Models: gemma-4-31B · gemma-4-31B-it · gemma-4-26B-A4B · gemma-4-E4B · gemma-4-E2B

I have a working implementation and would be happy to open a PR.

Alternatives

The Gemma 3 text patch (apply_liger_kernel_to_gemma3_text) cannot be reused directly because:

1. RMSNorm configuration differs (Gemma 3 uses offset=1.0, init_fn="zeros" while Gemma 4 uses offset=0.0, init_fn="ones")
2. The MLP class has a different __init__ signature ((config, layer_idx) with conditional double-wide logic)
3. The model type string is different (gemma4_text vs gemma3_text), so auto-patching would not trigger

Additional context

This issue covers the text-only model (Gemma4ForCausalLM). Multimodal support (Gemma4ForConditionalGeneration) with vision/audio encoders could be addressed in a follow-up issue.

Alternatives

No response

Additional context

No response

[minostauros]

It would be great if multimodal version is included as well. Thanks for your contribution!

[ruilin-gif]

Sure, will update my current issue and my local implementation to include multimodal version, I will open the PR soon

[dvdimitrov13]

Hi! First, thanks to @ruilin-gif for filing this and @lardinator for the excellent PR #1196 — the fused-linear-CE work there is exactly what we needed for the text path.

I'm new to open-source contribution, so I'd really welcome your advice on whether / how to proceed with what I'm about to describe.

Context: We've been training google/gemma-4-E4B-it (LoRA SFT on a Blackwell workstation, sm_120), which AutoModelForCausalLM instantiates as Gemma4ForConditionalGeneration — the multimodal wrapper class that PR #1196 explicitly scopes out. Without fused CE, the 262k-vocab fp32 logits allocation OOMs at long sequence length. PR #1196, once merged, won't help this class because its registration targets "gemma4_text" / Gemma4ForCausalLM only.

What we have locally: a working Gemma4ForConditionalGeneration.forward replacement that routes loss through LigerForCausalLMLoss — essentially PR #1196's causal_forward pattern adapted with:

• multimodal kwargs passthrough (pixel_values, input_features, mm_token_type_ids),
• config.get_text_config().hidden_size (hidden_size lives on text_config),
• final_logit_softcapping read from config.text_config,
• output class Gemma4CausalLMOutputWithPast preserved (with image_hidden_states / audio_hidden_states).

Numerical verification on the full 8B E4B-it model (bf16):

• Fused CE vs materialized-logits reference: abs diff 0.000000.
• End-to-end forward, patched vs stock SDPA: 100% top-1 token agreement, loss diff 0.0016, KL 3.5e-2.

No RoPE kernel swap (same rationale as #1196 — signature incompatibility).

What I'd like your guidance on:

1. Is this a PR you'd accept as a follow-up to [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) #1196? The issue body explicitly flags multimodal as a follow-up, but I want to confirm before investing.
2. Scope: add apply_liger_kernel_to_gemma4 (mm-class) alongside PR [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) #1196's apply_liger_kernel_to_gemma4_text, and register "gemma4" in MODEL_TYPE_TO_APPLY_LIGER_FN? Or would you prefer a single unified apply_liger_kernel_to_gemma4 that dispatches on class?
3. Caveat — PLE (Per-Layer Embeddings, used by E2B/E4B). My current patch doesn't explicitly handle PLE; it worked on E4B-it because PLE state flows through the base model forward unchanged. Should the follow-up scope to non-PLE variants, or should I verify PLE compatibility before opening the PR?
4. Ordering: I'd plan to wait for [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) #1196 to merge first so I can build on its file layout cleanly. Is that the right call, or would you prefer I open a PR earlier so the two can be reviewed together?
5. Anything I'm missing as a first-time contributor to this repo — CLA, test scaffolding expectations, maintainer norms, etc.?

Happy to share the current patch and verification methodology on request. Thanks in advance for any pointers.

---

🤖 Drafted with Claude Code (Claude Opus 4.7), reviewed and posted by me.

[Mecoli1219]

Hi @dvdimitrov13, thanks for the detailed write-up and for verifying numerics end-to-end. Yes, I think #1196 implementation looks great to me. We'll get it merged once it's ready.

We'd love multimodal support for Gemma4ForConditionalGeneration as well, so please feel free to contribute it as a follow-up.

For now, I'd prefer a single unified apply_liger_kernel_to_gemma4 that dispatches on class for simplicity, but feel free to propose it if you find a cleaner factoring during implementation.

I don't think we have PLE on liger today, and since you've already verified E4B that works correctly with PLE state flowing through unchanged, let's skip explicit PLE handling for now and focus on unblocking the OOM. We can revisit a dedicated PLE patch later as a separate effort.

Opening a draft PR now sounds great. It signals to the community what you're working on. You can rebase/merge main into your branch once #1196 lands.

Make sure you read the contribution guide. Please test on your local environment and specify the hardware you're using in the PR description. If you don't have access to a suitable GPU, feel free to tag me once the PR is ready and I'll run testing before merge.

[dvdimitrov13]

#take — for the multimodal slice (Gemma4ForConditionalGeneration) as the follow-up to #1196, per the discussion above. Will open a draft PR shortly so the scope is visible.

cc @Mecoli1219