Add Liger Kernel support for Gemma 4 Unified multimodal (gemma4_unified)

[lefft]

Disclosure: This PR was written by Fable 5 at my (@lefft) direction, and also reviewed by me. It has been validated in a production long-context (64k) SFT pipeline with gemma-4-12B-it on an 8xH200 node. Resulting checkpoints produce expected performance gains on internal benchmark tasks (e.g. https://huggingface.co/datasets/lefft/needleif-bench). Our internal fork of liger-kernel includes the code in both #1311 and #1312. I've split it into two here for reviewability.

Summary

Adds Liger Kernel support for Gemma 4 Unified multimodal models
(Gemma4UnifiedForConditionalGeneration, model_type: gemma4_unified):
apply_liger_kernel_to_gemma4_unified, with RMSNorm, GeGLU, CrossEntropyLoss, and
FusedLinearCrossEntropy.

Liger Kernel already supports the other Gemma 4 families — gemma4_text (dense text, #1196) and
gemma4 (omni multimodal, #1203). But google/gemma-4-12B / gemma-4-12B-it load as
Gemma4UnifiedForConditionalGeneration (model_type: gemma4_unified) from their own
modeling_gemma4_unified module, which the existing patches don't reach — _apply_liger_kernel
no-ops with no Liger kernels supported for model type: gemma4_unified. This is the PR that
unblocks the 12B checkpoints; text-only fine-tuning of them also lands here, since that's the class
they instantiate as.

Closes #1308.

Stacked on #1311 (gemma4_unified_text), which adds the text apply function this one
calls and the shared LigerGemma4UnifiedCausalLMOutputWithPast output class. Please review after
PR1 merges; the diff shown here is against main and will shrink to just the multimodal delta once
PR1 lands. This is the same text→multimodal sequence gemma4 followed (#1196 → #1203).

Details

• Reuses PR1's text-stack patches (RMSNorm / GeGLU class swaps) by calling
  apply_liger_kernel_to_gemma4_unified_text internally, then installs the multimodal
  multimodal_forward on Gemma4UnifiedForConditionalGeneration.
• Both forwards return the upstream-shared Gemma4UnifiedCausalLMOutputWithPast (extended with
  liger's token_accuracy / predicted_tokens), passing shared_kv_states through alongside the
  image/audio hidden states.
• The FusedLinearCrossEntropy win is unusually large here: the family shares the 262,144-token
  vocabulary, so the [B, T, 262144] logits tensor is ~32 GiB in bf16 at T=65,536 (and ~64 GiB once
  the loss path upcasts to fp32), OOMing even 141 GB cards after the forward's activations. Fused CE
  materializes only the loss scalar.
• final_logit_softcapping (30.0 on the 12B checkpoint) flows through both loss paths.
• Out of scope (mirroring the gemma4 sibling's scope): the gemma4_unified_assistant model_type,
  LayerNorm kernels for the encoder-free vision embedder, and the audio path in convergence tests.
  layer_norm is accepted for API compatibility but applies no kernels.
• The multimodal convergence rows run with FLCE disabled by design (matching mini_gemma4); the
  fused path through the CondGen class was validated separately — see Testing Done.
• Rebased onto current main (post-Fix convergence test failures #1313), with the new mini_gemma4_unified rows given the same
  treatment Fix convergence test failures #1313 gave the mini_gemma4 sibling, since both families hit the same two upstream
  issues:
  • the validate_inputs operator-precedence bug (audio is not None and self.audio_token is None or self.boa_token is None or ...) raises the audio ValueError with no audio inputs whenever
    boa/eoa tokens are unset — fixed the same way, by setting dummy placeholders on the mini
    processor (audio is not exercised in convergence);
  • with that unblocked, the eval top-k logprobs comparison is fragile to near-tie token swaps, so
    the two mini_gemma4_unified rows carry the same xfail(strict=False) as their sibling rows.
    They will auto-flag XPASS if a future transformers release realigns the numerics.

Testing Done

Environment: Modal H100, torch 2.13.0+cu130, triton 3.7.1; all suites re-run after rebasing onto
main @ ed08f6e, at transformers 5.10.1 and 5.14.1 (current .[dev] resolve). Results
identical at both versions: every gemma4-filtered command exits 0.

make checkstyle:

All checks passed!
317 files left unchanged
===== EXIT 0: make checkstyle

test/transformers/test_monkey_patch.py -k gemma4 — 4/4 (both new unified instance tests plus both
gemma4 omni sibling tests as regression check):

test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text PASSED
test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation PASSED
test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_unified_text PASSED
test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_unified_conditional_generation PASSED
======================= 4 passed, 60 deselected in 2.15s =======================

FLCE-through-Gemma4UnifiedForConditionalGeneration numerical check — 20 fp32 optimizer 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):

step 00  hf=7.667602  liger=7.667601  delta=1.43e-06
...
step 19  hf=7.672588  liger=7.676995  delta=4.41e-03
first-step delta 1.43e-06, max rel delta 9.66e-04
generate OK: (1, 12)
FLCE-CondGen CHECK PASS

Text convergence rows (inherited from the stacked text PR, re-run on this branch) pass alongside the
sibling rows:

test/convergence/bf16/test_mini_models.py: 2 passed, 35 deselected in 15.11s
test/convergence/bf16/test_mini_models_with_logits.py: 2 passed, 33 deselected in 11.96s

Multimodal convergence — the new mini_gemma4_unified rows now land as xfail in lockstep with the
mini_gemma4 sibling rows, so both files exit 0:

test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_gemma4-32-1e-05-dtype12-0.05-0.05-0.1-0.1-0.01-0.01] XFAIL
test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_gemma4_unified-32-1e-05-dtype13-0.05-0.05-0.1-0.1-0.01-0.01] XFAIL
================ 15 deselected, 2 xfailed, 4 warnings in 23.96s ================

test/convergence/fp32/test_mini_models_multimodal.py: 15 deselected, 2 xfailed, 4 warnings in 23.98s

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).

• Hardware Type: H100 (Modal)
• run make test to ensure correctness — on this branch's rebased HEAD (transformers 5.14.1): 3907 passed, 1181 skipped, 14 xfailed in 43:06 — zero failures
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence — gemma4-filtered rows as above; all commands exit 0 at both transformers versions

[lefft]

Rebased onto current main (ed08f6e) — restacked on the rebased #1311 — and re-validated.

#1313 landed after this PR was opened and gave the mini_gemma4 multimodal rows two things the new
mini_gemma4_unified rows needed as well, since both families hit the same upstream issues:

1. dummy boa/eoa tokens on the mini processor, working around the validate_inputs
   operator-precedence bug (audio is not None and self.audio_token is None or self.boa_token is None or ...) that raises the audio ValueError even with no audio inputs;
2. xfail(strict=False) on the convergence rows, whose eval top-k logprobs comparison is fragile
   to near-tie token swaps.

Applied both to the unified rows in the bf16 and fp32 files. With the processor unblocked, the new
rows now land as xfail in lockstep with their sibling rows instead of erroring out, and both
multimodal convergence files exit 0.

Re-ran everything on the rebased HEAD at transformers 5.10.1 and 5.14.1 (Modal H100, torch
2.13.0+cu130, triton 3.7.1). Every command exits 0 at both versions. Description updated with the
fresh logs.

• make checkstyle: pass
• test_monkey_patch.py -k gemma4: 4 passed
• FLCE-through-Gemma4UnifiedForConditionalGeneration check: pass (fused ≡ unfused at step 0, delta
  1.43e-06; max relative loss drift 9.66e-04 over 20 fp32 steps)
• text convergence rows: 2 passed in each of the two bf16 files
• multimodal convergence: 2 xfailed in each file (mini_gemma4 + mini_gemma4_unified)
• full make test on this branch's rebased HEAD: 3907 passed, 1181 skipped, 14 xfailed, 0 failed