diff --git a/README.md b/README.md index fecc68518..2844fc2ec 100644 --- a/README.md +++ b/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 | diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 26bdef91b..c859a7447 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/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", 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..5d85385c4 --- /dev/null +++ b/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, + ) 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..0ef6f718b 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/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, 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 c78a97477..5f876005f 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, + 7e-1, # logprobs_atol — matches mini_gemma4_text (same bf16 near-tie flips; #1313 recalibrated it to 7e-1) + 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_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py index fa7a53ec0..e42bd16fe 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, + 7e-1, # logprobs_atol — matches mini_gemma4_text (same bf16 near-tie flips; #1313 recalibrated it to 7e-1) + 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/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 25099f7be..4d9d8f5df 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,7 @@ 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_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 +2122,72 @@ 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}") + + 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..8df437bce 100644 --- a/test/utils.py +++ b/test/utils.py @@ -520,6 +520,21 @@ 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_gemma3(model_config: MiniModelConfig): """ Revert all Liger kernel patches applied to Gemma3.