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

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,26 +86,33 @@ 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)
@@ -108,6 +120,7 @@ class LigerRMSNormFunction(torch.autograd.Function):
         BLOCK_SIZE, num_warps = calculate_settings(n_cols)
 
         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.
@@ -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/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/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/transformers/monkey_patch.py

@@ -1,4 +1,5 @@
 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
@@ -98,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/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/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.1.dist-info → liger_kernel-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger-kernel
-Version: 0.0.1
+Version: 0.1.0
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2

{liger_kernel-0.0.1.dist-info → liger_kernel-0.1.0.dist-info}/RECORD

@@ -2,25 +2,26 @@ liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 liger_kernel/ops/cross_entropy.py,sha256=YTHKVyPW748EWtbWJeKdIe9S1dEq6i90_PbBuCD-9s0,9178
 liger_kernel/ops/fused_linear_cross_entropy.py,sha256=58MmDhLJGR5b8ixztkhR707yp0VY28oBRASFVwGbeV8,7346
 liger_kernel/ops/geglu.py,sha256=5tGinryOOYRpGtKwJ4B1ertwtzd81xdjevD3Ha7H1AY,3849
-liger_kernel/ops/rms_norm.py,sha256=Tyz5Ea7U8dNtNUpuRmT6qsV7PmDe0FuUFjaEPTsFu1E,4303
+liger_kernel/ops/rms_norm.py,sha256=AQ1jaCXUlrBazqAPg-Cpf2K5OsO4byDKcdfWsGy9-zI,4848
 liger_kernel/ops/rope.py,sha256=fYBct8gDQfKPZdMWlzkZZ8kBzh6nQ7DIpDsc7lZwM8c,8584
-liger_kernel/ops/swiglu.py,sha256=__QsfYxKyZHtRScm31zL3sAOVEblQFqKj2ll8I4Odqg,2835
+liger_kernel/ops/swiglu.py,sha256=MRbSIXsBLqlFr9ZdtuFqSjLJJ-716URmQIhxQ57GGEw,2915
 liger_kernel/ops/utils.py,sha256=vsFIywd8LQlVPRA3RPZOm5HyN8c0cS4NFEEnwjNw-MI,1427
-liger_kernel/transformers/__init__.py,sha256=7rOw9yZ8kNXO483Colx-EUq8GcTCvCZxrxF-S7pmkkU,172
+liger_kernel/transformers/__init__.py,sha256=nVvk0h7er3fdgubQF8Z8KjA3ew-q5oJHyJRg5cKmBoc,205
 liger_kernel/transformers/cross_entropy.py,sha256=G-L4EaUYVc25NKZ2jrlaG-d5YUvDqJdUlawPN7K1d1g,389
 liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=h0AW9ubFGfz4DBwgh2CLW8rpKo9PvxYpB6AUzjx-1b0,501
 liger_kernel/transformers/geglu.py,sha256=FrLBHZRdI68jw9RR6MSTE59-xCzueOwSRp9jL8y-j98,896
-liger_kernel/transformers/monkey_patch.py,sha256=9CilRC9pBBbQ8R1_4HLsZq2xfmxVC4xGx345vfejX6I,3914
+liger_kernel/transformers/monkey_patch.py,sha256=FjaRZVWm_ZMHO3NXc4IT6EpCTWJOdZKP72mZq01qbrA,5006
 liger_kernel/transformers/rms_norm.py,sha256=2LHfEctSpzuNRaoZ9uUECSFK8fZeIxIsHm9QbEHZvDQ,452
 liger_kernel/transformers/rope.py,sha256=m-ah8vZBYW8tfplTXCiAPMHJWlB1tdp_JPXJeWE-Boo,943
 liger_kernel/transformers/swiglu.py,sha256=8kt4MffEZT5vx3k0WA-GO-WPLv5kGdnu_nAwlJyMI2U,1516
+liger_kernel/transformers/trainer_integration.py,sha256=gt0fF-se2XiIB6PocHBPBuD6tLCOtQRcb20WfUS2ceA,1645
 liger_kernel/transformers/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/transformers/model/llama.py,sha256=4mfVTMrY7T-xiJeQJe02hBVnAwNCKlvLGp49gj6TWiU,5298
 liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
-liger_kernel/triton/monkey_patch.py,sha256=yRNaGdyG5PrwX5ed_MQdqtqvvpVvQ7ZD2FQ_9W1q9u8,1629
-liger_kernel-0.0.1.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel-0.0.1.dist-info/METADATA,sha256=2PhmP9NVtu0CsGG2_jnxCukPTMZx6vnzLpTQlJDrqq4,504
-liger_kernel-0.0.1.dist-info/NOTICE,sha256=BXkXY9aWvEy_7MAB57zDu1z8uMYT1i1l9B6EpHuBa8s,173
-liger_kernel-0.0.1.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-liger_kernel-0.0.1.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel-0.0.1.dist-info/RECORD,,
+liger_kernel/triton/monkey_patch.py,sha256=5BcGKTtdqeYchypBIBopGIWPx1-cFALz7sOKoEsqXJ0,1584
+liger_kernel-0.1.0.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel-0.1.0.dist-info/METADATA,sha256=E_OSiFz2sC4jmWO4VH3sTXWiR3Ev7qNy5oSLSWk-s8g,504
+liger_kernel-0.1.0.dist-info/NOTICE,sha256=BXkXY9aWvEy_7MAB57zDu1z8uMYT1i1l9B6EpHuBa8s,173
+liger_kernel-0.1.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+liger_kernel-0.1.0.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel-0.1.0.dist-info/RECORD,,