liger-kernel 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

liger_kernel/transformers/monkey_patch.py

@@ -1,9 +1,23 @@
+import inspect
+import logging
+from functools import partial
+
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
 from liger_kernel.transformers.geglu import LigerGEGLUMLP
-from liger_kernel.transformers.model.llama import lce_forward
+from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forward
+from liger_kernel.transformers.model.llama import lce_forward as llama_lce_forward
+from liger_kernel.transformers.model.mistral import lce_forward as mistral_lce_forward
+from liger_kernel.transformers.model.phi3 import lce_forward as phi3_lce_forward
+from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward
 from liger_kernel.transformers.rms_norm import LigerRMSNorm
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
-from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP, LigerSwiGLUMLP
+from liger_kernel.transformers.swiglu import (
+    LigerBlockSparseTop2MLP,
+    LigerPhi3SwiGLUMLP,
+    LigerSwiGLUMLP,
+)
+
+logger = logging.getLogger(__name__)
 
 
 def apply_liger_kernel_to_llama(
@@ -20,7 +34,7 @@ def apply_liger_kernel_to_llama(
         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 lienar cross entropy loss. Default is True.
+            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.
@@ -42,12 +56,13 @@ def apply_liger_kernel_to_llama(
     if cross_entropy:
         modeling_llama.CrossEntropyLoss = LigerCrossEntropyLoss
     if fused_linear_cross_entropy:
-        modeling_llama.LlamaForCausalLM.forward = lce_forward
+        modeling_llama.LlamaForCausalLM.forward = llama_lce_forward
 
 
 def apply_liger_kernel_to_mistral(
     rope: bool = True,
-    cross_entropy: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
     rms_norm: bool = True,
     swiglu: bool = True,
 ) -> None:
@@ -57,9 +72,17 @@ def apply_liger_kernel_to_mistral(
     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 True.
+        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.
         rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
         swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
     """
+    assert not (
+        cross_entropy and fused_linear_cross_entropy
+    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
 
     from transformers.models.mistral import modeling_mistral
 
@@ -69,6 +92,8 @@ def apply_liger_kernel_to_mistral(
         modeling_mistral.MistralRMSNorm = LigerRMSNorm
     if cross_entropy:
         modeling_mistral.CrossEntropyLoss = LigerCrossEntropyLoss
+    if fused_linear_cross_entropy:
+        modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
     if swiglu:
         modeling_mistral.MistralMLP = LigerSwiGLUMLP
 
@@ -94,7 +119,7 @@ def apply_liger_kernel_to_mixtral(
     if rope:
         modeling_mixtral.apply_rotary_pos_emb = liger_rotary_pos_emb
     if rms_norm:
-        modeling_mixtral.MistralRMSNorm = LigerRMSNorm
+        modeling_mixtral.MixtralRMSNorm = LigerRMSNorm
     if cross_entropy:
         modeling_mixtral.CrossEntropyLoss = LigerCrossEntropyLoss
     if swiglu:
@@ -104,12 +129,13 @@ def apply_liger_kernel_to_mixtral(
 def apply_liger_kernel_to_gemma(
     rope: bool = True,
     cross_entropy: bool = True,
+    fused_linear_cross_entropy: bool = True,
     rms_norm: bool = True,
     geglu: bool = True,
 ) -> None:
     """
-    Apply Liger kernels to replace original implementation in HuggingFace Gemma2 models
-    to make GPU go burrr.
+    Apply Liger kernels to replace original implementation in HuggingFace Gemma
+    (Gemma 1 and 1.1 supported, for Gemma2 please use `apply_liger_kernel_to_gemma2` ) to make GPU go burrr.
 
     Args:
         rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
@@ -117,14 +143,181 @@ def apply_liger_kernel_to_gemma(
         rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
         geglu (bool): Whether to apply Liger's GeGLU MLP. Default is True.
     """
-    # TODO(yundai424): add convergence test for gemma
+    assert not (
+        cross_entropy and fused_linear_cross_entropy
+    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+
     from transformers.models.gemma import modeling_gemma
 
     if rope:
         modeling_gemma.apply_rotary_pos_emb = liger_rotary_pos_emb
     if rms_norm:
-        modeling_gemma.GemmaRMSNorm = LigerRMSNorm
+        # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/gemma/modeling_gemma.py#L109
+        modeling_gemma.GemmaRMSNorm = partial(
+            LigerRMSNorm, offset=1.0, init_fn="zeros", casting_mode="gemma"
+        )
     if cross_entropy:
         modeling_gemma.CrossEntropyLoss = LigerCrossEntropyLoss
     if geglu:
         modeling_gemma.GemmaMLP = LigerGEGLUMLP
+    if fused_linear_cross_entropy:
+        modeling_gemma.GemmaForCausalLM.forward = gemma_lce_forward
+
+
+def apply_liger_kernel_to_gemma2(
+    rope: bool = True,
+    cross_entropy: bool = True,
+    rms_norm: bool = True,
+    geglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Gemma2
+    (for Gemma1 please use `apply_liger_kernel_to_gemma`) to make GPU go burrr.
+
+    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 True.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        geglu (bool): Whether to apply Liger's GeGLU MLP. Default is True.
+    """
+    from transformers.models.gemma2 import modeling_gemma2
+
+    if rope:
+        modeling_gemma2.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/gemma/modeling_gemma.py#L109
+        modeling_gemma2.Gemma2RMSNorm = partial(
+            LigerRMSNorm, offset=1.0, init_fn="zeros"
+        )
+    if cross_entropy:
+        modeling_gemma2.CrossEntropyLoss = LigerCrossEntropyLoss
+    if geglu:
+        modeling_gemma2.Gemma2MLP = LigerGEGLUMLP
+
+
+def apply_liger_kernel_to_qwen2(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Qwen2 models
+
+    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.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+    """
+    assert not (
+        cross_entropy and fused_linear_cross_entropy
+    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+
+    from transformers.models.qwen2 import modeling_qwen2
+
+    if rope:
+        modeling_qwen2.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        modeling_qwen2.Qwen2RMSNorm = LigerRMSNorm
+    if cross_entropy:
+        modeling_qwen2.CrossEntropyLoss = LigerCrossEntropyLoss
+    if fused_linear_cross_entropy:
+        modeling_qwen2.Qwen2ForCausalLM.forward = qwen2_lce_forward
+    if swiglu:
+        modeling_qwen2.Qwen2MLP = LigerSwiGLUMLP
+
+
+def apply_liger_kernel_to_phi3(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Phi3 models.
+
+    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.
+        swiglu (bool): Whether to apply Liger's SwiGLU Phi3MLP. Default is True.
+    """
+    assert not (
+        cross_entropy and fused_linear_cross_entropy
+    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+
+    from transformers.models.phi3 import modeling_phi3
+
+    if rope:
+        modeling_phi3.apply_rotary_pos_emb = liger_rotary_pos_emb  # Same as Gemma
+    if rms_norm:
+        modeling_phi3.Phi3RMSNorm = LigerRMSNorm  # Same as Llama
+    if swiglu:
+        modeling_phi3.Phi3MLP = LigerPhi3SwiGLUMLP
+    if cross_entropy:
+        modeling_phi3.CrossEntropyLoss = LigerCrossEntropyLoss
+    if fused_linear_cross_entropy:
+        modeling_phi3.Phi3ForCausalLM.forward = phi3_lce_forward
+
+
+# Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
+MODEL_TYPE_TO_APPLY_LIGER_FN = {
+    "gemma": apply_liger_kernel_to_gemma,
+    "gemma2": apply_liger_kernel_to_gemma2,
+    "llama": apply_liger_kernel_to_llama,
+    "mistral": apply_liger_kernel_to_mistral,
+    "mixtral": apply_liger_kernel_to_mixtral,
+    "qwen2": apply_liger_kernel_to_qwen2,
+    "phi3": apply_liger_kernel_to_phi3,
+}
+
+
+def _apply_liger_kernel(model_type: str = "", **kwargs) -> None:
+    """
+    Applies Liger kernels based on the specified model type. The custom
+    kernels for the specified model type will be applied with the provided
+    keyword arguments, otherwise the default configuration will be used.
+
+    Args:
+        - model_type: the model types as defined in transformers/models/auto/modeling_auto.py
+          and specified in the model's config.json
+        - kwargs: keyword arguments that are passed to the corresponding apply_liger_kernel_to_* function.
+    """
+
+    if not model_type:
+        logger.info("Model type was not provided. No Liger kernels will be applied.")
+        return
+
+    if model_type not in MODEL_TYPE_TO_APPLY_LIGER_FN.keys():
+        logger.info(
+            f"There are currently no Liger kernels supported for model type: {model_type}."
+        )
+        return
+
+    apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type]
+    apply_fn_signature = inspect.signature(apply_fn)
+
+    # Filter out the keyword arguments that are not supported by the apply function
+    applicable_kwargs = {
+        key: value
+        for key, value in kwargs.items()
+        if key in apply_fn_signature.parameters
+    }
+
+    logger.info(
+        f"Applying Liger kernels for model type: {model_type} with kwargs: {applicable_kwargs}"
+    )
+
+    # Apply the default combination of liger kernels available for the model
+    apply_fn(**applicable_kwargs)

liger_kernel/transformers/rms_norm.py

@@ -5,12 +5,28 @@ from liger_kernel.ops.rms_norm import LigerRMSNormFunction
 
 
 class LigerRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones"
+    ):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
+        assert init_fn in [
+            "ones",
+            "zeros",
+        ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
+        self.weight = nn.Parameter(
+            torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size)
+        )
+        self.variance_epsilon, self.offset, self.casting_mode = (
+            eps,
+            offset,
+            casting_mode,
+        )
 
     def forward(self, hidden_states):
         return LigerRMSNormFunction.apply(
-            hidden_states, self.weight, self.variance_epsilon
+            hidden_states,
+            self.weight,
+            self.variance_epsilon,
+            self.offset,
+            self.casting_mode,
         )

liger_kernel/transformers/swiglu.py

@@ -38,3 +38,27 @@ class LigerBlockSparseTop2MLP(nn.Module):
     def forward(self, x):
 
         return self.w2(LigerSiLUMulFunction.apply(self.w1(x), self.w3(x)))
+
+
+class LigerPhi3SwiGLUMLP(nn.Module):
+    """
+    Patch Phi3MLP to use LigerSiLUMulFunction
+    https://github.com/huggingface/transformers/blob/v4.41.0/src/transformers/models/phi3/modeling_phi3.py#L241
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_up_proj = nn.Linear(
+            self.hidden_size, 2 * self.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        if config.hidden_act not in ["silu", "swish"]:
+            raise ValueError(f"Activation function {config.hidden_act} not supported.")
+
+    def forward(self, x):
+        up_states = self.gate_up_proj(x)
+        gate, up_states = up_states.chunk(2, dim=-1)
+        return self.down_proj(LigerSiLUMulFunction.apply(gate, up_states))

liger_kernel/transformers/trainer_integration.py

@@ -1,45 +1,2 @@
-import logging
-
-from liger_kernel.transformers.monkey_patch import (
-    apply_liger_kernel_to_gemma,
-    apply_liger_kernel_to_llama,
-    apply_liger_kernel_to_mistral,
-    apply_liger_kernel_to_mixtral,
-)
-
-logger = logging.getLogger(__name__)
-
-# Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
-MODEL_TYPE_TO_APPLY_LIGER_FN = {
-    "gemma": apply_liger_kernel_to_gemma,
-    "llama": apply_liger_kernel_to_llama,
-    "mistral": apply_liger_kernel_to_mistral,
-    "mixtral": apply_liger_kernel_to_mixtral,
-}
-
-
-def _apply_liger_kernel(model_type: str = "", **kwargs) -> None:
-    """
-    Applies Liger kernels based on the specified model type. The custom
-    kernels for the specified model type will be applied with the provided
-    keyword arguments, otherwise the default configuration will be used.
-
-    Args:
-        - model_type: the model types as defined in transformers/models/auto/modeling_auto.py
-          and specified in the model's config.json
-        - kwargs: keyword arguments that are passed to the corresponding apply_liger_kernel_to_* function.
-    """
-
-    if not model_type:
-        logger.info("Model type was not provided. No Liger kernels will be applied.")
-        return
-
-    if model_type not in MODEL_TYPE_TO_APPLY_LIGER_FN.keys():
-        logger.info(
-            f"There are currently no Liger kernels supported for model type: {model_type}."
-        )
-        return
-
-    logger.info(f"Applying Liger kernels for model type: {model_type}.")
-    # Apply the default combination of liger kernels available for the model
-    MODEL_TYPE_TO_APPLY_LIGER_FN[model_type](**kwargs)
+# To not break HF Trainer integration
+from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401