Add Liger Kernel support for Gemma 4 Unified text (gemma4_unified_text)

README.md

@@ -308,6 +308,7 @@ loss.backward()
 | 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         |
+| Gemma4 Unified (Text)      | `liger_kernel.transformers.apply_liger_kernel_to_gemma4_unified_text`   | 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        |

src/liger_kernel/transformers/__init__.py

@@ -48,6 +48,7 @@
     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_gemma4_unified_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
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v_moe  # noqa: F401
@@ -125,6 +126,7 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_gemma3_text",
         "apply_liger_kernel_to_gemma4",
         "apply_liger_kernel_to_gemma4_text",
+        "apply_liger_kernel_to_gemma4_unified_text",
         "apply_liger_kernel_to_glm4",
         "apply_liger_kernel_to_glm4v",
         "apply_liger_kernel_to_glm4v_moe",
@@ -216,6 +218,7 @@ def __getattr__(name: str):
             "apply_liger_kernel_to_gemma3_text",
             "apply_liger_kernel_to_gemma4",
             "apply_liger_kernel_to_gemma4_text",
+            "apply_liger_kernel_to_gemma4_unified_text",
             "apply_liger_kernel_to_glm4",
             "apply_liger_kernel_to_glm4v",
             "apply_liger_kernel_to_glm4v_moe",

src/liger_kernel/transformers/geglu.py

@@ -28,6 +28,8 @@ class LigerGEGLUMLPForGemma4(LigerGEGLUMLP):
     HF's Gemma4TextMLP conditionally doubles intermediate_size for KV-shared layers
     when ``config.use_double_wide_mlp=True``. This subclass replicates that logic
     so the class-level swap works for all Gemma 4 variants (31B text, future MoE).
