From decb1b73d7451fb6e8208b0daa25da40bbd00d56 Mon Sep 17 00:00:00 2001 From: Dimitar Dimitrov <60075474+dvdimitrov13@users.noreply.github.com> Date: Wed, 20 May 2026 22:37:33 +0100 Subject: [PATCH] [Gemma 4] Add multimodal support (apply_liger_kernel_to_gemma4 for Gemma4ForConditionalGeneration) (#1203) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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) --- README.md | 1 + src/liger_kernel/transformers/__init__.py | 3 + src/liger_kernel/transformers/model/gemma4.py | 146 ++++++++++++++++++ .../transformers/model/output_classes.py | 13 ++ src/liger_kernel/transformers/monkey_patch.py | 102 +++++++++++- .../bf16/test_mini_models_multimodal.py | 110 +++++++++++++ .../fp32/test_mini_models_multimodal.py | 109 +++++++++++++ .../gemma-4-e4b-it/tokenizer_config.json | 90 +++++++++++ test/transformers/test_monkey_patch.py | 83 ++++++++++ test/utils.py | 15 ++ 10 files changed, 671 insertions(+), 1 deletion(-) create mode 100644 test/resources/fake_configs/Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json diff --git a/README.md b/README.md index 9d5af6d2..229dec8e 100644 --- a/README.md +++ b/README.md @@ -260,6 +260,7 @@ loss.backward() | Gemma3 (Text) | `liger_kernel.transformers.apply_liger_kernel_to_gemma3_text` | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Gemma3 (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_gemma3` | LayerNorm, RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Gemma4 (Text) | `liger_kernel.transformers.apply_liger_kernel_to_gemma4_text` | RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy | +| Gemma4 (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_gemma4` | RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Paligemma, Paligemma2, & Paligemma2 Mix | `liger_kernel.transformers.apply_liger_kernel_to_paligemma` | LayerNorm, RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen2, Qwen2.5, & QwQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen2-VL, & QVQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl` | RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 6879ad74..038c5a1d 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -44,6 +44,7 @@ if TYPE_CHECKING: from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text # noqa: F401 + from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma4 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma4_text # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4 # noqa: F401 from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v # noqa: F401 @@ -120,6 +121,7 @@ def __getattr__(name: str): "apply_liger_kernel_to_gemma2", "apply_liger_kernel_to_gemma3", "apply_liger_kernel_to_gemma3_text", + "apply_liger_kernel_to_gemma4", "apply_liger_kernel_to_gemma4_text", "apply_liger_kernel_to_glm4", "apply_liger_kernel_to_glm4v", @@ -209,6 +211,7 @@ if _TRANSFORMERS_AVAILABLE: "apply_liger_kernel_to_gemma2", "apply_liger_kernel_to_gemma3", "apply_liger_kernel_to_gemma3_text", + "apply_liger_kernel_to_gemma4", "apply_liger_kernel_to_gemma4_text", "apply_liger_kernel_to_glm4", "apply_liger_kernel_to_glm4v", diff --git a/src/liger_kernel/transformers/model/gemma4.py b/src/liger_kernel/transformers/model/gemma4.py index 50e8e849..6c31f72c 100644 --- a/src/liger_kernel/transformers/model/gemma4.py +++ b/src/liger_kernel/transformers/model/gemma4.py @@ -11,6 +11,14 @@ from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast +try: + from liger_kernel.transformers.model.output_classes import LigerGemma4CausalLMOutputWithPast +except ImportError: + # Older transformers without gemma4 — multimodal_forward is then unreachable + # because monkey_patch.apply_liger_kernel_to_gemma4 imports gemma4 modules + # behind the same try/except. + LigerGemma4CausalLMOutputWithPast = None + logger = logging.get_logger(__name__) @@ -149,3 +157,141 @@ def causal_forward( token_accuracy=token_accuracy, predicted_tokens=predicted_tokens, ) + + +def multimodal_forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + input_features: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + input_features_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + image_position_ids: Optional[torch.LongTensor] = None, + video_position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + mm_token_type_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + skip_logits: Optional[bool] = None, + **lm_kwargs, +): + r"""Fused-linear-cross-entropy forward for ``Gemma4ForConditionalGeneration``. + + Mirrors :func:`liger_kernel.transformers.model.gemma3.multimodal_forward` + with Gemma 4-specific kwargs (``pixel_values_videos``, ``input_features``, + ``image_position_ids``, ``video_position_ids``, ``mm_token_type_ids``) and + output fields (``image_hidden_states``, ``audio_hidden_states``). + + The win on Gemma 4 multimodal is large: vocab=262,144 means the (B, T, V) + fp32 logits tensor is ~17 GB at T=8192 in bf16 (and another ~34 GB once the + loss path upcasts), OOMing 96 GB cards. Routing loss through + ``LigerForCausalLMLoss`` materializes only the loss scalar. + + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either + be in `[0, ..., config.text_config.vocab_size]` or -100 (see `input_ids` + docstring). Tokens with indices set to `-100` are ignored. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, + calculate logits for all `input_ids` (special case). If a `torch.Tensor`, + must be 1D corresponding to the indices to keep in the sequence-length + dimension. + """ + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + input_features=input_features, + attention_mask=attention_mask, + input_features_mask=input_features_mask, + position_ids=position_ids, + image_position_ids=image_position_ids, + video_position_ids=video_position_ids, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + inputs_embeds=inputs_embeds, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + **lm_kwargs, + ) + + hidden_states = outputs[0] + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + kept_hidden_states = hidden_states[:, slice_indices, :] + + text_cfg = self.config.get_text_config() + softcap = getattr(text_cfg, "final_logit_softcapping", None) + shift_labels = lm_kwargs.pop("shift_labels", None) + + loss = None + logits = None + token_accuracy = None + predicted_tokens = None + + if skip_logits is None: + skip_logits = self.training and (labels is not None or shift_labels is not None) + + if skip_logits: + result = LigerForCausalLMLoss( + hidden_states=kept_hidden_states, + lm_head_weight=self.lm_head.weight, + labels=labels, + shift_labels=shift_labels, + hidden_size=text_cfg.hidden_size, + final_logit_softcapping=softcap, + **lm_kwargs, + ) + loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result) + else: + logits = self.lm_head(kept_hidden_states) + if softcap is not None: + logits = logits / softcap + logits = torch.tanh(logits) + logits = logits * softcap + if labels is not None or shift_labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + shift_labels=shift_labels, + vocab_size=text_cfg.vocab_size, + **lm_kwargs, + ) + + if not return_dict: + output_tuple = (logits,) + outputs[1:] + output_tuple = (loss,) + output_tuple if loss is not None else output_tuple + output_tuple = output_tuple + (token_accuracy,) if token_accuracy is not None else output_tuple + output_tuple = output_tuple + (predicted_tokens,) if predicted_tokens is not None else output_tuple + return output_tuple + + return LigerGemma4CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=getattr(outputs, "image_hidden_states", None), + audio_hidden_states=getattr(outputs, "audio_hidden_states", None), + token_accuracy=token_accuracy, + predicted_tokens=predicted_tokens, + ) diff --git a/src/liger_kernel/transformers/model/output_classes.py b/src/liger_kernel/transformers/model/output_classes.py index 4fb5aaee..c2ef4a7c 100644 --- a/src/liger_kernel/transformers/model/output_classes.py +++ b/src/liger_kernel/transformers/model/output_classes.py @@ -19,6 +19,11 @@ try: except Exception: _Gemma3CausalLMOutputWithPast = None +try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4CausalLMOutputWithPast as _Gemma4CausalLMOutputWithPast +except Exception: + _Gemma4CausalLMOutputWithPast = None + try: from transformers.models.glm4v_moe.modeling_glm4v_moe import ( Glm4vMoeCausalLMOutputWithPast as _Glm4vMoeCausalLMOutputWithPast, @@ -108,6 +113,14 @@ if _Gemma3CausalLMOutputWithPast is not None: predicted_tokens: Optional[torch.LongTensor] = None +if _Gemma4CausalLMOutputWithPast is not None: + + @dataclass + class LigerGemma4CausalLMOutputWithPast(_Gemma4CausalLMOutputWithPast): + token_accuracy: Optional[torch.FloatTensor] = None + predicted_tokens: Optional[torch.LongTensor] = None + + if _Glm4vMoeCausalLMOutputWithPast is not None: @dataclass diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index c85c0b0b..83391cec 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -1356,7 +1356,11 @@ def apply_liger_kernel_to_gemma4_text( # The model instance already exists, so we need to additionally patch the # instance variables that reference already-instantiated modules - causal_lm_types = tuple(cls for cls in (Gemma4ForCausalLM, Gemma4TextForCausalLM) if cls is not None) + # `isinstance(cls, type)` filter (rather than `cls is not None`) so we + # also drop the MagicMock the test harness substitutes for + # `Gemma4TextForCausalLM` under `with patch("...modeling_gemma4")` — + # an `isinstance` check against a non-class entry raises TypeError. + causal_lm_types = tuple(cls for cls in (Gemma4ForCausalLM, Gemma4TextForCausalLM) if isinstance(cls, type)) if isinstance(model, causal_lm_types) or isinstance(model, Gemma4TextModel): # get the base model from the model instance base_model = model.model if isinstance(model, causal_lm_types) else model @@ -1385,6 +1389,101 @@ def apply_liger_kernel_to_gemma4_text( raise TypeError("The model must be Gemma4ForCausalLM, Gemma4TextForCausalLM, or Gemma4TextModel.") +def apply_liger_kernel_to_gemma4( + rope: bool = False, + cross_entropy: bool = False, + fused_linear_cross_entropy: bool = True, + layer_norm: bool = False, + rms_norm: bool = True, + geglu: bool = True, + model: PreTrainedModel = None, +) -> None: + """ + Apply Liger kernels to replace original implementation in HuggingFace Gemma4 + multimodal models (`Gemma4ForConditionalGeneration`). + + For text-only Gemma 4 (`Gemma4ForCausalLM`, `Gemma4TextForCausalLM`, + `Gemma4TextModel`), use :func:`apply_liger_kernel_to_gemma4_text` instead. + + The primary win is the fused-linear-cross-entropy path on the multimodal + class: with vocab=262,144, the (B, T, V) fp32 logits tensor is ~17 GB at + T=8192 in bf16 (and ~34 GB once the loss path upcasts), OOMing 96 GB cards. + Fused CE materializes only the loss scalar. + + Out of scope (deferred to future PRs): + - Vision and audio tower kernel swaps. Gemma 4's vision and audio towers + are loaded via `AutoModel.from_config(config.vision_config)` and + `AutoModel.from_config(config.audio_config)` respectively, so their + module classes are polymorphic. A safe class-level swap would need to + enumerate supported tower types — out of scope here; FLCE on the LM head + is what unblocks training OOM. + - PLE (Per-Layer Embeddings) kernels. PLE state passes through the inner + forward unchanged; verified end-to-end on E4B-it without explicit + handling. + - Gemma 4 MoE expert kernels (`Gemma4TextExperts`); guarded out via the + same `enable_moe_block` check used in the text path. + + Args: + rope (bool): Currently a no-op (HF's apply_rotary_pos_emb signature is + incompatible with Liger's fused variant). Default False. + cross_entropy (bool): Whether to apply Liger's cross entropy loss. + Default False. Mutually exclusive with `fused_linear_cross_entropy`. + fused_linear_cross_entropy (bool): Fused linear CE for memory + efficiency. Default True. + rms_norm (bool): Whether to apply Liger's RMSNorm. Default True. + geglu (bool): Whether to apply Liger's GeGLU MLP. Default True. + model (PreTrainedModel): An already-instantiated model to patch + in-place. Default None. + """ + assert not (cross_entropy and fused_linear_cross_entropy), ( + "cross_entropy and fused_linear_cross_entropy cannot both be True." + ) + + from transformers.models.gemma4 import modeling_gemma4 + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration + + from liger_kernel.transformers.model.gemma4 import multimodal_forward + + if model is not None and not isinstance(model, Gemma4ForConditionalGeneration): + raise TypeError("The model must be Gemma4ForConditionalGeneration.") + + # Class-level patches for the text decoder layers (RMSNorm, GeGLU MLP). + # We disable FLCE here because the multimodal class needs its own forward + # (handles pixel_values / input_features / mm_token_type_ids / etc.) — we + # install that below. + apply_liger_kernel_to_gemma4_text( + rope=rope, + cross_entropy=False, + fused_linear_cross_entropy=False, + rms_norm=rms_norm, + geglu=geglu, + ) + + if cross_entropy: + from transformers.loss.loss_utils import nn + + nn.functional.cross_entropy = liger_cross_entropy + + if fused_linear_cross_entropy: + if model is not None: + model.forward = MethodType(multimodal_forward, model) + else: + modeling_gemma4.Gemma4ForConditionalGeneration.forward = multimodal_forward + + if model is not None: + # Recurse into the language model for instance-level RMSNorm / GeGLU + # patching. (The class-level swap above already covers freshly + # instantiated modules; this catches the already-built ones.) + apply_liger_kernel_to_gemma4_text( + rope=rope, + cross_entropy=False, + fused_linear_cross_entropy=False, + rms_norm=rms_norm, + geglu=geglu, + model=model.model.language_model, + ) + + def apply_liger_kernel_to_paligemma( rope: bool = True, cross_entropy: bool = False, @@ -3364,6 +3463,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = { "gemma3_text": apply_liger_kernel_to_gemma3_text, "gemma3": apply_liger_kernel_to_gemma3, "gemma4_text": apply_liger_kernel_to_gemma4_text, + "gemma4": apply_liger_kernel_to_gemma4, "glm4": apply_liger_kernel_to_glm4, "glm4v": apply_liger_kernel_to_glm4v, "glm4v_moe": apply_liger_kernel_to_glm4v_moe, diff --git a/test/convergence/bf16/test_mini_models_multimodal.py b/test/convergence/bf16/test_mini_models_multimodal.py index 1a59cab7..265b8261 100644 --- a/test/convergence/bf16/test_mini_models_multimodal.py +++ b/test/convergence/bf16/test_mini_models_multimodal.py @@ -13,6 +13,7 @@ from transformers import PreTrainedTokenizerFast from transformers.models.siglip.configuration_siglip import SiglipVisionConfig from liger_kernel.transformers import apply_liger_kernel_to_gemma3 +from liger_kernel.transformers import apply_liger_kernel_to_gemma4 from liger_kernel.transformers import apply_liger_kernel_to_internvl from liger_kernel.transformers import apply_liger_kernel_to_llama4 from liger_kernel.transformers import apply_liger_kernel_to_llava @@ -40,6 +41,7 @@ from test.utils import load_tokenizer_config from test.utils import multimodal_collate_fn from test.utils import require_deterministic from test.utils import revert_liger_kernel_to_gemma3 +from test.utils import revert_liger_kernel_to_gemma4 from test.utils import revert_liger_kernel_to_internvl from test.utils import revert_liger_kernel_to_llama4 from test.utils import revert_liger_kernel_to_llava @@ -210,6 +212,21 @@ try: except ImportError: GEMMA3_AVAILABLE = False +try: + # Gemma4 multimodal requires transformers>=5.5.0 + from transformers.models.gemma4.configuration_gemma4 import Gemma4Config + from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig + from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig + from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor + from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration + from transformers.models.gemma4.processing_gemma4 import Gemma4Processor + from transformers.models.gemma4.video_processing_gemma4 import Gemma4VideoProcessor + + GEMMA4_AVAILABLE = True +except ImportError: + GEMMA4_AVAILABLE = False + try: from transformers.models.llama4.configuration_llama4 import Llama4Config from transformers.models.llama4.configuration_llama4 import Llama4TextConfig @@ -529,6 +546,52 @@ if GEMMA3_AVAILABLE: ), ) +if GEMMA4_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4"] = MiniModelConfig( + liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_gemma4, fused_linear_cross_entropy=False), + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4, + model_class=Gemma4ForConditionalGeneration, + mini_model_config=Gemma4Config( + text_config=Gemma4TextConfig( + vocab_size=32000, # 262144 + hidden_size=1024, # 3072 + intermediate_size=2048, # 24576 + num_hidden_layers=4, # 28 + num_attention_heads=4, # 16 + num_key_value_heads=2, # 8 + head_dim=128, # 256 + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + num_kv_shared_layers=0, + use_double_wide_mlp=False, + enable_moe_block=False, + hidden_size_per_layer_input=0, + ), + vision_config=Gemma4VisionConfig( + hidden_size=512, # 768 + intermediate_size=1024, # 3072 + num_hidden_layers=2, # 16 + num_attention_heads=4, # 12 + num_key_value_heads=4, # 12 + head_dim=128, # 64 + hidden_activation="gelu_pytorch_tanh", + patch_size=16, + position_embedding_size=1024, + ), + audio_config=None, # Audio out of scope (vision/audio tower follow-up PR) + image_token_id=5, # matches tokenizer "5": "" + boi_token_id=4, # matches tokenizer "4": "" + eoi_token_id=6, # matches tokenizer "6": "" + video_token_id=7, # dummy, video not tested + boa_token_id=4, # dummy, audio not tested + eoa_token_index=6, # dummy, audio not tested + attn_implementation="eager", + ), + ) + if QWEN2_VL_AVAILABLE: MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig( liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_qwen2_vl, fused_linear_cross_entropy=False), @@ -1231,6 +1294,34 @@ def create_processor(model_name: str): image_processor = Gemma3ImageProcessor() return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer) + elif model_name.startswith("mini_gemma4"): + tokenizer_config = load_tokenizer_config( + os.path.join( + FAKE_CONFIGS_PATH, + "Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json", + ) + ) + tokenizer_base = train_bpe_tokenizer( + [ + token.content + for key, token in sorted( + tokenizer_config["added_tokens_decoder"].items(), + key=lambda x: int(x[0]), + ) + ] + ) + fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config) + # Audio/video processors constructed with defaults; the convergence + # path only feeds image+text, so the audio/video branches in + # Gemma4Processor.__call__ are not exercised. Audio coverage is + # deferred to the vision/audio tower follow-up PR. + return Gemma4Processor( + feature_extractor=Gemma4AudioFeatureExtractor(), + image_processor=Gemma4ImageProcessor(), + tokenizer=fast_tokenizer, + video_processor=Gemma4VideoProcessor(), + ) + else: raise ValueError(f"Processor not available for model {model_name}") @@ -1662,6 +1753,25 @@ def run_mini_model_multimodal( ), ], ), + pytest.param( + "mini_gemma4", + 32, + 1e-5, + torch.bfloat16, + 5e-2, + 5e-2, + 1e-1, + 1e-1, + 1e-2, + 1e-2, + marks=[ + pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + pytest.mark.skipif( + not GEMMA4_AVAILABLE, + reason="Gemma4 not available in this version of transformers", + ), + ], + ), pytest.param( "mini_qwen3_5", 32, diff --git a/test/convergence/fp32/test_mini_models_multimodal.py b/test/convergence/fp32/test_mini_models_multimodal.py index df4f6bf1..9fd8d207 100644 --- a/test/convergence/fp32/test_mini_models_multimodal.py +++ b/test/convergence/fp32/test_mini_models_multimodal.py @@ -14,6 +14,7 @@ from transformers import PreTrainedTokenizerFast from transformers.models.siglip.configuration_siglip import SiglipVisionConfig from liger_kernel.transformers import apply_liger_kernel_to_gemma3 +from liger_kernel.transformers import apply_liger_kernel_to_gemma4 from liger_kernel.transformers import apply_liger_kernel_to_internvl from liger_kernel.transformers import apply_liger_kernel_to_llama4 from liger_kernel.transformers import apply_liger_kernel_to_llava @@ -41,6 +42,7 @@ from test.utils import load_tokenizer_config from test.utils import multimodal_collate_fn from test.utils import require_deterministic from test.utils import revert_liger_kernel_to_gemma3 +from test.utils import revert_liger_kernel_to_gemma4 from test.utils import revert_liger_kernel_to_internvl from test.utils import revert_liger_kernel_to_llama4 from test.utils import revert_liger_kernel_to_llava @@ -251,6 +253,21 @@ try: except ImportError: GEMMA3_AVAILABLE = False +try: + # Gemma4 multimodal requires transformers>=5.5.0 + from transformers.models.gemma4.configuration_gemma4 import Gemma4Config + from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig + from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig + from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor + from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration + from transformers.models.gemma4.processing_gemma4 import Gemma4Processor + from transformers.models.gemma4.video_processing_gemma4 import Gemma4VideoProcessor + + GEMMA4_AVAILABLE = True +except ImportError: + GEMMA4_AVAILABLE = False + try: # InternVL is only available in transformers>=4.52.1 from transformers.models.got_ocr2.image_processing_got_ocr2_fast import GotOcr2ImageProcessorFast @@ -558,6 +575,52 @@ if GEMMA3_AVAILABLE: ), ) +if GEMMA4_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4"] = MiniModelConfig( + liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_gemma4, fused_linear_cross_entropy=False), + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4, + model_class=Gemma4ForConditionalGeneration, + mini_model_config=Gemma4Config( + text_config=Gemma4TextConfig( + vocab_size=32000, # 262144 + hidden_size=1024, # 3072 + intermediate_size=2048, # 24576 + num_hidden_layers=4, # 28 + num_attention_heads=4, # 16 + num_key_value_heads=2, # 8 + head_dim=128, # 256 + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + num_kv_shared_layers=0, + use_double_wide_mlp=False, + enable_moe_block=False, + hidden_size_per_layer_input=0, + ), + vision_config=Gemma4VisionConfig( + hidden_size=512, # 768 + intermediate_size=1024, # 3072 + num_hidden_layers=2, # 16 + num_attention_heads=4, # 12 + num_key_value_heads=4, # 12 + head_dim=128, # 64 + hidden_activation="gelu_pytorch_tanh", + patch_size=16, + position_embedding_size=1024, + ), + audio_config=None, # Audio out of scope (vision/audio tower follow-up PR) + image_token_id=5, # matches tokenizer "5": "" + boi_token_id=4, # matches tokenizer "4": "" + eoi_token_id=6, # matches tokenizer "6": "" + video_token_id=7, # dummy, video not tested + boa_token_id=4, # dummy, audio not tested + eoa_token_index=6, # dummy, audio not tested + attn_implementation="eager", + ), + ) + if QWEN2_VL_AVAILABLE: MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig( liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_qwen2_vl, fused_linear_cross_entropy=False), @@ -1348,6 +1411,34 @@ def create_processor(model_name: str): image_processor = Gemma3ImageProcessor() return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer) + elif model_name.startswith("mini_gemma4"): + tokenizer_config = load_tokenizer_config( + os.path.join( + FAKE_CONFIGS_PATH, + "Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json", + ) + ) + tokenizer_base = train_bpe_tokenizer( + [ + token.content + for key, token in sorted( + tokenizer_config["added_tokens_decoder"].items(), + key=lambda x: int(x[0]), + ) + ] + ) + fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config) + # Audio/video processors constructed with defaults; the convergence + # path only feeds image+text, so the audio/video branches in + # Gemma4Processor.__call__ are not exercised. Audio coverage is + # deferred to the vision/audio tower follow-up PR. + return Gemma4Processor( + feature_extractor=Gemma4AudioFeatureExtractor(), + image_processor=Gemma4ImageProcessor(), + tokenizer=fast_tokenizer, + video_processor=Gemma4VideoProcessor(), + ) + else: raise ValueError(f"Processor not available for model {model_name}") @@ -1764,6 +1855,24 @@ def run_mini_model_multimodal( ), ], ), + pytest.param( + "mini_gemma4", + 32, + 1e-5, + torch.float32, + 1e-8, + 1e-4, + 5e-3, + 1e-5, + 5e-3, + 1e-5, + marks=[ + pytest.mark.skipif( + not GEMMA4_AVAILABLE, + reason="Gemma4 not available in this version of transformers", + ), + ], + ), pytest.param( "mini_qwen3_5", 32, diff --git a/test/resources/fake_configs/Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json b/test/resources/fake_configs/Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json new file mode 100644 index 00000000..6d7bdd67 --- /dev/null +++ b/test/resources/fake_configs/Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json @@ -0,0 +1,90 @@ +{ + "add_bos_token": true, + "add_eos_token": false, + "added_tokens_decoder": { + "0": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "1": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "2": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "3": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "4": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "5": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "6": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "7": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + } + }, + "boi_token": "", + "bos_token": "", + "chat_template": "{{ bos_token }}\n{%- if messages[0]['role'] == 'system' -%}\n {%- if messages[0]['content'] is string -%}\n {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}\n {%- else -%}\n {%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}\n {%- endif -%}\n {%- set loop_messages = messages[1:] -%}\n{%- else -%}\n {%- set first_user_prefix = \"\" -%}\n {%- set loop_messages = messages -%}\n{%- endif -%}\n{%- for message in loop_messages -%}\n {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}\n {{ raise_exception(\"Conversation roles must alternate user/assistant/user/assistant/...\") }}\n {%- endif -%}\n {%- if (message['role'] == 'assistant') -%}\n {%- set role = \"model\" -%}\n {%- else -%}\n {%- set role = message['role'] -%}\n {%- endif -%}\n {{ '' + role + '\n' + (first_user_prefix if loop.first else \"\") }}\n {%- if message['content'] is string -%}\n {{ message['content'] | trim }}\n {%- elif message['content'] is iterable -%}\n {%- for item in message['content'] -%}\n {%- if item['type'] == 'image' -%}\n {{ '' }}\n {%- elif item['type'] == 'text' -%}\n {{ item['text'] | trim }}\n {%- endif -%}\n {%- endfor -%}\n {%- else -%}\n {{ raise_exception(\"Invalid content type\") }}\n {%- endif -%}\n {{ '\n' }}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n {{'model\n'}}\n{%- endif -%}\n", + "clean_up_tokenization_spaces": false, + "eoi_token": "", + "eos_token": "", + "extra_special_tokens": { + "boi_token": "", + "eoi_token": "", + "image_token": "" + }, + "image_token": "", + "model_max_length": 1000000000000000019884624838656, + "pad_token": "", + "processor_class": "Gemma4Processor", + "sp_model_kwargs": null, + "spaces_between_special_tokens": false, + "tokenizer_class": "GemmaTokenizer", + "unk_token": "", + "use_default_system_prompt": false +} \ No newline at end of file diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 13073ce4..a879b285 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -1965,6 +1965,89 @@ def test_apply_liger_kernel_to_instance_for_gemma4_text(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +@pytest.mark.skipif(not is_gemma4_available(), reason="gemma4 module not available") +def test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation(): + # Ensure any monkey patching is cleaned up for subsequent tests + with patch("transformers.models.gemma4.modeling_gemma4"): + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration + + from liger_kernel.transformers.model.gemma4 import multimodal_forward as gemma4_multimodal_forward + + # Minimal dense-path text config — same knobs pinned off as the + # text-only test below (no PLE, MoE, KV-share, double-wide MLP). + text_config = transformers.models.gemma4.configuration_gemma4.Gemma4TextConfig( + dtype=torch.bfloat16, + rms_norm_eps=1e-5, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + num_kv_shared_layers=0, + use_double_wide_mlp=False, + enable_moe_block=False, + hidden_size_per_layer_input=0, + ) + # Vision/audio configs left as None — Gemma4Model wraps both towers in + # `if config._config is not None`, so a None-towers model still + # constructs as Gemma4ForConditionalGeneration and exercises the + # multimodal forward we're patching. The towers themselves are + # polymorphic (AutoModel.from_config) and not in this PR's scope. + config = transformers.models.gemma4.configuration_gemma4.Gemma4Config( + text_config=text_config, + vision_config=None, + audio_config=None, + ) + + dummy_model_instance = Gemma4ForConditionalGeneration._from_config(config) + assert isinstance(dummy_model_instance, Gemma4ForConditionalGeneration) + + # Pre-patch: forward and language-model norms must NOT be Liger. + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(gemma4_multimodal_forward) + assert inspect.getsource(dummy_model_instance.model.language_model.norm.forward) != inspect.getsource( + LigerRMSNorm.forward + ) + for layer in dummy_model_instance.model.language_model.layers: + assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerGEGLUMLP.forward) + assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.pre_feedforward_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.post_feedforward_layernorm.forward) != inspect.getsource( + LigerRMSNorm.forward + ) + assert inspect.getsource(layer.self_attn.q_norm.forward) != inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.self_attn.k_norm.forward) != inspect.getsource(LigerRMSNorm.forward) + + _apply_liger_kernel_to_instance(model=dummy_model_instance) + + # Post-patch: top-level forward is multimodal_forward, language_model + # norms / MLPs are Liger. + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(gemma4_multimodal_forward) + assert inspect.getsource(dummy_model_instance.model.language_model.norm.forward) == inspect.getsource( + LigerRMSNorm.forward + ) + for layer in dummy_model_instance.model.language_model.layers: + assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerGEGLUMLP.forward) + assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.pre_feedforward_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.post_feedforward_layernorm.forward) == inspect.getsource( + LigerRMSNorm.forward + ) + assert inspect.getsource(layer.self_attn.q_norm.forward) == inspect.getsource(LigerRMSNorm.forward) + assert inspect.getsource(layer.self_attn.k_norm.forward) == inspect.getsource(LigerRMSNorm.forward) + v_norm = getattr(layer.self_attn, "v_norm", None) + if v_norm is not None: + # with_scale=False → intentionally not patched. + assert inspect.getsource(v_norm.forward) != inspect.getsource(LigerRMSNorm.forward) + + try: + print(dummy_model_instance) + except Exception as e: + pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") + + def test_apply_liger_kernel_to_instance_for_qwen2(): # Ensure any monkey patching is cleaned up for subsequent tests with patch("transformers.models.qwen2.modeling_qwen2"): diff --git a/test/utils.py b/test/utils.py index 0af7c47c..831f00e8 100644 --- a/test/utils.py +++ b/test/utils.py @@ -505,6 +505,21 @@ def revert_liger_kernel_to_gemma4_text(model_config: MiniModelConfig): print("Liger kernel patches have been reverted.") +def revert_liger_kernel_to_gemma4(model_config: MiniModelConfig): + """Revert all Liger kernel patches applied to Gemma4 multimodal model.""" + + from transformers.models.gemma4 import modeling_gemma4 + + # Vision/audio towers are loaded via AutoModel.from_config, so their + # module classes are polymorphic — no class-level swap to revert there. + # Reloading modeling_gemma4 resets Gemma4RMSNorm / Gemma4TextMLP / + # Gemma4ForConditionalGeneration.forward, which is the surface the + # multimodal patch touches. + importlib.reload(modeling_gemma4) + model_config.model_class = modeling_gemma4.Gemma4ForConditionalGeneration + print("Liger kernel patches have been reverted.") + + def revert_liger_kernel_to_gemma3(model_config: MiniModelConfig): """ Revert all Liger kernel patches applied to Gemma3.