From a345a4f0d096ee1394ed1010c58cfe4bef61b515 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Thu, 19 Jun 2025 04:44:24 +0900
Subject: [PATCH] monkey patch liger kernel for gemma-3-1b-it

---
 src/ctx_to_lora/__init__.py     |   1 +
 src/ctx_to_lora/monkey_patch.py | 117 ++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 src/ctx_to_lora/monkey_patch.py

diff --git a/src/ctx_to_lora/__init__.py b/src/ctx_to_lora/__init__.py
index e69de29..115684e 100644
--- a/src/ctx_to_lora/__init__.py
+++ b/src/ctx_to_lora/__init__.py
@@ -0,0 +1 @@
+from . import monkey_patch as monkey_patch
diff --git a/src/ctx_to_lora/monkey_patch.py b/src/ctx_to_lora/monkey_patch.py
new file mode 100644
index 0000000..5040ff6
--- /dev/null
+++ b/src/ctx_to_lora/monkey_patch.py
@@ -0,0 +1,117 @@
+from functools import partial
+
+from liger_kernel.transformers import monkey_patch
+from liger_kernel.transformers.functional import liger_cross_entropy
+from liger_kernel.transformers.geglu import LigerGEGLUMLP
+from liger_kernel.transformers.monkey_patch import (
+    _bind_method_to_module,
+    _patch_rms_norm_module,
+)
+from liger_kernel.transformers.rope import liger_rotary_pos_emb
+from transformers import PreTrainedModel
+
+
+def apply_liger_kernel_to_gemma3_text_patched(
+    rope: bool = True,
+    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 Gemma3
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is True.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        geglu (bool): Whether to apply Liger's GeGLU MLP. Default is True.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+        loaded. Default is None.
+    """
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    ###
+    from liger_kernel.transformers.gema3_rms import LigerRMSNormForGemma3
+    from liger_kernel.transformers.model.gemma3 import causal_forward
+    from transformers.models.gemma3 import modeling_gemma3
+
+    ### PATCHED
+    from transformers.models.gemma3.modeling_gemma3 import (
+        Gemma3DecoderLayer,
+        Gemma3ForCausalLM,
+        Gemma3TextModel,
+    )
+
+    _patch_rms_norm_module_for_gemma3 = partial(
+        _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
+    )
+
+    if rope:
+        modeling_gemma3.apply_rotary_pos_emb = liger_rotary_pos_emb
+
+    if rms_norm:
+        modeling_gemma3.Gemma3RMSNorm = LigerRMSNormForGemma3
+
+    if geglu:
+        modeling_gemma3.Gemma3MLP = LigerGEGLUMLP
+
+    # 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:
+        modeling_gemma3.Gemma3ForCausalLM.forward = causal_forward
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+
+        ### PATCHED
+        if isinstance(model, Gemma3ForCausalLM) or isinstance(model, Gemma3TextModel):
+            # get the base model from the model instance
+            base_model = model.model if isinstance(model, Gemma3ForCausalLM) else model
+            ###
+
+            if rms_norm:
+                _patch_rms_norm_module_for_gemma3(base_model.norm)
+
+            for decoder_layer in base_model.layers:
+                decoder_layer: Gemma3DecoderLayer
+                if geglu:
+                    _bind_method_to_module(
+                        decoder_layer.mlp, "forward", LigerGEGLUMLP.forward
+                    )
+                if rms_norm:
+                    _patch_rms_norm_module_for_gemma3(decoder_layer.input_layernorm)
+                    _patch_rms_norm_module_for_gemma3(
+                        decoder_layer.post_attention_layernorm
+                    )
+                    _patch_rms_norm_module_for_gemma3(
+                        decoder_layer.pre_feedforward_layernorm
+                    )
+                    _patch_rms_norm_module_for_gemma3(
+                        decoder_layer.post_feedforward_layernorm
+                    )
+                    _patch_rms_norm_module_for_gemma3(decoder_layer.self_attn.q_norm)
+                    _patch_rms_norm_module_for_gemma3(decoder_layer.self_attn.k_norm)
+
+        else:
+            raise TypeError("The model must be Gemma3ForCausalLM.")
+
+
+monkey_patch.apply_liger_kernel_to_gemma3_text = (
+    apply_liger_kernel_to_gemma3_text_patched
+)
+monkey_patch.MODEL_TYPE_TO_APPLY_LIGER_FN["gemma3_text"] = (
+    apply_liger_kernel_to_gemma3_text_patched
+)