[Gemma 4] Add multimodal support (apply_liger_kernel_to_gemma4 for Gemma4ForConditionalGeneration)

[dvdimitrov13]

Summary

Follow-up to #1196 (the text path) — adds apply_liger_kernel_to_gemma4 for Gemma4ForConditionalGeneration (multimodal class; includes E2B / E4B / E4B-it which are loaded by AutoModelForCausalLM as the multimodal class even when only text is being trained).

Closes the multimodal half of #1186.

Why

Gemma 4's text vocab is 262,144. Without FLCE the (B, T, V) bf16 logits tensor is ~17 GB at T=8192 (and ~34 GB once the loss path upcasts to fp32 for cross-entropy), which OOMs even 96 GB cards on Gemma4ForConditionalGeneration SFT — the OOM that originally motivated #1186. Routing loss through LigerForCausalLMLoss materializes only the loss scalar.

Shape

A single unified entry point that dispatches on class, per @Mecoli1219's preference in #1186:

• apply_liger_kernel_to_gemma4(model=Gemma4ForConditionalGeneration_instance) — multimodal path. Class-level RMSNorm + GeGLU swaps via apply_liger_kernel_to_gemma4_text, FLCE forward via the new multimodal_forward, recurses into model.model.language_model for instance-level patches.
• apply_liger_kernel_to_gemma4(model=Gemma4ForCausalLM_instance) — routes to apply_liger_kernel_to_gemma4_text for backwards compatibility, so the same entry point works for either shape.
• Registry: adds "gemma4": apply_liger_kernel_to_gemma4 alongside the existing "gemma4_text".

Drive-by fixes

Two isinstance(model, tuple_with_None_filter) sites (one in our new dispatcher, one in the existing apply_liger_kernel_to_gemma4_text) raised TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union when called under with patch("transformers.models.gemma4.modeling_gemma4"): getattr(MagicMock_module, "Gemma4TextForCausalLM", None) returns a MagicMock (not None) because MagicMock auto-creates attributes, so the cls is not None filter let it slip into the isinstance tuple. The text-path version was dormant — its existing test passes a Gemma4ForCausalLM which short-circuits the isinstance match before reaching the bad entry — but the multimodal recursive call into the text path passes a Gemma4TextModel, so no early match. Both sites now use isinstance(cls, type) as the filter.

Out of scope (deferred)

• Vision / audio tower kernel swaps. Gemma 4 loads both via AutoModel.from_config(config.{vision,audio}_config), so the module classes are polymorphic. Out of scope here; FLCE on the LM head is what unblocks training OOM.
• Gemma4MultimodalEmbedder / projector norms. Analogous to gemma3's mm_soft_emb_norm patching. Skipped for the same minimal-surface reason.
• PLE (Per-Layer Embeddings). State passes through the inner forward unchanged; verified end-to-end on google/gemma-4-E4B-it.
• MoE experts (Gemma4TextExperts). Guarded out via the same enable_moe_block check used by the text path in [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) #1196.
• Multimodal mini convergence test (mini_gemma4). Would need new test/resources/fake_configs/Google/Gemma4/... scaffolding for an image / audio processor — PR [Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) #1196 followed the same pattern (only added mini_gemma4_text to the non-multimodal convergence files). Happy to add it as a follow-up if you'd prefer it bundled here — let me know.

Testing Done

Hardware

• Hardware Type: RTX 5090 (Blackwell, sm_120, 32 GB), Vast.ai instance
• CUDA 13.0, NVIDIA driver 590.48.01
• Python 3.10.12, torch 2.11.0+cu130, transformers 5.7.0.dev0 (built from huggingface/transformers@main — gemma4 requires ≥ 5.5.0)
• bf16 throughout for end-to-end numerics

End-to-end numerical equivalence on real google/gemma-4-E4B-it

Verified before authoring this PR with our internal verify_patch_equivalence.py (same shape as #1196's verification):

Axis	Threshold	Measured
Fused-CE vs reference loss diff (seq=512)	< 5e-3	0.0016
End-to-end top-1 token agreement (seq=256, patched vs stock SDPA)	> 99 %	>99 %
End-to-end loss diff (patched vs stock SDPA)	< 5e-3	0.0016
Logit-distribution KL (informational)	—	~3.5e-2
SDPA EFFICIENT_ATTENTION vs MATH cos-sim (seq=4096)	> 0.9999	>0.9999

Liger-Kernel test gates

• make checkstyle — All checks passed!, 267 files already formatted
• make test — see log
• make test-convergence — see logs (per file)

make test

3131 passed, 903 skipped, 12 xfailed, 3 failed in 35:42.

Our two new unit tests pass cleanly:

PASSED  test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text                       (0.70s)
PASSED  test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation     (0.22s)

All six LigerGEGLUMLPForGemma4 edge-case tests added by #1196 also pass.

The 3 unrelated failures are pre-existing on the parent branch (untouched by this PR) and reproduce on main:

FAILED  test/transformers/test_grpo_loss.py::test_grpo_loss_with_bias_correction_kl[...]
FAILED  test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_paligemma
FAILED  test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation

Note: test/transformers/test_fused_moe.py was excluded via --ignore because it triggers a pre-existing collection error on main (ImportError: cannot import name 'compute_routing_metadata' from 'liger_kernel.ops' — the symbol exists in src/liger_kernel/ops/fused_moe.py but isn't re-exported from liger_kernel.ops.__init__.py). Worth a separate one-line fix in a follow-up.

make test-convergence (per file)

File	Result
fp32/test_mini_models.py	32 passed, 3 skipped
fp32/test_mini_models_multimodal.py	8 passed, 4 skipped, 1 xfailed, 1 failed (mini_qwen2_vl)
fp32/test_mini_models_with_logits.py	29 passed, 3 skipped, 1 xpassed
bf16/test_mini_models.py	33 passed, 2 failed (mini_llama4, mini_gemma4_text)
bf16/test_mini_models_multimodal.py	11 passed, 1 skipped, 2 failed (mini_qwen2_vl, mini_llama4)
bf16/test_mini_models_with_logits.py	30 passed, 1 skipped, 2 failed (mini_llama4, mini_qwen3_moe)

mini_gemma4_text passes in bf16/test_mini_models_with_logits.py and fp32/test_mini_models.py (PR #1196's text path). It fails only in bf16/test_mini_models.py with the same Blackwell bf16 logprob drift @eqy and reviewers flagged on #1196 (review comment r4321013177); this is in PR #1196's territory, not introduced by this multimodal patch. The other failures (mini_llama4, mini_qwen2_vl, mini_qwen3_moe) are also in models we don't touch and are pre-existing test instability on consumer Blackwell.

cc @Mecoli1219 @lardinator @ruilin-gif

---

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