liger-kernel 0.1.1__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

{liger_kernel-0.1.1.dist-info → liger_kernel-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger-kernel
-Version: 0.1.1
+Version: 0.2.0
 Summary: Efficient Triton kernels for LLM Training
 Home-page: https://github.com/linkedin/Liger-Kernel
 License: BSD-2-Clause
@@ -21,22 +21,35 @@ License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2
 Requires-Dist: triton>=2.3.0
-Requires-Dist: transformers>=4.40.1
+Requires-Dist: transformers>=4.42.0
 Provides-Extra: dev
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"
 Requires-Dist: flake8>=4.0.1.1; extra == "dev"
 Requires-Dist: black>=24.4.2; extra == "dev"
 Requires-Dist: isort>=5.13.2; extra == "dev"
-Requires-Dist: pre-commit>=3.7.1; extra == "dev"
-Requires-Dist: torch-tb-profiler>=0.4.1; extra == "dev"
+Requires-Dist: pytest>=7.1.2; extra == "dev"
+Requires-Dist: datasets>=2.19.2; extra == "dev"
 
 # Liger Kernel: Efficient Triton Kernels for LLM Training
 
+
+
 [![Downloads](https://static.pepy.tech/badge/liger-kernel)](https://pepy.tech/project/liger-kernel) [![PyPI version](https://badge.fury.io/py/liger-kernel.svg)](https://badge.fury.io/py/liger-kernel) [![PyPI version](https://badge.fury.io/py/liger-kernel-nightly.svg)](https://badge.fury.io/py/liger-kernel-nightly)
+[![](https://dcbadge.vercel.app/api/server/cudamode?style=flat)](https://discord.gg/CX2YmNmn)
 
+<img src="https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/logo-banner.png">
 
 [Installation](#installation) | [Getting Started](#getting-started) | [Examples](#examples) | [APIs](#apis) | [Structure](#structure) | [Contributing](#contributing)
 
+<details>
+  <summary>Latest News 🔥</summary>
+
+  - [2024/8/31] CUDA MODE talk, [Liger-Kernel: Real-world Triton kernel for LLM Training](https://discord.gg/6CNeDAjq?event=1273323969788772455)
+  - [2024/8/23] Official release: check out our [X post](https://x.com/hsu_byron/status/1827072737673982056)
+
+</details>
+
+
 **Liger (Linkedin GPU Efficient Runtime) Kernel** is a collection of Triton kernels designed specifically for LLM training. It can effectively increase multi-GPU **training throughput by 20%** and reduces **memory usage by 60%**. We have implemented **Hugging Face Compatible** `RMSNorm`, `RoPE`, `SwiGLU`, `CrossEntropy`, `FusedLinearCrossEntropy`, and more to come. The kernel works out of the box with [Flash Attention](https://github.com/Dao-AILab/flash-attention), [PyTorch FSDP](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html), and [Microsoft DeepSpeed](https://github.com/microsoft/DeepSpeed). We welcome contributions from the community to gather the best kernels for LLM training.
 
 ## Supercharge Your Model with Liger Kernel
@@ -52,8 +65,8 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
 | ![Speed up](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/e2e-tps.png) | ![Memory](https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/e2e-memory.png) |
 
 > **Note:**
-> - Benchmark conditions: LLaMA 3-8B, Batch Size = 8, Data Type = `bf16`, Optimizer = AdamW, Gradient Checkpointing = True, Distributed Strategy = FSDP1 on 8 A100s. 
-> - Hugging Face models start to OOM at a 4K context length, whereas Hugging Face + Liger Kernel scales up to 16K.  
+> - Benchmark conditions: LLaMA 3-8B, Batch Size = 8, Data Type = `bf16`, Optimizer = AdamW, Gradient Checkpointing = True, Distributed Strategy = FSDP1 on 8 A100s.
+> - Hugging Face models start to OOM at a 4K context length, whereas Hugging Face + Liger Kernel scales up to 16K.
 
 ## Examples
 
@@ -91,12 +104,15 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
 
 - `torch >= 2.1.2`
 - `triton >= 2.3.0`
-- `transformers >= 4.40.1`
+- `transformers >= 4.42.0`
+
+> **Note:**
+> Our kernels inherit the full spectrum of hardware compatibility offered by [Triton](https://github.com/triton-lang/triton).
 
 To install the stable version:
 
 ```bash
-$ pip install liger-kernel 
+$ pip install liger-kernel
 ```
 
 To install the nightly version:
@@ -105,9 +121,30 @@ To install the nightly version:
 $ pip install liger-kernel-nightly
 ```
 
+To install from source:
+
+```bash
+git clone https://github.com/linkedin/Liger-Kernel.git
+cd Liger-Kernel
+pip install -e .
+```
 ## Getting Started
 
-### 1. Patch Existing Hugging Face Models
+There are a couple ways to apply Liger kernels, depending on the level of customization required.
+
+### 1. Use AutoLigerKernelForCausalLM
+
+Using the `AutoLigerKernelForCausalLM` is the simplest approach, as you don't have to import a model-specific patching API. If the model type is supported, the modeling code will be automatically patched using the default settings.
+
+```python
+from liger_kernel.transformers import AutoLigerKernelForCausalLM
+
+# This AutoModel wrapper class automatically monkey-patches the
+# model with the optimized Liger kernels if the model is supported.
+model = AutoLigerKernelForCausalLM.from_pretrained("path/to/some/model")
+```
+
+### 2. Apply Model-Specific Patching APIs
 
 Using the [patching APIs](#patching), you can swap Hugging Face models with optimized Liger Kernels.
 
@@ -115,13 +152,22 @@ Using the [patching APIs](#patching), you can swap Hugging Face models with opti
 import transformers
 from liger_kernel.transformers import apply_liger_kernel_to_llama
 
-model = transformers.AutoModelForCausalLM.from_pretrained("<some llama model>")
+model = transformers.AutoModelForCausalLM("path/to/llama/model")
 
 # Adding this line automatically monkey-patches the model with the optimized Liger kernels
-apply_liger_kernel_to_llama() 
+apply_liger_kernel_to_llama()
+
+# You could alternatively specify exactly which kernels are applied
+apply_liger_kernel_to_llama(
+  rope=True,
+  swiglu=True,
+  cross_entropy=True,
+  fused_linear_cross_entropy=False,
+  rms_norm=False
+)
 ```
 
-### 2. Compose Your Own Model
+### 3. Compose Your Own Model
 
 You can take individual [kernels](#kernels) to compose your models.
 
@@ -161,14 +207,26 @@ loss.backward()
 
 ## APIs
 
+### AutoModel
+
+| **AutoModel Variant** | **API** |
+|-----------|---------|
+| AutoModelForCausalLM | `liger_kernel.transformers.AutoLigerKernelForCausalLM` |
+
+
 ### Patching
 
 | **Model**   | **API**                                                      | **Supported Operations**                                                |
 |-------------|--------------------------------------------------------------|-------------------------------------------------------------------------|
-| LLaMA (2 & 3) | `liger_kernel.transformers.apply_liger_kernel_to_llama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
-| Mistral     | `liger_kernel.transformers.apply_liger_kernel_to_mistral`  | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss        |
+| LLaMA 2 & 3 | `liger_kernel.transformers.apply_liger_kernel_to_llama`   | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
+| Mistral     | `liger_kernel.transformers.apply_liger_kernel_to_mistral`  | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | Mixtral     | `liger_kernel.transformers.apply_liger_kernel_to_mixtral`  | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss        |
-| Gemma2      | `liger_kernel.transformers.apply_liger_kernel_to_gemma`    | RoPE, RMSNorm, GeGLU, CrossEntropyLoss         |
+| Gemma1      | `liger_kernel.transformers.apply_liger_kernel_to_gemma`    | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
+| Gemma2      | `liger_kernel.transformers.apply_liger_kernel_to_gemma2`   | RoPE, RMSNorm, GeGLU, CrossEntropyLoss         |
+| Qwen2       | `liger_kernel.transformers.apply_liger_kernel_to_qwen2`    | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
+| Phi3        | `liger_kernel.transformers.apply_liger_kernel_to_phi3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
+
+
 
 ### Kernels
 
@@ -182,11 +240,11 @@ loss.backward()
 | FusedLinearCrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
 
 - **RMSNorm**: [RMSNorm](https://arxiv.org/pdf/1910.07467), which normalizes activations using their root mean square, is implemented by fusing the normalization and scaling steps into a single Triton kernel, and achieves ~3X speedup with ~3X peak memory reduction.
-- **RoPE**: [Rotary Positional Embedding](https://arxiv.org/pdf/2104.09864) is implemented by fusing the query and key embedding rotary into a single kernel with inplace replacement, and achieves ~3X speedup with ~3X peak memory reduction. 
-- **SwiGLU**: [Swish Gated Linear Units](https://arxiv.org/pdf/2002.05202), given by 
+- **RoPE**: [Rotary Positional Embedding](https://arxiv.org/pdf/2104.09864) is implemented by fusing the query and key embedding rotary into a single kernel with inplace replacement, and achieves ~3X speedup with ~3X peak memory reduction.
+- **SwiGLU**: [Swish Gated Linear Units](https://arxiv.org/pdf/2002.05202), given by
 $$\text{SwiGLU}(x)=\text{Swish}_{\beta}(xW+b)\otimes(xV+c)$$
 , is implemented by fusing the elementwise multiplication (denoted by $\otimes$) into a single kernel with inplace replacement, and achieves parity speed with ~1.5X peak memory reduction.
-- **GeGLU**: [GELU Gated Linear Units](https://arxiv.org/pdf/2002.05202), given by 
+- **GeGLU**: [GELU Gated Linear Units](https://arxiv.org/pdf/2002.05202), given by
 $$\text{GeGLU}(x)=\text{GELU}(xW+b)\otimes(xV+c)$$
 , is implemented by fusing the elementwise multiplication into a single kernel with inplace replacement, and achieves parity speed with ~1.5X peak memory reduction. Note that the [tanh approximation form of GELU](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html) is used.
 - **CrossEntropy**: [Cross entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) is implemented by computing both the loss and gradient in the forward pass with inplace replacement of input to reduce the peak memory by avoiding simultaneous materialization of both input logits and gradient. It achieves >2X speedup and >4X memory reduction for common vocab sizes (e.g., 32K, 128K, etc.).
@@ -195,12 +253,12 @@ $$\text{GeGLU}(x)=\text{GELU}(xW+b)\otimes(xV+c)$$
 
 
 <!-- TODO: be more specific about batch size -->
-> **Note:**  
+> **Note:**
 > Reported speedups and memory reductions are with respect to the LLaMA 3-8B Hugging Face layer implementations. All models use 4K hidden size and 4K sequence length and are evaluated based on memory usage and wall time for the forward+backward pass on a single NVIDIA A100 80G GPU using small batch sizes. Liger kernels exhibit more efficient scaling to larger batch sizes, detailed further in the [Benchmark](./benchmark) folder.
 
 ## Note on ML Compiler
 
-### 1. Torch Compile
+### Torch Compile
 
 Since Liger Kernel is 100% Triton-based, it works seamlessly with [`torch.compile`](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html). In the following example, Liger Kernel can further optimize the model on top of Torch Compile, reducing the memory by more than half.
 
@@ -209,20 +267,17 @@ Since Liger Kernel is 100% Triton-based, it works seamlessly with [`torch.compil
 | Torch Compile                  | 3780                       | 66.4                   |
 | Torch Compile + Liger Kernel   | 3702                       | 31.0                   |
 
-> **Note:**  
+> **Note:**
 > 1. Benchmark conditions: LLaMA 3-8B, Batch Size = 8, Seq Len = 4096, Data Type = `bf16`, Optimizer = AdamW, Gradient Checkpointing = True, Distributed Strategy = FSDP1 on 8 A100s.
 > 2. Tested on torch `2.5.0.dev20240731+cu118`
 
-### 2. Lightning Thunder
-
-*WIP*
-
 ## Contributing
 
 [CONTRIBUTING GUIDE](https://github.com/linkedin/Liger-Kernel/blob/main/CONTRIBUTING.md)
 
 ## Acknowledgement
 
+- [@claire_yishan](https://twitter.com/claire_yishan) for the LOGO design
 - [flash-attn](https://github.com/Dao-AILab/flash-attention) and [Unsloth](https://github.com/unslothai/unsloth) for inspiration in Triton kernels for training
 - [tiny shakespeare dataset](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) by Andrej Karpathy for convergence testing
 - [Efficient Cross Entropy](https://github.com/mgmalek/efficient_cross_entropy) for lm_head + cross entropy inspiration
@@ -232,6 +287,10 @@ Since Liger Kernel is 100% Triton-based, it works seamlessly with [`torch.compil
 
 [BSD 2-CLAUSE](https://github.com/linkedin/Liger-Kernel/blob/main/LICENSE)
 
+## Contact
+
+- For collaboration, please send email to byhsu@linkedin.com
+
 ## Cite this work
 
 Biblatex entry:
@@ -243,3 +302,6 @@ Biblatex entry:
   year   = {2024}
 }
 ```
+
+## Star History
+[![Star History Chart](https://api.star-history.com/svg?repos=linkedin/Liger-Kernel&type=Date)](https://star-history.com/#linkedin/Liger-Kernel&Date)

liger_kernel-0.2.0.dist-info/RECORD

@@ -0,0 +1,33 @@
+liger_kernel/env_report.py,sha256=LFUJ6UMkFFGPBYXBlqHFGy4bhsemEpSI-_1edSazlHI,1130
+liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liger_kernel/ops/cross_entropy.py,sha256=6-jI03Yw_B8gHLmqxTOMpbFMRJhNNaE4DKpmowWYrTE,9177
+liger_kernel/ops/fused_linear_cross_entropy.py,sha256=RCStJoBsgiAa03ZupWhZEHjnAbm52iNiMDsGs6VDtnY,8385
+liger_kernel/ops/geglu.py,sha256=f8j9qnZgqvXFDFStZ5WbtRPDuNo9FBdVdXX7ufMHNpE,4052
+liger_kernel/ops/rms_norm.py,sha256=B0FDElKiTygv1CdG3NzbeOeC7nj_-2vmNIg4RHistHI,9517
+liger_kernel/ops/rope.py,sha256=8TOkpjmeekQEp1x6OAXTAWwoTTcEhNHSk9GnjuhW-Cw,8570
+liger_kernel/ops/swiglu.py,sha256=MRbSIXsBLqlFr9ZdtuFqSjLJJ-716URmQIhxQ57GGEw,2915
+liger_kernel/ops/utils.py,sha256=vsFIywd8LQlVPRA3RPZOm5HyN8c0cS4NFEEnwjNw-MI,1427
+liger_kernel/transformers/__init__.py,sha256=Um9ZRvT289MVFoGmliSva3q3YLRDqwYmBLxIj0rD9nI,403
+liger_kernel/transformers/auto_model.py,sha256=WQyaORi2zPIWTLhuAWCRPIzyHd5T4my4yGHQrt1-uBA,1247
+liger_kernel/transformers/cross_entropy.py,sha256=G-L4EaUYVc25NKZ2jrlaG-d5YUvDqJdUlawPN7K1d1g,389
+liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=X6ni--b5F2GRxL46PrgjsvQuWEcKp3Z2cELdMRR0oyY,518
+liger_kernel/transformers/geglu.py,sha256=QcrME_8ooIn0xa59LaC0aoOdRrBIFd11Y0bAyF0NfCw,1130
+liger_kernel/transformers/monkey_patch.py,sha256=vbYmT7povRJQA49Ra7GGD3hMN2lVTXmsxO9R8lfQSoI,13022
+liger_kernel/transformers/rms_norm.py,sha256=YxBSn2bIfh24De8Xb7QhhmdG3taauj_qJNvEjlazonU,912
+liger_kernel/transformers/rope.py,sha256=m-ah8vZBYW8tfplTXCiAPMHJWlB1tdp_JPXJeWE-Boo,943
+liger_kernel/transformers/swiglu.py,sha256=0-tVJ8xEYfhxnduc16PflXFj8sZPxdx9sHUn3hfwCI4,2468
+liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
+liger_kernel/transformers/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liger_kernel/transformers/model/gemma.py,sha256=EcdkGbSj_qroTDFl0Sc_HLyDyY0xcDhwrgkM_wkXnw8,4987
+liger_kernel/transformers/model/llama.py,sha256=6McXLi_Bt35WuxaJ_0CzEnOtayHXiPw5vjiDsaQKdJU,5323
+liger_kernel/transformers/model/mistral.py,sha256=_MQJrDntlxBO5cJwgTjr2rk2nNd5FAXVnzcTg_PEekQ,5079
+liger_kernel/transformers/model/phi3.py,sha256=zmjOsVV5TjKJ0U2dCm6W-8WCx1toKoh2Wm2PZu3XOIw,4927
+liger_kernel/transformers/model/qwen2.py,sha256=Va4uiZaVzCG2V7XKDfHjZyYTre5vPQM02j83jnnhono,4873
+liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
+liger_kernel/triton/monkey_patch.py,sha256=5BcGKTtdqeYchypBIBopGIWPx1-cFALz7sOKoEsqXJ0,1584
+liger_kernel-0.2.0.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel-0.2.0.dist-info/METADATA,sha256=eAzgrCGOn_jpoiH8VomABxzt7CThKfBqOIKH5Qmfm3w,17049
+liger_kernel-0.2.0.dist-info/NOTICE,sha256=BXkXY9aWvEy_7MAB57zDu1z8uMYT1i1l9B6EpHuBa8s,173
+liger_kernel-0.2.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+liger_kernel-0.2.0.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel-0.2.0.dist-info/RECORD,,