diff --git a/README.md b/README.md index fecc68518..5a56db567 100644 --- a/README.md +++ b/README.md @@ -308,6 +308,8 @@ 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 | +| Gemma4 Unified (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_gemma4_unified` | 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 26bdef91b..fe6871325 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -48,6 +48,8 @@ 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 # 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 +127,8 @@ 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", + "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 +220,8 @@ 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", + "apply_liger_kernel_to_gemma4_unified_text", "apply_liger_kernel_to_glm4", "apply_liger_kernel_to_glm4v", "apply_liger_kernel_to_glm4v_moe", diff --git a/src/liger_kernel/transformers/geglu.py b/src/liger_kernel/transformers/geglu.py index 32d4ff528..34ddae13c 100644 --- a/src/liger_kernel/transformers/geglu.py +++ b/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 """ diff --git a/src/liger_kernel/transformers/model/gemma4_unified.py b/src/liger_kernel/transformers/model/gemma4_unified.py new file mode 100644 index 000000000..9925cec8e --- /dev/null +++ b/src/liger_kernel/transformers/model/gemma4_unified.py @@ -0,0 +1,293 @@ +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 — both forwards are then + # unreachable because monkey_patch.apply_liger_kernel_to_gemma4_unified* + # 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?" + ```""" + + # Unlike gemma4 (omni), upstream Gemma4UnifiedForCausalLM.forward carries no + # eager-attention training recommendation, so no warning is mirrored here. + 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, + ) + + +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 ``Gemma4UnifiedForConditionalGeneration``. + + Mirrors :func:`liger_kernel.transformers.model.gemma4.multimodal_forward`. + Gemma 4 Unified shares one output class between the causal and multimodal + models, so this forward passes ``shared_kv_states`` through alongside the + image/audio hidden states. + + The win at long context is large: vocab=262,144 means the (B, T, V) 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. + 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 LigerGemma4UnifiedCausalLMOutputWithPast( + 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), + shared_kv_states=getattr(outputs, "shared_kv_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 c2ef4a7c7..c8c0dad8c 100644 --- a/src/liger_kernel/transformers/model/output_classes.py +++ b/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 diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 4d33d7e41..f96ef7c3a 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -1484,6 +1484,239 @@ 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_gemma4_unified( + 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 + Unified multimodal models (`Gemma4UnifiedForConditionalGeneration`). + + This is the class google/gemma-4-12B-it loads as (model_type + "gemma4_unified"), so text-only fine-tuning of the unified checkpoints also + lands here. For the standalone text classes (`Gemma4UnifiedForCausalLM`, + `Gemma4UnifiedTextModel`), use + :func:`apply_liger_kernel_to_gemma4_unified_text` instead. + + The primary win is the fused-linear-cross-entropy path: with vocab=262,144, + the (B, T, V) logits tensor is ~32 GiB in bf16 at T=65,536 (and ~64 GiB + once the loss path upcasts to fp32). Fused CE materializes only the loss + scalar. + + Out of scope (deferred to future PRs): + - LayerNorm kernels for the vision embedder. Unlike gemma4 (omni), there + are no AutoModel vision/audio towers here — images and audio are embedded + by the encoder-free `Gemma4UnifiedVisionEmbedder` / + `Gemma4UnifiedMultimodalEmbedder` (LayerNorm + Dense + scale-free + RMSNorm). The class-level RMSNorm swap covers the embedder RMSNorms for + fresh construction; `LigerRMSNormForGemma4` falls back to the exact torch + implementation for `with_scale=False` modules, preserving semantics. + + 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. + layer_norm (bool): Accepted for API compatibility; no LayerNorm + kernels are applied (the vision embedder is out of scope). + Default False. + 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_unified import modeling_gemma4_unified + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration + + from liger_kernel.transformers.model.gemma4_unified import multimodal_forward + + if model is not None and not isinstance(model, Gemma4UnifiedForConditionalGeneration): + raise TypeError("The model must be Gemma4UnifiedForConditionalGeneration.") + + # 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_unified_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_unified.Gemma4UnifiedForConditionalGeneration.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_unified_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, @@ -3539,6 +3772,8 @@ 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, + "gemma4_unified": apply_liger_kernel_to_gemma4_unified, "glm4": apply_liger_kernel_to_glm4, "glm4v": apply_liger_kernel_to_glm4v, "glm4v_moe": apply_liger_kernel_to_glm4v_moe, diff --git a/src/liger_kernel/transformers/rms_norm.py b/src/liger_kernel/transformers/rms_norm.py index 03a01a574..3cd893781 100644 --- a/src/liger_kernel/transformers/rms_norm.py +++ b/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__( diff --git a/test/convergence/bf16/test_mini_models.py b/test/convergence/bf16/test_mini_models.py index 88dabcc5a..5e78041c9 100644 --- a/test/convergence/bf16/test_mini_models.py +++ b/test/convergence/bf16/test_mini_models.py @@ -31,6 +31,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_gemma2 from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text from liger_kernel.transformers import apply_liger_kernel_to_gemma4_text +from liger_kernel.transformers import apply_liger_kernel_to_gemma4_unified_text from liger_kernel.transformers import apply_liger_kernel_to_glm4 from liger_kernel.transformers import apply_liger_kernel_to_glm4v from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe @@ -75,6 +76,7 @@ from test.utils import revert_liger_kernel_to_gemma2 from test.utils import revert_liger_kernel_to_gemma3_text from test.utils import revert_liger_kernel_to_gemma4_text +from test.utils import revert_liger_kernel_to_gemma4_unified_text from test.utils import revert_liger_kernel_to_glm4 from test.utils import revert_liger_kernel_to_glm4v from test.utils import revert_liger_kernel_to_glm4v_moe @@ -373,6 +375,15 @@ except ImportError: GEMMA4_AVAILABLE = False +try: + # Gemma4 Unified is only available in transformers>=5.10.0 + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedTextConfig + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForCausalLM + + GEMMA4_UNIFIED_AVAILABLE = True +except ImportError: + GEMMA4_UNIFIED_AVAILABLE = False + device = infer_device() @@ -816,6 +827,55 @@ ), ) +if GEMMA4_UNIFIED_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4_unified_text"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_gemma4_unified_text, + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4_unified_text, + model_class=Gemma4UnifiedForCausalLM, + mini_model_config=Gemma4UnifiedTextConfig( + # Shrunk from Gemma 4 12B (num_hidden_layers=48, hidden_size=3840, + # vocab_size=262144). Layer types mirror the 12B pattern + # (5 sliding, 1 full, repeat). + vocab_size=32000, + hidden_size=1024, + intermediate_size=2048, + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=256, + # Mini-sized to match head_dim; the 12B ships 512 for full-attention layers. + global_head_dim=256, + hidden_activation="gelu_pytorch_tanh", + max_position_embeddings=8192, + initializer_range=0.02, + rms_norm_eps=1e-06, + use_cache=True, + pad_token_id=0, + bos_token_id=2, + eos_token_id=1, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + attn_implementation="eager", + # 12B ships final_logit_softcapping=30.0 — exercises the FLCE softcap path. + final_logit_softcapping=30.0, + sliding_window=1024, + # Match 12B: every Nth layer is full_attention. + layer_types=[ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + # Defaults on 12B, pinned explicitly. Unlike gemma4 (omni) there + # are no PLE / MoE fields on the unified text config. + num_kv_shared_layers=0, + use_double_wide_mlp=False, + ), + ) + if MLLAMA_AVAILABLE: MINI_MODEL_SETUPS["mini_mllama"] = MiniModelConfig( @@ -2353,6 +2413,25 @@ def run_mini_model( ), ], ), + pytest.param( + "mini_gemma4_unified_text", + 32, + 1e-5, + torch.bfloat16, + 5e-2, # loss_atol — same 6-layer bf16 drift budget as mini_gemma4_text + 1e-2, + 5e-1, # logprobs_atol — same bf16 near-tie flips as mini_gemma4_text + 1e-2, + 1e-2, + 1e-2, + marks=[ + pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + pytest.mark.skipif( + not GEMMA4_UNIFIED_AVAILABLE, + reason="Gemma4 Unified not available in this version of transformers", + ), + ], + ), pytest.param( "mini_falcon_h1", 32, diff --git a/test/convergence/bf16/test_mini_models_multimodal.py b/test/convergence/bf16/test_mini_models_multimodal.py index 265b8261e..b29b1f62b 100644 --- a/test/convergence/bf16/test_mini_models_multimodal.py +++ b/test/convergence/bf16/test_mini_models_multimodal.py @@ -14,6 +14,7 @@ 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_gemma4_unified 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 @@ -42,6 +43,7 @@ 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_gemma4_unified 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 @@ -227,6 +229,21 @@ except ImportError: GEMMA4_AVAILABLE = False +try: + # Gemma4 Unified multimodal requires transformers>=5.10.0 + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedConfig + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedTextConfig + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedVisionConfig + from transformers.models.gemma4_unified.feature_extraction_gemma4_unified import Gemma4UnifiedAudioFeatureExtractor + from transformers.models.gemma4_unified.image_processing_gemma4_unified import Gemma4UnifiedImageProcessor + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration + from transformers.models.gemma4_unified.processing_gemma4_unified import Gemma4UnifiedProcessor + from transformers.models.gemma4_unified.video_processing_gemma4_unified import Gemma4UnifiedVideoProcessor + + GEMMA4_UNIFIED_AVAILABLE = True +except ImportError: + GEMMA4_UNIFIED_AVAILABLE = False + try: from transformers.models.llama4.configuration_llama4 import Llama4Config from transformers.models.llama4.configuration_llama4 import Llama4TextConfig @@ -592,6 +609,54 @@ ), ) +if GEMMA4_UNIFIED_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4_unified"] = MiniModelConfig( + liger_kernel_patch_func=functools.partial( + apply_liger_kernel_to_gemma4_unified, fused_linear_cross_entropy=False + ), + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4_unified, + model_class=Gemma4UnifiedForConditionalGeneration, + mini_model_config=Gemma4UnifiedConfig( + text_config=Gemma4UnifiedTextConfig( + vocab_size=32000, # 262144 + hidden_size=1024, # 3840 + intermediate_size=2048, # 15360 + num_hidden_layers=4, # 48 + num_attention_heads=4, # 16 + num_key_value_heads=2, # 8 + head_dim=128, # 256 + global_head_dim=128, # 512 — mini-sized to match head_dim + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + # Defaults on 12B. Unlike gemma4 (omni) there are no PLE / MoE + # fields on the unified text config. + num_kv_shared_layers=0, + use_double_wide_mlp=False, + ), + # Encoder-free vision embedder (no AutoModel tower). patch_size and + # pooling_kernel_size are pinned by Gemma4UnifiedImageProcessor + # defaults, which emit raw patches of dim + # (patch_size * pooling_kernel_size)^2 * 3 = 6912. + vision_config=Gemma4UnifiedVisionConfig( + patch_size=16, + pooling_kernel_size=3, + mm_embed_dim=256, # 3840 + output_proj_dims=256, # must equal mm_embed_dim + ), + audio_config=None, # Audio out of scope (mirrors mini_gemma4) + 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), @@ -1294,6 +1359,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_unified"): + # Must precede the broader "mini_gemma4" prefix check below. + tokenizer_config = load_tokenizer_config( + os.path.join( + FAKE_CONFIGS_PATH, + "Google/Gemma4Unified/gemma-4-12B-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 + # Gemma4UnifiedProcessor.__call__ are not exercised. + return Gemma4UnifiedProcessor( + feature_extractor=Gemma4UnifiedAudioFeatureExtractor(), + image_processor=Gemma4UnifiedImageProcessor(), + tokenizer=fast_tokenizer, + video_processor=Gemma4UnifiedVideoProcessor(), + ) + elif model_name.startswith("mini_gemma4"): tokenizer_config = load_tokenizer_config( os.path.join( @@ -1772,6 +1865,25 @@ def run_mini_model_multimodal( ), ], ), + pytest.param( + "mini_gemma4_unified", + 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_UNIFIED_AVAILABLE, + reason="Gemma4 Unified not available in this version of transformers", + ), + ], + ), pytest.param( "mini_qwen3_5", 32, diff --git a/test/convergence/bf16/test_mini_models_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py index 8c530dba4..aaa485c01 100644 --- a/test/convergence/bf16/test_mini_models_with_logits.py +++ b/test/convergence/bf16/test_mini_models_with_logits.py @@ -31,6 +31,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_gemma2 from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text from liger_kernel.transformers import apply_liger_kernel_to_gemma4_text +from liger_kernel.transformers import apply_liger_kernel_to_gemma4_unified_text from liger_kernel.transformers import apply_liger_kernel_to_glm4 from liger_kernel.transformers import apply_liger_kernel_to_glm4v from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe @@ -73,6 +74,7 @@ from test.utils import revert_liger_kernel_to_gemma2 from test.utils import revert_liger_kernel_to_gemma3_text from test.utils import revert_liger_kernel_to_gemma4_text +from test.utils import revert_liger_kernel_to_gemma4_unified_text from test.utils import revert_liger_kernel_to_glm4 from test.utils import revert_liger_kernel_to_glm4v from test.utils import revert_liger_kernel_to_glm4v_moe @@ -266,6 +268,15 @@ except ImportError: GEMMA4_AVAILABLE = False +try: + # Gemma4 Unified is only available in transformers>=5.10.0 + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedTextConfig + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForCausalLM + + GEMMA4_UNIFIED_AVAILABLE = True +except ImportError: + GEMMA4_UNIFIED_AVAILABLE = False + try: # Smollm3 is only available in transformers>=4.53.0 from transformers.models.smollm3.configuration_smollm3 import SmolLM3Config @@ -872,6 +883,55 @@ ), ) +if GEMMA4_UNIFIED_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4_unified_text"] = MiniModelConfig( + liger_kernel_patch_func=apply_liger_kernel_to_gemma4_unified_text, + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4_unified_text, + model_class=Gemma4UnifiedForCausalLM, + mini_model_config=Gemma4UnifiedTextConfig( + # Shrunk from Gemma 4 12B (num_hidden_layers=48, hidden_size=3840, + # vocab_size=262144). Layer types mirror the 12B pattern + # (5 sliding, 1 full, repeat). + vocab_size=32000, + hidden_size=1024, + intermediate_size=2048, + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=256, + # Mini-sized to match head_dim; the 12B ships 512 for full-attention layers. + global_head_dim=256, + hidden_activation="gelu_pytorch_tanh", + max_position_embeddings=8192, + initializer_range=0.02, + rms_norm_eps=1e-06, + use_cache=True, + pad_token_id=0, + bos_token_id=2, + eos_token_id=1, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + attn_implementation="eager", + # 12B ships final_logit_softcapping=30.0 — exercises the softcap path. + final_logit_softcapping=30.0, + sliding_window=1024, + # Match 12B: every Nth layer is full_attention. + layer_types=[ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + # Defaults on 12B, pinned explicitly. Unlike gemma4 (omni) there + # are no PLE / MoE fields on the unified text config. + num_kv_shared_layers=0, + use_double_wide_mlp=False, + ), + ) + if MLLAMA_AVAILABLE: MINI_MODEL_SETUPS["mini_mllama"] = MiniModelConfig( liger_kernel_patch_func=apply_liger_kernel_to_mllama, @@ -2184,6 +2244,25 @@ def run_mini_model( ), ], ), + pytest.param( + "mini_gemma4_unified_text", + 32, + 1e-5, + torch.bfloat16, + 5e-2, # loss_atol — same 6-layer bf16 drift budget as mini_gemma4_text + 5e-2, + 5e-1, # logprobs_atol — same bf16 near-tie flips as mini_gemma4_text + 1e-2, + 1e-2, + 1e-2, + marks=[ + pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + pytest.mark.skipif( + not GEMMA4_UNIFIED_AVAILABLE, + reason="Gemma4 Unified not available in this version of transformers", + ), + ], + ), pytest.param( "mini_smollm3", 32, diff --git a/test/convergence/fp32/test_mini_models_multimodal.py b/test/convergence/fp32/test_mini_models_multimodal.py index 9fd8d207e..b8c6b9822 100644 --- a/test/convergence/fp32/test_mini_models_multimodal.py +++ b/test/convergence/fp32/test_mini_models_multimodal.py @@ -15,6 +15,7 @@ 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_gemma4_unified 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 @@ -43,6 +44,7 @@ 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_gemma4_unified 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 @@ -268,6 +270,21 @@ except ImportError: GEMMA4_AVAILABLE = False +try: + # Gemma4 Unified multimodal requires transformers>=5.10.0 + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedConfig + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedTextConfig + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedVisionConfig + from transformers.models.gemma4_unified.feature_extraction_gemma4_unified import Gemma4UnifiedAudioFeatureExtractor + from transformers.models.gemma4_unified.image_processing_gemma4_unified import Gemma4UnifiedImageProcessor + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration + from transformers.models.gemma4_unified.processing_gemma4_unified import Gemma4UnifiedProcessor + from transformers.models.gemma4_unified.video_processing_gemma4_unified import Gemma4UnifiedVideoProcessor + + GEMMA4_UNIFIED_AVAILABLE = True +except ImportError: + GEMMA4_UNIFIED_AVAILABLE = False + try: # InternVL is only available in transformers>=4.52.1 from transformers.models.got_ocr2.image_processing_got_ocr2_fast import GotOcr2ImageProcessorFast @@ -621,6 +638,54 @@ ), ) +if GEMMA4_UNIFIED_AVAILABLE: + MINI_MODEL_SETUPS["mini_gemma4_unified"] = MiniModelConfig( + liger_kernel_patch_func=functools.partial( + apply_liger_kernel_to_gemma4_unified, fused_linear_cross_entropy=False + ), + liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4_unified, + model_class=Gemma4UnifiedForConditionalGeneration, + mini_model_config=Gemma4UnifiedConfig( + text_config=Gemma4UnifiedTextConfig( + vocab_size=32000, # 262144 + hidden_size=1024, # 3840 + intermediate_size=2048, # 15360 + num_hidden_layers=4, # 48 + num_attention_heads=4, # 16 + num_key_value_heads=2, # 8 + head_dim=128, # 256 + global_head_dim=128, # 512 — mini-sized to match head_dim + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=True, + attention_bias=False, + attention_dropout=0.0, + # Defaults on 12B. Unlike gemma4 (omni) there are no PLE / MoE + # fields on the unified text config. + num_kv_shared_layers=0, + use_double_wide_mlp=False, + ), + # Encoder-free vision embedder (no AutoModel tower). patch_size and + # pooling_kernel_size are pinned by Gemma4UnifiedImageProcessor + # defaults, which emit raw patches of dim + # (patch_size * pooling_kernel_size)^2 * 3 = 6912. + vision_config=Gemma4UnifiedVisionConfig( + patch_size=16, + pooling_kernel_size=3, + mm_embed_dim=256, # 3840 + output_proj_dims=256, # must equal mm_embed_dim + ), + audio_config=None, # Audio out of scope (mirrors mini_gemma4) + 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), @@ -1411,6 +1476,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_unified"): + # Must precede the broader "mini_gemma4" prefix check below. + tokenizer_config = load_tokenizer_config( + os.path.join( + FAKE_CONFIGS_PATH, + "Google/Gemma4Unified/gemma-4-12B-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 + # Gemma4UnifiedProcessor.__call__ are not exercised. + return Gemma4UnifiedProcessor( + feature_extractor=Gemma4UnifiedAudioFeatureExtractor(), + image_processor=Gemma4UnifiedImageProcessor(), + tokenizer=fast_tokenizer, + video_processor=Gemma4UnifiedVideoProcessor(), + ) + elif model_name.startswith("mini_gemma4"): tokenizer_config = load_tokenizer_config( os.path.join( @@ -1873,6 +1966,24 @@ def run_mini_model_multimodal( ), ], ), + pytest.param( + "mini_gemma4_unified", + 32, + 1e-5, + torch.float32, + 1e-8, + 1e-4, + 5e-3, + 1e-5, + 5e-3, + 1e-5, + marks=[ + pytest.mark.skipif( + not GEMMA4_UNIFIED_AVAILABLE, + reason="Gemma4 Unified not available in this version of transformers", + ), + ], + ), pytest.param( "mini_qwen3_5", 32, diff --git a/test/resources/fake_configs/Google/Gemma4Unified/gemma-4-12B-it/tokenizer_config.json b/test/resources/fake_configs/Google/Gemma4Unified/gemma-4-12B-it/tokenizer_config.json new file mode 100644 index 000000000..98ec9ae8e --- /dev/null +++ b/test/resources/fake_configs/Google/Gemma4Unified/gemma-4-12B-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": "Gemma4UnifiedProcessor", + "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 25099f7be..7160f8267 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -201,6 +201,16 @@ def is_gemma4_available(): return False +def is_gemma4_unified_available(): + # Requires transformers>=5.10.0. + try: + import transformers.models.gemma4_unified # noqa: F401 + + return True + except ImportError: + return False + + def is_paligemma_available(): try: import transformers.models.paligemma # noqa: F401 @@ -290,6 +300,8 @@ def test_import_from_root(): from liger_kernel.transformers import apply_liger_kernel_to_gemma3 # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_gemma4_text # noqa: F401 + from liger_kernel.transformers import apply_liger_kernel_to_gemma4_unified # noqa: F401 + from liger_kernel.transformers import apply_liger_kernel_to_gemma4_unified_text # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_glm4 # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_glm4v # noqa: F401 from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe # noqa: F401 @@ -2111,6 +2123,154 @@ def test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +@pytest.mark.skipif(not is_gemma4_unified_available(), reason="gemma4_unified module not available") +def test_apply_liger_kernel_to_instance_for_gemma4_unified_text(): + # Ensure any monkey patching is cleaned up for subsequent tests + with patch("transformers.models.gemma4_unified.modeling_gemma4_unified"): + from liger_kernel.transformers.model.gemma4_unified import causal_forward as gemma4_unified_causal_forward + + # Instantiate a dummy model. Unlike gemma4 (omni) there are no PLE or + # MoE knobs to pin off — the unified text stack is dense-only. + config = transformers.models.gemma4_unified.configuration_gemma4_unified.Gemma4UnifiedTextConfig( + 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, + ) + dummy_model_instance = AutoModelForCausalLM.from_config(config) + + # Pre-patch assertions + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(gemma4_unified_causal_forward) + assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward) + # q_norm / k_norm are only present on non-KV-shared layers; we pin + # num_kv_shared_layers=0 in the config above so every layer has them. + for layer in dummy_model_instance.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 kernels to the instance + _apply_liger_kernel_to_instance(model=dummy_model_instance) + + # Post-patch assertions + assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(gemma4_unified_causal_forward) + assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward) + for layer in dummy_model_instance.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 is scale-free (with_scale=False); _maybe_patch_scaled_norm + # intentionally skips it, so the instance must retain the HF forward. + v_norm = getattr(layer.self_attn, "v_norm", None) + if v_norm is not None: + 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}") + + +@pytest.mark.skipif(not is_gemma4_unified_available(), reason="gemma4_unified module not available") +def test_apply_liger_kernel_to_instance_for_gemma4_unified_conditional_generation(): + # Ensure any monkey patching is cleaned up for subsequent tests + with patch("transformers.models.gemma4_unified.modeling_gemma4_unified"): + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration + + from liger_kernel.transformers.model.gemma4_unified import ( + multimodal_forward as gemma4_unified_multimodal_forward, + ) + + # Minimal dense text config — same shape as the text-only test above. + text_config = transformers.models.gemma4_unified.configuration_gemma4_unified.Gemma4UnifiedTextConfig( + 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, + ) + # Vision/audio configs left as None — the unified model builds its + # encoder-free embedders inline behind `if config._config is not + # None`, so a None-towers model still constructs as + # Gemma4UnifiedForConditionalGeneration and exercises the multimodal + # forward we're patching. + config = transformers.models.gemma4_unified.configuration_gemma4_unified.Gemma4UnifiedConfig( + text_config=text_config, + vision_config=None, + audio_config=None, + ) + + dummy_model_instance = Gemma4UnifiedForConditionalGeneration._from_config(config) + assert isinstance(dummy_model_instance, Gemma4UnifiedForConditionalGeneration) + + # Pre-patch: forward and language-model norms must NOT be Liger. + assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(gemma4_unified_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_unified_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 52142ac71..7916826b1 100644 --- a/test/utils.py +++ b/test/utils.py @@ -520,6 +520,35 @@ def revert_liger_kernel_to_gemma4(model_config: MiniModelConfig): print("Liger kernel patches have been reverted.") +def revert_liger_kernel_to_gemma4_unified_text(model_config: MiniModelConfig): + """Revert all Liger kernel patches applied to Gemma4 Unified text model.""" + + from transformers.models.gemma4_unified import modeling_gemma4_unified + + # Only modeling_gemma4_unified needs reloading: the class-level swaps + # (Gemma4UnifiedRMSNorm, Gemma4UnifiedTextMLP) are reassignments on this + # module, and reloading resets them to the original HF classes. + importlib.reload(modeling_gemma4_unified) + + model_config.model_class = modeling_gemma4_unified.Gemma4UnifiedForCausalLM + + print("Liger kernel patches have been reverted.") + + +def revert_liger_kernel_to_gemma4_unified(model_config: MiniModelConfig): + """Revert all Liger kernel patches applied to Gemma4 Unified multimodal model.""" + + from transformers.models.gemma4_unified import modeling_gemma4_unified + + # Reloading modeling_gemma4_unified resets Gemma4UnifiedRMSNorm / + # Gemma4UnifiedTextMLP / Gemma4UnifiedForConditionalGeneration.forward, + # which is the surface the multimodal patch touches (the encoder-free + # vision/audio embedders live in the same module). + importlib.reload(modeling_gemma4_unified) + model_config.model_class = modeling_gemma4_unified.Gemma4UnifiedForConditionalGeneration + print("Liger kernel patches have been reverted.") + + def revert_liger_kernel_to_gemma3(model_config: MiniModelConfig): """ Revert all Liger kernel patches applied to Gemma3.