+    Also swapped in for ``Gemma4UnifiedTextMLP`` (gemma4_unified), which is
+    implementation-identical including the double-wide handling.
 
     See: https://github.com/huggingface/transformers/blob/74a2a4d0c/src/transformers/models/gemma4/modeling_gemma4.py#L1030-L1035
     """

src/liger_kernel/transformers/model/gemma4_unified.py

@@ -0,0 +1,151 @@
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+
+from transformers.cache_utils import Cache
+
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
+
+try:
+    from liger_kernel.transformers.model.output_classes import LigerGemma4UnifiedCausalLMOutputWithPast
+except ImportError:
+    # Older transformers without gemma4_unified — the forward is then
+    # unreachable because monkey_patch.apply_liger_kernel_to_gemma4_unified_text
+    # imports gemma4_unified modules behind the same try/except.
+    LigerGemma4UnifiedCausalLMOutputWithPast = None
+
+
+def causal_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Cache] = 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,
+    **loss_kwargs,
+) -> Union[Tuple, "LigerGemma4UnifiedCausalLMOutputWithPast"]:
+    r"""
+        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.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        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). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Fused-linear-cross-entropy forward for Gemma4UnifiedForCausalLM. Mirrors
+    liger's gemma4 causal_forward. google/gemma-4-12B-it ships
+    final_logit_softcapping=30.0, so the softcap flows through both the fused
+    and non-fused loss paths.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, Gemma4UnifiedForCausalLM
+
+    >>> model = Gemma4UnifiedForCausalLM.from_pretrained("google/gemma-4-12B-it")  # illustrative slug
+    >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-12B-it")
+
+    >>> prompt = "What is your favorite condiment?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "What is your favorite condiment?"
+    ```"""
+
+    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,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        **loss_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, :]
+    shift_labels = loss_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:
+        # final_logit_softcapping via getattr: some future Gemma 4 variants may omit the attribute entirely.
+        result = LigerForCausalLMLoss(
+            hidden_states=kept_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            shift_labels=shift_labels,
+            hidden_size=self.config.hidden_size,
+            final_logit_softcapping=getattr(self.config, "final_logit_softcapping", None),
+            **loss_kwargs,
+        )
+        loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)
+    else:
+        logits = self.lm_head(kept_hidden_states)
+        final_logit_softcapping = getattr(self.config, "final_logit_softcapping", None)
+        if final_logit_softcapping is not None:
+            logits = logits / final_logit_softcapping
+            logits = torch.tanh(logits)
+            logits = logits * final_logit_softcapping
+        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=self.vocab_size,
+                **loss_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 LigerGemma4UnifiedCausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        shared_kv_states=getattr(outputs, "shared_kv_states", None),
+        token_accuracy=token_accuracy,
+        predicted_tokens=predicted_tokens,
+    )

src/liger_kernel/transformers/model/output_classes.py

@@ -24,6 +24,13 @@
 except Exception:
     _Gemma4CausalLMOutputWithPast = None
 
+try:
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+        Gemma4UnifiedCausalLMOutputWithPast as _Gemma4UnifiedCausalLMOutputWithPast,
+    )
+except Exception:
+    _Gemma4UnifiedCausalLMOutputWithPast = None
+
 try:
     from transformers.models.glm4v_moe.modeling_glm4v_moe import (
         Glm4vMoeCausalLMOutputWithPast as _Glm4vMoeCausalLMOutputWithPast,
@@ -121,6 +128,14 @@ class LigerGemma4CausalLMOutputWithPast(_Gemma4CausalLMOutputWithPast):
         predicted_tokens: Optional[torch.LongTensor] = None
 
 
+if _Gemma4UnifiedCausalLMOutputWithPast is not None:
+
+    @dataclass
+    class LigerGemma4UnifiedCausalLMOutputWithPast(_Gemma4UnifiedCausalLMOutputWithPast):
+        token_accuracy: Optional[torch.FloatTensor] = None
+        predicted_tokens: Optional[torch.LongTensor] = None
+
+
 if _Glm4vMoeCausalLMOutputWithPast is not None:
 
     @dataclass

src/liger_kernel/transformers/monkey_patch.py

@@ -1484,6 +1484,142 @@ def apply_liger_kernel_to_gemma4(
         )
 
 
+def apply_liger_kernel_to_gemma4_unified_text(
+    rope: bool = False,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    geglu: bool = True,
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Gemma4
+    Unified text models (Gemma4UnifiedForCausalLM / Gemma4UnifiedTextModel).
+
+    Primary target: google/gemma-4-12B-it (model_type "gemma4_unified", text
+    stack "gemma4_unified_text"). Unlike gemma4 (omni), the unified text stack
+    has no PLE and no MoE; every decoder layer is a plain
+    (norm, attn, norm, mlp, norm) stack with optional KV sharing (the 12B
+    config sets num_kv_shared_layers=0).
+
+    Known limitation: rope kernel swap is a no-op on Gemma 4 Unified — HF's
+    apply_rotary_pos_emb takes a single tensor at a time, which is incompatible
+    with Liger's (q, k, cos, sin) signature. HF's plain pytorch rope stays in
+    place. The large training-memory win (fused linear cross-entropy: a ~32 GiB
+    bf16 logits tensor eliminated at seq 65536 / vocab 262144) is unaffected.
+
+    Args:
+        rope (bool): Currently a no-op for Gemma 4 Unified (HF uses single-tensor
+            apply_rotary_pos_emb incompatible with Liger). Default False.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default False.
+        fused_linear_cross_entropy (bool): Fused linear CE for memory efficiency. Default True.
+            Mutually exclusive with `cross_entropy`.
+        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.
+    """
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.gemma4_unified import modeling_gemma4_unified
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForCausalLM
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedTextDecoderLayer
+    from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedTextModel
+
+    from liger_kernel.transformers.model.gemma4_unified import causal_forward
+    from liger_kernel.transformers.rms_norm import LigerRMSNormForGemma4
+
+    # Gemma4UnifiedRMSNorm is identical to Gemma4RMSNorm (ones-init, no +1
+    # offset, fp32 compute), so the Gemma4 wrapper classes are reused; only
+    # the patch targets differ (modeling_gemma4_unified.*).
+    _patch_rms_norm_module_for_gemma4_unified = partial(
+        _patch_rms_norm_module, offset=0.0, casting_mode="gemma", in_place=False
+    )
+
+    def _maybe_patch_scaled_norm(module):
+        """Patch only Gemma4UnifiedRMSNorm modules that carry a weight.
+
+        Attention's ``v_norm`` and the multimodal embedder norms are
+        instantiated with ``with_scale=False`` — no weight exists, so Liger's
+        weight-multiplying kernel cannot apply. We leave these as HF's
+        scale-free RMSNorm (kernelized copies already swapped at the class
+        level via LigerRMSNormForGemma4 which also handles with_scale=False
+        correctly in its forward).
+        """
+        if module is None:
+            return
+        if not getattr(module, "with_scale", True):
+            return
+        _patch_rms_norm_module_for_gemma4_unified(module)
+
+    if rope:
+        # HF's Gemma 4 Unified apply_rotary_pos_emb has signature
+        #   apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=2)
+        # (single tensor at a time) whereas liger_rotary_pos_emb takes
+        # (q, k, cos, sin, ...). Until a Gemma-4-specific rope wrapper exists,
+        # leave HF's plain pytorch rope in place. Emit a single warning so
+        # callers flipping rope on aren't silently ignored.
+        logger.warning_once(
+            "rope=True is currently a no-op for Gemma 4 Unified: HF's "
+            "apply_rotary_pos_emb uses a single-tensor signature that is "
+            "incompatible with liger_rotary_pos_emb. Skipping rope kernel swap."
+        )
+
+    if rms_norm:
+        modeling_gemma4_unified.Gemma4UnifiedRMSNorm = LigerRMSNormForGemma4
+
+    if geglu:
+        # Gemma4UnifiedTextMLP is constructed with (config, layer_idx); the
+        # wrapper subclass accepts and discards layer_idx so the class-level
+        # swap doesn't crash model construction.
+        modeling_gemma4_unified.Gemma4UnifiedTextMLP = LigerGEGLUMLPForGemma4
+
+    # Handle loss function
+    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 None:
+            modeling_gemma4_unified.Gemma4UnifiedForCausalLM.forward = causal_forward
+        elif isinstance(model, Gemma4UnifiedForCausalLM):
+            model.forward = MethodType(causal_forward, model)
+        # A bare Gemma4UnifiedTextModel has no lm_head / loss path, so the
+        # causal-LM forward cannot be bound to it; the norm/MLP instance
+        # patching below still applies.
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+        if isinstance(model, (Gemma4UnifiedForCausalLM, Gemma4UnifiedTextModel)):
+            # get the base model from the model instance
+            base_model = model.model if isinstance(model, Gemma4UnifiedForCausalLM) else model
+
+            if rms_norm:
+                _maybe_patch_scaled_norm(base_model.norm)
+
+            for decoder_layer in base_model.layers:
+                decoder_layer: Gemma4UnifiedTextDecoderLayer
+                if geglu:
+                    _bind_method_to_module(decoder_layer.mlp, "forward", LigerGEGLUMLP.forward)
+                if rms_norm:
+                    _maybe_patch_scaled_norm(decoder_layer.input_layernorm)
+                    _maybe_patch_scaled_norm(decoder_layer.post_attention_layernorm)
+                    _maybe_patch_scaled_norm(decoder_layer.pre_feedforward_layernorm)
+                    _maybe_patch_scaled_norm(decoder_layer.post_feedforward_layernorm)
+                    # k_norm / v_norm exist only on non-KV-shared layers, so stay
+                    # defensive with getattr. v_norm is scale-free
+                    # (with_scale=False) on all layers so the helper
+                    # intentionally leaves it untouched.
+                    _maybe_patch_scaled_norm(getattr(decoder_layer.self_attn, "q_norm", None))
+                    _maybe_patch_scaled_norm(getattr(decoder_layer.self_attn, "k_norm", None))
+                    _maybe_patch_scaled_norm(getattr(decoder_layer.self_attn, "v_norm", None))
+        else:
+            raise TypeError("The model must be Gemma4UnifiedForCausalLM or Gemma4UnifiedTextModel.")
+
+
 def apply_liger_kernel_to_paligemma(
     rope: bool = True,
     cross_entropy: bool = False,
@@ -3539,6 +3675,7 @@ def __init__(self, hidden_size, eps=1e-6, **kwargs):
     "gemma3": apply_liger_kernel_to_gemma3,
     "gemma4_text": apply_liger_kernel_to_gemma4_text,
     "gemma4": apply_liger_kernel_to_gemma4,
+    "gemma4_unified_text": apply_liger_kernel_to_gemma4_unified_text,
     "glm4": apply_liger_kernel_to_glm4,
     "glm4v": apply_liger_kernel_to_glm4v,
     "glm4v_moe": apply_liger_kernel_to_glm4v_moe,

src/liger_kernel/transformers/rms_norm.py

@@ -84,6 +84,9 @@ class LigerRMSNormForGemma4(LigerRMSNorm):
 
     When ``with_scale=False`` the Liger kernel has no weight to multiply by,
     so we fall back to a plain torch implementation that matches HF exactly.
+
+    Also swapped in for ``Gemma4UnifiedRMSNorm`` (gemma4_unified), which is
+    implementation-identical to ``Gemma4RMSNorm``.
     """
 
     def __init__(