mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Dimitar DimitrovandClaude Opus 4.7
decb1b73d7
## 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 #1196. - **Multimodal mini convergence test (`mini_gemma4`).** Would need new `test/resources/fake_configs/Google/Gemma4/...` scaffolding for an image / audio processor — PR #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 - [x] `make checkstyle` — `All checks passed!`, 267 files already formatted - [x] `make test` — see log - [x] `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 ``` #### `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](https://claude.com/claude-code) (Claude Opus 4.7), reviewed and posted by me. --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>