Add Liger Kernel support for Gemma 4 Unified (gemma4_unified)

[lefft]

Summary

Adds Liger Kernel support for the Gemma 4 Unified models (model_type: gemma4_unified / gemma4_unified_text, e.g. google/gemma-4-12B-it): apply_liger_kernel_to_gemma4_unified and apply_liger_kernel_to_gemma4_unified_text with RMSNorm, GeGLU, CrossEntropyLoss, and FusedLinearCrossEntropy, mirroring the gemma4 (omni) support from linkedin#1203.

Resolves linkedin#1308.

These checkpoints currently no-op with no Liger kernels supported for model type: gemma4_unified — the gemma4 (omni) patches can't reach them because the unified models live in their own transformers module (separate class objects). The family shares the 262,144-token vocabulary, so the FusedLinearCrossEntropy win is unusually large at any sequence length: the [B, T, 262144] logits tensor is ~4 GiB in bf16 per 8k-token row (~32 GiB at 64k), doubling again when the loss path upcasts to fp32.

Details

• No new kernels. Gemma4UnifiedRMSNorm / Gemma4UnifiedTextMLP are implementation-identical to gemma4's (ones-init no-offset RMSNorm with fp32 compute and a with_scale=False variant; GeGLU with the double-wide KV-shared handling), so the existing LigerRMSNormForGemma4 / LigerGEGLUMLPForGemma4 wrappers are reused — only the patch targets change (modeling_gemma4_unified.*). Cross-module application of the gemma4 patch functions is not possible (different class objects), hence the parallel apply functions.
• Both forwards return the upstream-shared Gemma4UnifiedCausalLMOutputWithPast (extended with liger's token_accuracy / predicted_tokens), passing shared_kv_states through. final_logit_softcapping (30.0 on the 12B checkpoint) flows through both the fused and non-fused loss paths.
• The FLCE instance rebind is guarded to Gemma4UnifiedForCausalLM — a bare Gemma4UnifiedTextModel has no lm_head, so binding the causal forward to it would crash on first use. (The gemma4 sibling has the unguarded version of this; happy to fix it in a follow-up.)
• rope remains an accepted no-op (upstream uses a single-tensor apply_rotary_pos_emb, incompatible with liger_rotary_pos_emb), matching gemma4.
• Out of scope: the gemma4_unified_assistant model_type, LayerNorm kernels for the encoder-free vision embedder, and the audio path in convergence tests (all mirroring the sibling's scope).
• Test-selection note: pytest -k gemma4 now also matches the gemma4_unified tests; use -k "gemma4 and not unified" to isolate the omni family.

Testing Done

Environment: Modal H100, torch 2.13.0, triton 3.7.1; suites run at both transformers 5.10.1 (the version we deploy) and 5.14.1 (current floated resolve of .[dev]).

New-code results (both transformers versions unless noted):

• make checkstyle: pass.
• test/transformers/test_monkey_patch.py -k gemma4: 4/4 pass (both new instance tests + both sibling tests).
• Convergence, new text rows: mini_gemma4_unified_text passes bf16/test_mini_models.py and bf16/test_mini_models_with_logits.py.
• FLCE-through-Gemma4UnifiedForConditionalGeneration numerical check (20 fp32 training steps, HF vs liger from identical init; the multimodal convergence harness runs with FLCE disabled by design, so this covers the fused path through the CondGen class): fused ≡ unfused at step 0 (delta 1.4e-06), max relative loss drift 9.7e-04 over 20 optimizer steps, generate() clean after patching.
• Real-workload validation: a 64k-sequence full fine-tune of google/gemma-4-12B-it (transformers 5.10.1, torch 2.13.0, FSDP full-shard + flex_attention, 8×H100-class) that OOMs at step 1 without this port (128.6/139.8 GiB per GPU already used after the forward; the unfused loss then tries to materialize ~32 GiB of logits) trains cleanly with it — finite, decreasing loss (1.98 → 0.91 over the first 15 steps).

Pre-existing failures in these environments, reproduced on upstream main without this PR (control run at bb715c9):

• mini_gemma4_text (omni sibling) bf16 convergence fails at transformers ≥5.10 (tolerance drift vs the version it was calibrated on).

• Multimodal convergence fails for BOTH families: at 5.10.1 due to an operator-precedence bug in the upstream processors (audio is not None and self.audio_token is None or self.boa_token is None or ... raises "Audio inputs were provided" with no audio when the test tokenizer lacks boa_token); at 5.14.1 due to family-wide logprob drift. The new mini_gemma4_unified rows behave identically to the sibling rows in these environments — they are included for parity and will pass wherever the sibling's do.

• Hardware Type: H100 (Modal)

• run make test to ensure correctness — full suite (transformers 5.14.1, torch 2.13.0): 3905 passed, 1181 skipped, 4 failed; all 4 failures are bf16-tolerance flakes in kernels untouched by this diff (test_poly_norm ×2, test_fused_add_rms_norm ×2) on this torch/GPU combo — every gemma4/gemma4_unified unit test passes

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence (gemma4-filtered runs as above; full-matrix behavior documented)

[lefft]

Review — style/convention, verbosity, sibling parity, framing

Reviewed against the direct sibling (gemma4 omni, linkedin#1203 + text linkedin#1196) and five other model-support PRs. Overall this is a faithful, almost-purely-additive mirror of the gemma4 port (+1216 / −0, 14 files). It's in good shape. Notes below, grouped by review dimension.

1. Size & shape vs. sibling model-support PRs

PR	+adds	−dels	files	scope
this PR	1216	0	14	gemma4_unified text + multimodal
linkedin#1196 gemma4 (text)	722	0	11	gemma4 text only
linkedin#1203 gemma4 (omni/multimodal)	671	1	10	gemma4 multimodal only
linkedin#1166 Ministral	597	96	10	dense text
linkedin#1165 Nemotron	431	0	10	dense text
linkedin#253 Pixtral	499	0	8	multimodal
linkedin#949 gpt-oss	434	1	8	MoE text

This is the largest single model-support PR in the set, but that's because it bundles text + multimodal in one PR, which gemma4 shipped as two (linkedin#1196 + linkedin#1203 = ~1393 adds combined). Per-scope this is actually slightly leaner than the sibling, and it's 0-deletion / additive, so risk is low. Worth deciding deliberately: keep as one PR, or split into ..._text then ..._unified to mirror the gemma4 review sequence and halve per-PR reviewer load. Either is defensible; the gemma4 precedent is the split.

2. Convention / template

• PR template: Testing Done + Hardware Type + the 3 checkboxes are present. make test is unchecked ("in flight"). Before flipping out of draft, check it (or state why N/A), and consider pasting the raw test_monkey_patch.py -k gemma4 and convergence pass logs — docs/contributing.md explicitly asks contributors to paste successful test logs, and [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) linkedin/Liger-Kernel#1196 did (checkstyle + pytest snippets). Right now results are described in prose only.
• Registry / __init__ ordering: matches existing conventions — __all__/imports are alphabetical (gemma4 < gemma4_text < gemma4_unified < gemma4_unified_text); MODEL_TYPE_TO_APPLY_LIGER_FN follows the file's text-variant-first grouping (..._text before base), same as the gemma4 pair. ✅
• create_processor prefix ordering: the startswith("mini_gemma4_unified") branch is correctly placed before the broader startswith("mini_gemma4") branch, with a comment saying so. Good catch — that collision would otherwise route unified through the omni tokenizer. ✅
• Shared-class reuse documented: docstring notes added to LigerRMSNormForGemma4 / LigerGEGLUMLPForGemma4 explaining they're reused for gemma4_unified. Matches liger's preference for documenting cross-model kernel reuse. ✅

3. Deliberate divergences from the sibling (both look correct — calling them out so review is explicit)

• FLCE instance-rebind is guarded to Gemma4UnifiedForCausalLM (monkey_patch.py, apply_liger_kernel_to_gemma4_unified_text), skipping a bare Gemma4UnifiedTextModel that has no lm_head. The gemma4 sibling binds causal_forward unconditionally when model is not None (monkey_patch.py:1347-1348), which would crash on a bare Gemma4TextModel + default fused_linear_cross_entropy=True. This PR's version is more correct. The PR body already flags it and offers a sibling follow-up — good. One judgment call for the maintainers: keep the improvement here + fix gemma4 separately, or align both in this PR. Flagging so a "why doesn't this match gemma4 exactly?" comment doesn't stall review.
• causal_forward returns LigerGemma4UnifiedCausalLMOutputWithPast (carrying shared_kv_states), whereas the gemma4 sibling's causal_forward returns the generic LigerCausalLMOutputWithPast (gemma4.py:151). The unified causal model emits shared_kv_states, so the generic class would silently drop it — the model-specific class is the faithful choice. getattr(outputs, "shared_kv_states", None) keeps it safe if absent. Justified; the PR body mentions it, though stating the reason (shared_kv_states) inline where the output class is introduced would save reviewers the inference.

4. Verbosity

Comment density is at parity with the sibling gemma4.py / apply_liger_kernel_to_gemma4, which is the right calibration for this repo — I would not strip the mirrored comments (the HF-style labels/logits_to_keep docstrings, the _maybe_patch_scaled_norm explainer, the rope no-op block all have 1:1 analogues in the sibling). No excessive edge-case handling — the getattr(..., with_scale, True) / getattr(self.config, "final_logit_softcapping", None) guards are all present in the sibling too.

The only comments that exceed the sibling are the net-new "Unlike gemma4 (omni), …" comparison notes (e.g. gemma4_unified.py:138-139 on the omitted eager-attn warning; the "no PLE / MoE fields" config comments). These are genuinely useful during review but read as scaffolding long-term. Optional: trim the 1–2 most meta ones before merge. Not blocking.

5. Framing (feature-first, not platform-first)

Mostly clean — the motivation is generalized correctly (memory-gated CE at vocab=262,144; the concrete 64k-FFT-OOMs-at-step-1 → trains-cleanly data point is exactly the right upstream justification for FLCE, and reads as a general training story, not an internal one). One small leak to neutralize:

• "transformers 5.10.1 (the version we deploy)" (PR body; also in the issue). Drop "the version we deploy" — testing at two transformers versions is a strength on its own; state it neutrally, e.g. "run at transformers 5.10.1 and 5.14.1 (current .[dev] resolve)." No other us/our/platform references in the diff or commit.

Minor

• test/.../gemma-4-12B-it/tokenizer_config.json has no trailing newline (\ No newline at end of file). make checkstyle passed, so likely fine, but worth a look if pre-commit's end-of-file-fixer runs in CI.
• The -k gemma4 collision note in the PR body (-k "gemma4 and not unified") is a nice courtesy — keep it.

[lefft]

Thanks — integrated the review and restructured for upstream. This draft is now superseded by two
stacked PRs against linkedin/Liger-Kernel, mirroring how gemma4 shipped (linkedin#1196 text → linkedin#1203 mm):

• Text: Add Liger Kernel support for Gemma 4 Unified text (gemma4_unified_text) linkedin/Liger-Kernel#1311 (gemma4_unified_text)
• Multimodal: Add Liger Kernel support for Gemma 4 Unified multimodal (gemma4_unified) linkedin/Liger-Kernel#1312 (gemma4_unified, stacked on the text PR)

Changes vs this draft: split text/multimodal into the two PRs above; guarded FLCE rebind retained
and called out; neutralized the "version we deploy" framing; real checkstyle/monkey-patch logs
pasted into each PR. Upstream review continues on those two PRs; this one stays open for other work.