liger-kernel 0.0.0__tar.gz → 0.1.0__tar.gz

liger_kernel-0.1.0/LICENSE

@@ -0,0 +1,23 @@
+BSD 2-CLAUSE LICENSE
+Copyright 2024 LinkedIn Corporation 
+All Rights Reserved.
+Redistribution and use in source and binary forms, with or
+without modification, are permitted provided that the following
+conditions are met:
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above
+copyright notice, this list of conditions and the following
+disclaimer in the documentation and/or other materials provided
+with the distribution.
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

liger_kernel-0.1.0/NOTICE

@@ -0,0 +1,4 @@
+Copyright 2024 LinkedIn Corporation
+All Rights Reserved.
+
+Licensed under the BSD 2-Clause License (the "License").  See License in the project root for license information.

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/PKG-INFO

@@ -1,4 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel
-Version: 0.0.0
+Version: 0.1.0
 Provides-Extra: dev
+License-File: LICENSE
+License-File: NOTICE

liger_kernel-0.1.0/README.md

@@ -0,0 +1,205 @@
+# 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)
+
+
+[Installation](#installation) | [Getting Started](#getting-started) | [Examples](#examples) | [APIs](#apis) | [Structure](#structure) | [Contributing](#contributing)
+
+**Liger (Linkedin GPU Efficient Runtime) Kernel** is a collection of Triton kernels designed specifically for LLM training. We have implemented **Hugging Face Compatible** `RMSNorm`, `RoPE`, `SwiGLU`, `CrossEntropy`, `FusedLinearCrossEntropy`, and more to come. It can effectively increase multi-GPU **training throughput by 20%** and reduces **memory usage by 60%**. The kernel works out of the box with [flash attention](https://github.com/Dao-AILab/flash-attention), PyTorch FSDP, and Microsoft DeepSpeed. We welcome contributions from the community to gather the best kernels for LLM training.
+
+
+
+## Supercharge Your Model with Liger Kernel
+
+![Banner](/docs/images/banner.GIF)
+
+Gain +20% throughput and reduce memory usage by 60%. Achieve longer context lengths and larger batch sizes. It’s also useful if you want to scale up your model to multi-head training or large vocabulary sizes.
+
+| Speed Up                 | Memory Reduction        |
+|--------------------------|-------------------------|
+| ![Speed up](docs/images/e2e-tps.png) | ![Memory](docs/images/e2e-memory.png) |
+
+
+> - 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 Liger Kernel scales up to 16K.  
+
+## Examples
+
+### Basic
+
+| **Example**                                    | **Description**                                                                                   | **Lightning Studio** |
+|------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
+| [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface)      | Train llama3 8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP | TBA                  |
+| [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning)         | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 | TBA                  |
+
+### Advanced
+
+| **Example**                                    | **Description**                                                                                   | **Lightning Studio** |
+|------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
+| [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa)        | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP      | TBA                  |
+
+## Key Features
+
+- **Ease of use:** Simply patch your Hugging Face model with one line of code, or compose your own model using our kernels.
+- **Time- and memory-efficient:** In the same spirit as Flash-Attn, but for layers like **RMSNorm**, **RoPE**, **CrossEntropy**! Increases multi-GPU training throughput by 20% and reduces memory usage by 60% with **kernel fusion**, **in-place replacement**, and **chunking** techniques.
+- **Exact:** Exact kernels—no approximations. Both forward and backward are implemented with rigorous unit and convergence testing to ensure accuracy.
+- **Lightweight:** The kernels have minimal dependencies, requiring only Torch and Triton—no extra libraries needed! Say goodbye to dependency headaches!
+- **Multi-GPU supported:** Compatible with multi-GPU setups (PyTorch FSDP and DeepSpeed).
+
+## Target Audiences
+
+- **Researchers**: Looking to compose models using efficient and reliable kernels for frontier experiments.
+- **ML Practitioners**: Focused on maximizing GPU training efficiency with optimal, high-performance kernels.
+- **Curious Novices**: Eager to learn how to write reliable Triton kernels to enhance training efficiency.
+
+
+## Installation
+
+### Dependencies
+
+- `torch >= 2.1.2`
+- `triton >= 2.3.0`
+- `transformers >= 4.40.1`
+
+To install the stable version:
+
+```bash
+$ pip install liger-kernel 
+```
+
+To install the nightly version:
+
+```bash
+$ pip install liger-kernel-nightly
+```
+
+## Getting Started
+
+### 1. Patch Existing Hugging Face Models
+
+Using [patching APIs](#patching), you can swap Hugging Face models with optimized Liger Kernels.
+
+```python
+import transformers
+from liger_kernel.transformers import apply_liger_kernel_to_llama
+
+model = transformers.AutoModelForCausalLM.from_pretrained("<some llama model>")
+
+# By adding this line, it automatically monkey patches the model with the optimized kernels
+apply_liger_kernel_to_llama() 
+```
+
+### 2. Compose Your Own Model
+
+You can take individual [kernels](#kernels) to compose your models.
+
+```python
+from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
+import torch.nn as nn
+import torch
+
+model = nn.Linear(128, 256).to("cuda")
+
+# LigerFusedLinearCrossEntropyLoss fuses linear and cross entropy layers together and performs chunk-by-chunk computation to reduce memory
+loss_fn = LigerFusedLinearCrossEntropyLoss()
+
+input = torch.randn(4, 128, requires_grad=True, device="cuda")
+target = torch.empty(4, dtype=torch.long, device="cuda").random_(256)
+
+loss = loss_fn(model.weight, input, target)
+loss.backward()
+```
+
+
+## Structure
+
+### Source Code
+
+- `ops/`: Core Triton operations.
+- `transformers/`: PyTorch `nn.Module` implementations built on Triton operations, compliant with the `transformers` API.
+
+### Tests
+
+- `transformers/`: Correctness tests for the Triton-based layers.
+- `convergence/`: Patches Hugging Face models with all kernels, runs multiple iterations, and compares weights, logits, and loss layer by layer.
+
+### Benchmark
+
+- `benchmark/`: Execution time and memory benchmarks compared to Hugging Face layers.
+
+## APIs
+
+### 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        |
+| 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         |
+
+### Kernels
+
+| **Kernel**                      | **API**                                                     |
+|---------------------------------|-------------------------------------------------------------|
+| RMSNorm                         | `liger_kernel.transformers.LigerRMSNorm`                    |
+| RoPE                            | `liger_kernel.transformers.liger_rotary_pos_emb`            |
+| SwiGLU                          | `liger_kernel.transformers.LigerSwiGLUMLP`                  |
+| GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
+| CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
+| FusedLinearCrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+
+- **RMSNorm**: RMSNorm, which normalizes tensor activations using their root mean square, is accelerated by fusing the normalization and scaling steps into a single triton kernel, achieved ~3X speedup with ~3X peak memory reduction. [RMSNorm Paper](https://arxiv.org/pdf/1910.07467)
+- **RoPE**: Fused the operations of query and key embedding rotary into a single kernel with inplace replacement, achieved ~3X speedup with ~3X peak memory reduction. [RoPE Paper](https://arxiv.org/pdf/2104.09864)
+- **SwiGLU**: Leveraging the fused triton kernel for the elementwise transformation in $$SwiGLU_{\beta=1}$$ ($$\sigma(A) \odot B$$) with inplace replacement, achieved parity speed with ~1.5X peak memory reduction. [SwiGLU Paper](https://arxiv.org/pdf/2002.05202)
+- **GeGLU**: Leveraging the fused triton kernel for the elementwise transformation in GeGLU with [tanh approximation form of GELU](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html) and inplace replacement, achieved parity speed with ~1.5X peak memory reduction. [GeGLU paper](https://arxiv.org/pdf/2002.05202)
+- **CrossEntropy**: Computes both loss and the gradient in the forward path with inplace replacement of input to reduce the peak memory (avoid the materialization of both input logits and gradient), achieved >2X speedup and >4X memory reduction for common vocab sizes. [PyTorch CrossEntropyLoss Documentation](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)
+- **FusedLinearCrossEntropy**: Further improves upon the basic Liger Cross Entropy kernel on reducing the peak memory usage by fusing the model last output head layer with the CE loss and chunking the input for block-wise loss and gradient calculation, inspired by [Efficient Cross Entropy](https://github.com/mgmalek/efficient_cross_entropy), achieved >4X memory reduction for 128k vocab size. **This is highly effective for large batch size, large sequence length, and large vocab size model** Please refer to the [Medusa example](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) for individual kernel usage.
+
+
+> * Reported speedups and memory reductions are compared with Llama3 8B Hugging Face layer implementations with 4k default hidden size and 4k sequence length for single forward and backward pass on single NVIDIA A100 80G GPU with small batch sizes. Liger kernels exhibits more efficient scaling to larger batch sizes of tokens. See [Benchmark](./benchmark) folder for details.
+
+## Note on ML Compiler
+
+### 1. Torch Compile
+
+Since Liger Kernel is 100% Triton-based, it works seamlessly with Torch Compile. In the following example, Liger Kernel can further optimize the model on top of Torch Compile, reducing the memory by more than half.
+
+| Configuration                  | Throughput (tokens/sec) | Memory Reserved (MB) |
+|--------------------------------|----------------------------|-------------------------|
+| Torch Compile                  | 3780                       | 66358                   |
+| Torch Compile + Liger Kernel   | 3702                       | 31000                   |
+
+> **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
+
+- [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) for convergence testing by andrej karpathy
+
+
+## License
+
+[BSD 2-CLAUSE](https://github.com/linkedin/Liger-Kernel/blob/main/LICENSE)
+
+## Cite this work
+
+Biblatex entry:
+```bib
+@software{liger2024,
+  title  = {Liger-Kernel: Efficient Triton Kernels for LLM Training},
+  author = {Hsu, Pin-Lun and Dai, Yun and Kothapalli, Vignesh and Song, Qingquan and Tang, Shao and Zhu, Siyu},
+  url    = {https://github.com/linkedin/Liger-Kernel},
+  year   = {2024}
+}
+```

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/setup.py

@@ -1,6 +1,6 @@
 from setuptools import find_namespace_packages, setup
 
-__version__ = "0.0.0"
+__version__ = "0.1.0"
 
 setup(
     name="liger_kernel",

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/cross_entropy.py

@@ -17,7 +17,7 @@ def liger_cross_entropy_kernel(
     BLOCK_SIZE: tl.constexpr,
 ):
     """
-    This kernel computes both cross entropy loss and the gradient of the _input.
+    This kernel computes both cross entropy loss and the gradient of the input.
     We only consider hard label + mean reduction for now. Please refer to https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html for the math.
 
     Parameters:
@@ -34,7 +34,7 @@ def liger_cross_entropy_kernel(
     """
 
     # https://github.com/triton-lang/triton/issues/1058
-    # Essentially if B*T*V is too large, program_id * stride will overflow out of int32
+    # If B*T*V is too large, program_id * stride will overflow out of int32, so we convert to int64
     program_id = tl.program_id(0).to(tl.int64)
 
     # 1. Load Y_ptr first because if the target is ignore_index, we can return right away
@@ -90,13 +90,7 @@ def liger_cross_entropy_kernel(
     tl.debug_barrier()
 
     # 5. Calculate the loss
-    # Old Approach: Problematic LogSoftmax
-    # min of bfloat16 and float32 is 1e-38, so we set a value larger than that but small enough
-    # This will overflow if X_y * n_non_ignore is too small. Even if we add a tiny epsilon, it will still overflow
-    # loss = -tl.log(X_y * n_non_ignore)
 
-    # New Approach: Safe LogSoftmax
-    # Therefore, we propose to use safe logsoftmax by reordering the formula.
     # loss = log (softmax(X_y)) = log ((e ^ (X_y - max(X)) / sum(e ^ (X - max(X))))
     #      = (X_y - max(X)) - log(sum(e ^ (X - max(X))))
     # sum(e ^ (X - max(X))) must >= 1 because the max term is e ^ 0 = 1
@@ -114,7 +108,7 @@ def liger_cross_entropy_kernel(
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 65536 // 2  # manual tune a bit
+MAX_FUSED_SIZE = 65536 // 2  # the best size we found by manually tuning
 
 
 @triton.jit
@@ -184,28 +178,6 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         n_non_ignore = (target != ignore_index).sum().item()
 
         # ensure _input and target are contiguous in the last dimension
-        # there are examples that are NOT contiguous overall but contiguous in the last dimension
-        ####################################################################
-        # tensor = torch.arange(1, 21).reshape(5, -1)
-        # print(tensor)
-        # tensor([[ 1,  2,  3,  4],
-        # [ 5,  6,  7,  8],
-        # [ 9, 10, 11, 12],
-        # [13, 14, 15, 16],
-        # [17, 18, 19, 20]])
-        # print(tensor.is_contiguous())
-        # True
-        # slice = tensor[::2, :]
-        # print(slice)
-        # tensor([[ 1,  2,  3,  4],
-        # [ 9, 10, 11, 12],
-        # [17, 18, 19, 20]])
-        # print(slice.is_contiguous())
-        # False
-        # print(slice.stride())
-        # (8, 1)
-        # slice is NOT a contiguous tensor but is contiguous in the last dimension, CE kernel can execute because the stride is 8, and each triton program will jump by 8
-        ####################################################################
         if _input.stride(-1) != 1:
             _input = _input.contiguous()
         if target.stride(-1) != 1:
@@ -252,10 +224,9 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
         if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
             pass
+
         # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
         # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-        # Although the Brew trainer should only perform backward once, it encounters this issue.
-        # https://github.com/triton-lang/triton/issues/4004
         else:
             BT, V = _input.shape
             n_rows = BT

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -1,8 +1,3 @@
-"""Fusing the last linear layer with cross-entropy loss
-
-Reference: https://github.com/mgmalek/efficient_cross_entropy
-"""
-
 import torch
 import triton
 
@@ -11,13 +6,16 @@ from liger_kernel.ops.cross_entropy import element_mul, liger_cross_entropy_kern
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 65536 // 2  # manual tune a bit
+MAX_FUSED_SIZE = 65536 // 2
 
 
 class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, _input, linear, target, ignore_index):
         """
+        Fusing the last linear layer with cross-entropy loss
+            Reference: https://github.com/mgmalek/efficient_cross_entropy
+
         Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding
         the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can
         compute the gradient at the forward pass. By doing so, we don't have to store the _input and target
@@ -54,6 +52,8 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
 
         grad_linear = torch.zeros_like(linear, device=device)
         grad_input = torch.zeros_like(_input, device=device)
+
+        # we use fp32 for loss accumulator
         loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
 
         total_n_non_ignore = (target != ignore_index).sum().item()

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/geglu.py

@@ -1,8 +1,19 @@
+import operator
+
 import torch
 import triton
 import triton.language as tl
 
-from liger_kernel.ops.utils import calculate_settings, ensure_contiguous
+from liger_kernel.ops.utils import (
+    calculate_settings,
+    compare_version,
+    ensure_contiguous,
+)
+
+if compare_version("triton", operator.ge, "3.0.0"):
+    from triton.language.extra.libdevice import tanh
+else:
+    from triton.language.math import tanh
 
 
 @triton.jit
@@ -26,7 +37,7 @@ def _geglu_tanh_forward_kernel(
     sqrt_2_over_pi = 0.7978845608028654  # sqrt(2 / pi)
     a_cubed = a_row * a_row * a_row
     tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
-    tanh_result = tl.math.tanh(tanh_arg)
+    tanh_result = tanh(tanh_arg)
     geglu_a = 0.5 * a_row * (1 + tanh_result)
     c_row = geglu_a * b_row
     tl.store(c + col_offsets, c_row, mask=mask)
@@ -54,7 +65,7 @@ def _geglu_tanh_backward_kernel(
     sqrt_2_over_pi = 0.7978845608028654  # sqrt(2 / pi)
     a_cubed = a_row * a_row * a_row
     tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
-    tanh_result = tl.math.tanh(tanh_arg)
+    tanh_result = tanh(tanh_arg)
     geglu_a = 0.5 * a_row * (1 + tanh_result)
 
     db_row = dc_row * geglu_a

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/rms_norm.py

@@ -20,9 +20,12 @@ def _rms_norm_forward(
     BLOCK_SIZE: tl.constexpr,
 ):
     """
+    y_i = (x_i / (RMS)) * wi, RMS = sqrt(sum(x_i^2) / N)
+
     Reference:
     1. https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
     2. https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/unsloth/kernels/rms_layernorm.py#L22
+    3. https://arxiv.org/pdf/1910.07467
     """
 
     row_idx = tl.program_id(0)
@@ -36,16 +39,17 @@ def _rms_norm_forward(
     X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
     W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
 
-    row_var = tl.sum(X_row * X_row, axis=0) / n_cols
-    inv_var = tl.math.rsqrt(row_var + eps)
+    mean_square = tl.sum(X_row * X_row, axis=0) / n_cols
+    inv_rms = tl.math.rsqrt(mean_square + eps)
 
-    # trick: row_var is tiny compared to X_row because it just has one per row we can save 4 ops (*, sum, /, rqrt) if we cache it
-    tl.store(r_ptr, inv_var)
+    # We can save time by caching rms with minimal memory overhead
+    # because rms is much smaller compared to X_row, as rms is for each row.
+    # However, on the computation side, it can save 4 operations (*, sum, /, sqrt).
+    tl.store(r_ptr, inv_rms)
 
-    normed = X_row * inv_var
+    Y_row = X_row * inv_rms * W_row
 
-    output = normed * W_row
-    tl.store(Y_ptr + col_offsets, output, mask=mask)
+    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
 
 
 @triton.jit
@@ -65,9 +69,10 @@ def _rms_norm_backward(
     BLOCK_SIZE: tl.constexpr,
 ):
     """
-    dx = (1 / var(x)) * (dy * w - (1/N) * (dy * w) dot x) * x
-    dw = sum(dy * (x / var(x)))
+    dx = (1 / RMS) * [dy * w  - (1 / N) * (1 / RMS^2) * ((dy * w) dot x) * x]. * means element-wise multiplication, whileas dot means dot product
+    dw = sum(dy * (x / RMS)). summation over BxT dimension
     """
+
     row_idx = tl.program_id(0)
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
@@ -81,34 +86,42 @@ def _rms_norm_backward(
     X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
     W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
 
-    # Get saved row variance
-    inv_var = tl.load(r_ptr)
-
-    normed = X_row * inv_var
+    # Get cached rms
+    inv_rms_row = tl.load(r_ptr)
 
-    dY_W = dY_row * W_row
-    dY_normed = dY_row * normed
-
-    rowsum_dY_normed = tl.sum(dY_W * normed, axis=0)
-    output = inv_var / n_cols * (n_cols * dY_W - normed * rowsum_dY_normed)
-    tl.store(dY_ptr + col_offsets, output, mask=mask)
+    dX_row = (inv_rms_row) * (
+        dY_row * W_row
+        - (1 / n_cols)
+        * inv_rms_row
+        * inv_rms_row
+        * tl.sum(dY_row * W_row * X_row, axis=0)
+        * X_row
+    )
+    tl.store(dY_ptr + col_offsets, dX_row, mask=mask)
 
     # calculate the gradient of W
-    tl.store(dW_ptr + col_offsets, dY_normed, mask=mask)
+    dW_row = dY_row * X_row * inv_rms_row
+    tl.store(dW_ptr + col_offsets, dW_row, mask=mask)
 
 
 class LigerRMSNormFunction(torch.autograd.Function):
     @staticmethod
     @ensure_contiguous
     def forward(ctx, X, W, eps):
+        """
+        X: (B, T, H) or (BxT, H)
+        W: (H,)
+        """
+
         shape = X.shape
         dim = shape[-1]
         X = X.view(-1, dim)
         n_rows, n_cols = X.shape
         BLOCK_SIZE, num_warps = calculate_settings(n_cols)
 
-        Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device="cuda")
-        r = torch.empty(n_rows, dtype=X.dtype, device="cuda")
+        Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
+        # r is to cache (1/rms) for each row
+        r = torch.empty(n_rows, dtype=X.dtype, device=X.device)
 
         # Check constraints.
         assert (
@@ -139,6 +152,10 @@ class LigerRMSNormFunction(torch.autograd.Function):
     @staticmethod
     @ensure_contiguous
     def backward(ctx, dY):
+        """
+        Y: (B, T, H) or (BxT, H)
+        """
+
         shape = dY.shape
         dim = shape[-1]
         dY = dY.view(-1, dim)
@@ -146,6 +163,7 @@ class LigerRMSNormFunction(torch.autograd.Function):
         n_rows, n_cols = dY.shape
         dW = torch.zeros_like(X)
 
+        # Here we use dY to store the value of dX to save memory
         _rms_norm_backward[(n_rows,)](
             dY,
             dY.stride(0),

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/swiglu.py

@@ -12,43 +12,43 @@ def silu(x):
 
 @triton.jit
 def _swiglu_forward_kernel(
-    a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
     program_id = tl.program_id(0)
 
     # locate start index
-    a += program_id * stride
-    b += program_id * stride
-    c += program_id * stride
+    a_ptr += program_id * stride
+    b_ptr += program_id * stride
+    c_ptr += program_id * stride
 
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
     # sigmoid requires type float32
-    a_row = tl.load(a + col_offsets, mask=mask, other=0).to(tl.float32)
-    b_row = tl.load(b + col_offsets, mask=mask, other=0)
+    a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
+    b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
     c_row = silu(a_row) * b_row
-    tl.store(c + col_offsets, c_row, mask=mask)
+    tl.store(c_ptr + col_offsets, c_row, mask=mask)
 
 
 @triton.jit
 def _swiglu_backward_kernel(
-    dc, a, b, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
     program_id = tl.program_id(0)
 
     # locate start index
-    dc += program_id * stride
-    a += program_id * stride
-    b += program_id * stride
+    dc_ptr += program_id * stride
+    a_ptr += program_id * stride
+    b_ptr += program_id * stride
 
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    dc_row = tl.load(dc + col_offsets, mask=mask, other=0)
+    dc_row = tl.load(dc_ptr + col_offsets, mask=mask, other=0)
     # sigmoid requires type float32
-    a_row = tl.load(a + col_offsets, mask=mask, other=0).to(tl.float32)
-    b_row = tl.load(b + col_offsets, mask=mask, other=0)
+    a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
+    b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
 
     # recomputation to save memory
     sig_a = tl.sigmoid(a_row)
@@ -56,8 +56,8 @@ def _swiglu_backward_kernel(
     db_row = dc_row * silu_a
     da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row
 
-    tl.store(a + col_offsets, da_row, mask=mask)
-    tl.store(b + col_offsets, db_row, mask=mask)
+    tl.store(a_ptr + col_offsets, da_row, mask=mask)
+    tl.store(b_ptr + col_offsets, db_row, mask=mask)
 
 
 class LigerSiLUMulFunction(torch.autograd.Function):

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/ops/utils.py

@@ -1,7 +1,10 @@
 import functools
+import importlib
+from typing import Callable
 
 import torch
 import triton
+from packaging.version import Version
 
 
 def ensure_contiguous(fn):
@@ -36,3 +39,12 @@ def calculate_settings(n):
     elif BLOCK_SIZE >= 2048:
         num_warps = 8
     return BLOCK_SIZE, num_warps
+
+
+def compare_version(package: str, operator: Callable, target: str):
+    try:
+        pkg = importlib.import_module(package)
+    except ImportError:
+        return False
+    pkg_version = Version(pkg.__version__)
+    return operator(pkg_version, Version(target))

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/transformers/__init__.py

@@ -1,4 +1,5 @@
 from liger_kernel.transformers.monkey_patch import (  # noqa: F401
+    apply_liger_kernel_to_gemma,
     apply_liger_kernel_to_llama,
     apply_liger_kernel_to_mistral,
     apply_liger_kernel_to_mixtral,

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/transformers/model/llama.py

@@ -37,6 +37,9 @@ def lce_forward(
     cache_position: Optional[torch.LongTensor] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
+    Copy paste llama forward but replace torch cross entropy with liger fused linear cross entropy
+
+
     Args:
         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
             Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/transformers/monkey_patch.py

@@ -1,27 +1,26 @@
 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.rms_norm import LigerRMSNorm
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP, LigerSwiGLUMLP
 
 
-# TODO: probably rename utils.py as hf_patcher.py to be more descriptive
 def apply_liger_kernel_to_llama(
     rope: bool = True,
-    cross_entropy: bool = True,
-    fused_linear_cross_entropy: bool = False,
+    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 Llama models (2 and 3)
-    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.
+        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 False.
+            Whether to apply Liger's fused lienar 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.
@@ -54,7 +53,6 @@ def apply_liger_kernel_to_mistral(
 ) -> None:
     """
     Apply Liger kernels to replace original implementation in HuggingFace Mistral models
-    to make GPU go burrr.
 
     Args:
         rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
@@ -83,12 +81,12 @@ def apply_liger_kernel_to_mixtral(
 ) -> None:
     """
     Apply Liger kernels to replace original implementation in HuggingFace Mixtral models
-    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.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
     """
 
     from transformers.models.mixtral import modeling_mixtral
@@ -101,3 +99,32 @@ def apply_liger_kernel_to_mixtral(
         modeling_mixtral.CrossEntropyLoss = LigerCrossEntropyLoss
     if swiglu:
         modeling_mixtral.MixtralBlockSparseTop2MLP = LigerBlockSparseTop2MLP
+
+
+def apply_liger_kernel_to_gemma(
+    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 models
+    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.
+    """
+    # TODO(yundai424): add convergence test for gemma
+    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
+    if cross_entropy:
+        modeling_gemma.CrossEntropyLoss = LigerCrossEntropyLoss
+    if geglu:
+        modeling_gemma.GemmaMLP = LigerGEGLUMLP

liger_kernel-0.1.0/src/liger_kernel/transformers/trainer_integration.py

@@ -0,0 +1,45 @@
+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)

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel/triton/monkey_patch.py

@@ -1,12 +1,10 @@
 import os
 import random
 
-from overrides import override
 from triton.runtime.cache import FileCacheManager
 
 
 class LigerTritonFileCacheManager(FileCacheManager):
-    @override
     def put(self, data, filename, binary=True) -> str:
         if not self.cache_dir:
             raise RuntimeError("Could not create or locate cache dir")

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel.egg-info/PKG-INFO

@@ -1,4 +1,6 @@
 Metadata-Version: 2.1
 Name: liger-kernel
-Version: 0.0.0
+Version: 0.1.0
 Provides-Extra: dev
+License-File: LICENSE
+License-File: NOTICE

{liger_kernel-0.0.0 → liger_kernel-0.1.0}/src/liger_kernel.egg-info/SOURCES.txt

@@ -1,3 +1,6 @@
+LICENSE
+NOTICE
+README.md
 setup.py
 src/liger_kernel.egg-info/PKG-INFO
 src/liger_kernel.egg-info/SOURCES.txt
@@ -20,6 +23,7 @@ src/liger_kernel/transformers/monkey_patch.py
 src/liger_kernel/transformers/rms_norm.py
 src/liger_kernel/transformers/rope.py
 src/liger_kernel/transformers/swiglu.py
+src/liger_kernel/transformers/trainer_integration.py
 src/liger_kernel/transformers/model/__init__.py
 src/liger_kernel/transformers/model/llama.py
 src/liger_kernel/triton/__init__.py