liger-kernel 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

liger_kernel/ops/rms_norm.py

@@ -10,6 +10,7 @@ https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddec
 Modifications made by Yanning Chen, 2024.
 """
 
+import math
 import operator
 
 import torch
@@ -20,6 +21,7 @@ from liger_kernel.ops.utils import (
     calculate_settings,
     compare_version,
     ensure_contiguous,
+    torch_to_triton_dtype,
 )
 
 if compare_version("triton", operator.ge, "3.0.0"):
@@ -84,6 +86,10 @@ def _rms_norm_forward_kernel(
         W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
+    if casting_mode == _CASTING_MODE_NONE:
+        eps = eps.to(X_row_dtype)
+        offset = offset.to(X_row_dtype)
+
     mean_square = tl.sum(X_row * X_row, axis=0) / n_cols
     rstd = rsqrt(mean_square + eps)
 
@@ -100,6 +106,9 @@ def _rms_norm_forward_kernel(
 
     Y_row = X_row * (offset + W_row)
 
+    if casting_mode == _CASTING_MODE_GEMMA:
+        Y_row = Y_row.to(X_row_dtype)
+
     tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
 
 
@@ -109,14 +118,17 @@ def _rms_norm_backward_kernel(
     dY_row_stride,
     X_ptr,
     X_row_stride,
+    X_dtype: tl.constexpr,
     W_ptr,
     W_row_stride,
     RSTD_ptr,
     RSTD_row_stride,
     dW_ptr,
     dW_row_stride,
+    n_rows,
     n_cols,
     offset,
+    rows_per_program: tl.constexpr,
     casting_mode: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
@@ -125,54 +137,60 @@ def _rms_norm_backward_kernel(
     dw = sum(dy * (x / RMS)). summation over BxT dimension
     """
 
-    row_idx = tl.program_id(0)
+    row_block_id = tl.program_id(0)
+    row_start = row_block_id * rows_per_program
+    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    dY_ptr += row_idx * dY_row_stride
-    X_ptr += row_idx * X_row_stride
-    RSTD_ptr += row_idx * RSTD_row_stride
-    dW_ptr += row_idx * dW_row_stride
+    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0)
-    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
-    original_x_dtype = X_row.dtype
-
-    # Get cached rms
-    rstd_row = tl.load(RSTD_ptr)
+    dY_ptr += row_start * dY_row_stride
+    X_ptr += row_start * X_row_stride
+    RSTD_ptr += row_start
 
+    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
     W_row = W_row + offset
 
-    X_row = X_row.to(tl.float32)
+    for _ in range(row_start, row_end):
+        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
 
-    # Different bacward graphs for different casting modes
-    if casting_mode == _CASTING_MODE_LLAMA:
-        m = (dY_row * W_row).to(tl.float32)
+        # Get cached rms
+        rstd_row = tl.load(RSTD_ptr)
 
-    elif casting_mode == _CASTING_MODE_GEMMA:
-        dY_row, W_row = (
-            dY_row.to(tl.float32),
-            W_row.to(tl.float32),
-        )
+        X_row = X_row.to(tl.float32)
 
-    m = dY_row * W_row
+        # Different bacward graphs for different casting modes
+        if casting_mode == _CASTING_MODE_LLAMA:
+            m = (dY_row * W_row).to(tl.float32)
 
-    dX_row = rstd_row * m
+        elif casting_mode == _CASTING_MODE_GEMMA:
+            dY_row = dY_row.to(tl.float32)
+            m = dY_row * W_row
+        else:
+            m = dY_row * W_row
 
-    dX_row += (rstd_row) * (
-        -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
-    )
+        dX_row = rstd_row * m
 
-    # calculate the gradient of W
-    if casting_mode == _CASTING_MODE_LLAMA:
-        dW_row = dY_row * (X_row * rstd_row).to(original_x_dtype)
-    else:
-        # here X_row is already in fp32 (see previous if block)
-        dW_row = dY_row * (X_row * rstd_row)
+        dX_row += (rstd_row) * (
+            -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
+        )
+
+        # calculate the gradient of W
+        if casting_mode == _CASTING_MODE_LLAMA:
+            dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
+        else:
+            # here X_row is already in fp32 (see previous if block)
+            dW_row += dY_row * (X_row * rstd_row)
 
-    tl.store(dY_ptr + col_offsets, dX_row, mask=mask)
-    tl.store(dW_ptr + col_offsets, dW_row, mask=mask)
+        tl.store(dY_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
+
+        dY_ptr += dY_row_stride
+        X_ptr += X_row_stride
+        RSTD_ptr += RSTD_row_stride
+
+    tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
 
 _str_to_casting_mode = {
@@ -238,31 +256,38 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     dim = shape[-1]
     dY = dY.view(-1, dim)
     n_rows, n_cols = dY.shape
-    dW = torch.empty_like(
-        X,
-        dtype=(torch.float32 if casting_mode == _CASTING_MODE_GEMMA.value else W.dtype),
-    )
 
+    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    # fp32 for numerical stability especially.
+    _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+
+    if n_cols > BLOCK_SIZE:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    rows_per_program = math.ceil(n_rows / sm_count)
+    grid = (sm_count,)
     # Here we use dY to store the value of dX to save memory
-    _rms_norm_backward_kernel[(n_rows,)](
+    _rms_norm_backward_kernel[grid](
         dY,
         dY.stride(0),
         X,
         X.stride(0),
+        torch_to_triton_dtype[X.dtype],
         W,
         W.stride(0),
         RSTD,
         RSTD.stride(0),
-        dW,
-        dW.stride(0),
+        _dW,
+        _dW.stride(0),
+        n_rows,
         n_cols,
         offset,
+        rows_per_program,
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
     )
     dX = dY.view(*shape)
-    dW = torch.sum(dW, dim=0).to(W.dtype)
+    dW = _dW.sum(dim=0).to(W.dtype)
     return dX, dW
 
 

liger_kernel/ops/swiglu.py

@@ -14,7 +14,7 @@ def silu(x):
 def _swiglu_forward_kernel(
     a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     a_ptr += program_id * stride
@@ -35,7 +35,7 @@ def _swiglu_forward_kernel(
 def _swiglu_backward_kernel(
     dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     dc_ptr += program_id * stride

liger_kernel/ops/utils.py

@@ -12,13 +12,19 @@ Modifications made by Yanning Chen, 2024.
 
 import functools
 import importlib
+import operator
 from typing import Callable
 
 import torch
 import triton
+import triton.language as tl
 from packaging.version import Version
 
 
+def is_hip() -> bool:
+    return torch.version.hip is not None
+
+
 def ensure_contiguous(fn):
     @functools.wraps(fn)
     def wrapper(ctx, *args, **kwargs):
@@ -45,7 +51,7 @@ def calculate_settings(n):
 
     num_warps = 4
     if BLOCK_SIZE >= 32768:
-        num_warps = 32
+        num_warps = 32 if not is_hip() else 16
     elif BLOCK_SIZE >= 8192:
         num_warps = 16
     elif BLOCK_SIZE >= 2048:
@@ -60,3 +66,58 @@ def compare_version(package: str, operator: Callable, target: str):
         return False
     pkg_version = Version(pkg.__version__)
     return operator(pkg_version, Version(target))
+
+
+def get_amp_custom_fwd_bwd() -> Callable:
+    if compare_version("torch", operator.ge, "2.4.0"):
+        return (
+            functools.partial(torch.amp.custom_fwd, device_type="cuda"),
+            functools.partial(torch.amp.custom_bwd, device_type="cuda"),
+        )
+    return torch.cuda.amp.custom_fwd, torch.cuda.amp.custom_bwd
+
+
+amp_custom_fwd, amp_custom_bwd = get_amp_custom_fwd_bwd()
+
+
+torch_to_triton_dtype = {
+    torch.float32: tl.float32,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+@triton.jit
+def element_mul_kernel(
+    X_ptr,
+    X_stride,
+    grad_output_ptr,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr.
+    The multiplication is performed in-place on the tensor pointed by X_ptr.
+
+    Parameters:
+    X_ptr: Pointer to the input tensor.
+    X_stride (int): The stride of the input tensor.
+    grad_output_ptr: Pointer to the gradient output value.
+    n_cols (int): The number of columns in the input tensor.
+    BLOCK_SIZE (int): The block size for Triton operations.
+    """
+
+    # Get the program ID and convert it to int64 to avoid overflow
+    program_id = tl.program_id(0).to(tl.int64)
+
+    # Locate the start index
+    X_ptr += program_id * X_stride
+
+    # Load the gradient output value
+    grad_output = tl.load(grad_output_ptr)
+
+    # Perform the element-wise multiplication
+    for i in range(0, n_cols, BLOCK_SIZE):
+        X_offsets = i + tl.arange(0, BLOCK_SIZE)
+        X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols)
+        tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols)

liger_kernel/transformers/__init__.py

@@ -5,7 +5,9 @@ from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noq
 from liger_kernel.transformers.fused_linear_cross_entropy import (  # noqa: F401
     LigerFusedLinearCrossEntropyLoss,
 )
+from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD  # noqa: F401
 from liger_kernel.transformers.geglu import LigerGEGLUMLP  # noqa: F401
+from liger_kernel.transformers.jsd import LigerJSD  # noqa: F401
 from liger_kernel.transformers.layer_norm import LigerLayerNorm  # noqa: F401
 from liger_kernel.transformers.monkey_patch import (  # noqa: F401
     _apply_liger_kernel,
@@ -15,6 +17,7 @@ from liger_kernel.transformers.monkey_patch import (  # noqa: F401
     apply_liger_kernel_to_llama,
     apply_liger_kernel_to_mistral,
     apply_liger_kernel_to_mixtral,
+    apply_liger_kernel_to_mllama,
     apply_liger_kernel_to_phi3,
     apply_liger_kernel_to_qwen2,
     apply_liger_kernel_to_qwen2_vl,

liger_kernel/transformers/auto_model.py

@@ -1,6 +1,11 @@
+import inspect
+
 from transformers import AutoConfig, AutoModelForCausalLM
 
-from liger_kernel.transformers.monkey_patch import _apply_liger_kernel
+from liger_kernel.transformers.monkey_patch import (
+    MODEL_TYPE_TO_APPLY_LIGER_FN,
+    _apply_liger_kernel,
+)
 
 
 def _get_model_config(model_dir, **model_init_kwargs):
@@ -21,13 +26,20 @@ class AutoLigerKernelForCausalLM(AutoModelForCausalLM):
         # Determine the model type and apply the Liger Kernel if applicable
         # Note: _apply_liger_kernel will only pass relevant kwargs to the apply_liger_kernel_to_* function
         model_type = model_config.model_type
+
         _apply_liger_kernel(model_type, **kwargs)
 
-        # Retain only the keyword args present in the model configuration
-        for k in list(kwargs.keys()):
-            if k not in model_config.__dict__:
-                del kwargs[k]
+        # Filter out kwargs that were passed to the apply_liger_* function, which will cause
+        # model initialization errors otherwise
+        apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type]
+        apply_fn_signature = inspect.signature(apply_fn)
+
+        applicable_kwargs = {
+            key: value
+            for key, value in kwargs.items()
+            if key not in apply_fn_signature.parameters
+        }
 
         return super().from_pretrained(
-            pretrained_model_name_or_path, *model_args, **kwargs
+            pretrained_model_name_or_path, *model_args, **applicable_kwargs
         )

liger_kernel/transformers/functional.py

@@ -2,7 +2,9 @@ from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
 from liger_kernel.ops.fused_linear_cross_entropy import (
     LigerFusedLinearCrossEntropyFunction,
 )
+from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
 from liger_kernel.ops.geglu import LigerGELUMulFunction
+from liger_kernel.ops.jsd import LigerJSDFunction
 from liger_kernel.ops.kl_div import LigerKLDivLossFunction
 from liger_kernel.ops.layer_norm import LigerLayerNormFunction
 from liger_kernel.ops.rms_norm import LigerRMSNormFunction
@@ -17,3 +19,5 @@ liger_rms_norm = LigerRMSNormFunction.apply
 liger_rope = LigerRopeFunction.apply
 liger_layer_norm = LigerLayerNormFunction.apply
 liger_kl_div = LigerKLDivLossFunction.apply
+liger_jsd = LigerJSDFunction.apply
+liger_fused_linear_jsd = LigerFusedLinearJSDFunction.apply

liger_kernel/transformers/fused_linear_jsd.py

@@ -0,0 +1,98 @@
+from typing import Optional
+
+import torch
+
+from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
+
+
+class LigerFusedLinearJSD(torch.nn.Module):
+    r"""Fusing the last linear layer with generalized JSD
+
+    Handle the forward and backward pass of the final linear layer via JSD by avoiding
+    the materialization of the large logits tensor.
+
+    Args:
+        jsd_beta (float): coefficient beta of generalized JSD in the open interval (0, 1). Default: `0.5`
+        ignore_index (int): The index to ignore in the target. Default: `-100`
+        temperature (float): temperature in softmax function to control the output probability distribution. Default: `1.0`
+
+    Shape:
+        - student_input: :math:`(BT, H)`, where B is batch size, T is sequence length, H is hidden dimension.
+        - student_weight: :math:`(V, H)`, where V is vocab size.
+        - teacher_input: :math:`(BT, H')`, where H' is hidden dimension of the teacher model.
+        - teacher_weight: :math:`(V, H')`, where hidden size H and H' can be different.
+        - shift_labels: :math:`(BT,)`
+        - Output: a scalar.
+
+    Examples:
+    ```python
+    >>> (B, T, H_s, H_t, V) = (2, 2, 3, 5, 10)
+    >>> fused_jsd = LigerFusedLinearJSD(jsd_beta=0.1, temperature=2.0)
+    >>> # generate inputs and weights
+    >>> student_input = torch.rand(B * T, H_s, device="cuda", requires_grad=True)
+    >>> student_lin = torch.nn.Linear(H_s, V, bias=False, device="cuda")
+    >>> # teacher input doesn't require grad, hidden_dim can be different from student's
+    >>> teacher_input = torch.rand(B * T, H_t, device="cuda")
+    >>> teacher_lin = torch.nn.Linear(H_t, V, bias=False, device="cuda")
+    >>> output = fused_jsd(student_input, student_lin.weight, teacher_input, teacher_lin.weight)
+    >>> output.backward()
+    >>>
+    >>> # Example with labels for supervised fine-tuning (SFT) context:
+    >>>
+    >>> # Assume hidden_states, lm_heads and corresponding labels are given
+    >>> student_lm_head = torch.nn.Linear(H_s, V, bias=False)
+    >>> student_hidden_states = torch.randn(B * T, H_s, requires_grad=True).log_softmax(dim=-1)
+    >>> teacher_lm_head = torch.nn.Linear(H_t, V, bias=False)
+    >>> teacher_hidden_states = torch.randn(B * T, H_t).log_softmax(dim=-1)
+    >>> labels = torch.randint(0, V, (B * T,), torch.long)
+    >>>
+    >>> # Shift so that tokens < n predict n
+    >>> shift_student_hidden_states = student_hidden_states[..., :-1, :].contiguous()
+    >>> shift_teacher_hidden_states = teacher_hidden_states[..., :-1, :].contiguous()
+    >>> shift_labels = labels[..., 1:].contiguous()
+    >>>
+    >>> # Flatten tokens
+    >>> shift_student_hidden_states = shift_student_hidden_states.view(-1, V)
+    >>> shift_teacher_hidden_states = shift_teacher_hidden_states.view(-1, V)
+    >>> shift_labels = shift_labels.view(-1)
+    >>>
+    >>> # Calculate loss
+    >>> loss_fct = LigerJSD(beta=0.1)
+    >>> loss = loss_fct(
+    >>>     shift_studetn_hidden_states,
+    >>>     student_lm_head.weight,
+    >>>     shift_teacher_hidden_states,
+    >>>     teacher_lm_head.weight,
+    >>>     shift_labels
+    >>> )
+    ```
+    """
+
+    def __init__(self, jsd_beta=0.5, ignore_index=-100, temperature=1.0):
+        super().__init__()
+        assert (
+            jsd_beta > 0 and jsd_beta < 1
+        ), f"beta must be greater than 0 and less than 1. Got: {jsd_beta}"
+        assert temperature != 0, "temperature cannot be 0."
+        self.jsd_beta = jsd_beta
+        self.temperature = temperature
+        self.ignore_index = ignore_index
+
+    def forward(
+        self,
+        student_input: torch.Tensor,
+        student_weight: torch.Tensor,
+        teacher_input: torch.Tensor,
+        teacher_weight: torch.Tensor,
+        shift_labels: Optional[torch.LongTensor],
+    ):
+        return LigerFusedLinearJSDFunction.apply(
+            student_input,
+            student_weight,
+            teacher_input,
+            teacher_weight,
+            shift_labels,
+            self.jsd_beta,
+            self.ignore_index,
+            self.temperature,
+        )

liger_kernel/transformers/jsd.py

@@ -0,0 +1,75 @@
+from typing import Optional
+
+import torch
+
+from liger_kernel.ops.jsd import LigerJSDFunction
+
+
+class LigerJSD(torch.nn.Module):
+    r"""The generalized Jensen-Shannon Divergence.
+    .. math::
+    JSD(\beta)(P || Q)
+        = \beta * KLDiv(P || (\beta * P + (1 - \beta) * Q)) + (1 - \beta) * KLDiv(Q || (\beta * P + (1 - \beta) * Q))
+    .. note::
+    As all the other losses in PyTorch, this function expects the first argument,
+    :attr:`log_q`, to be the predictions, the output of the student model in log-space,
+    and the second, :attr:`log_p`, to be the observations, the output of the teacher model in log-space.
+    This differs from the standard mathematical notation :math:`JSD(P || Q)` where
+    :math:`P` denotes the teacher model and :math:`Q` denotes the student model.
+
+    Args:
+        beta (float): coefficient beta of generalized JSD in the open interval (0, 1). Default: `0.5`
+        ignore_index (int): The index to ignore in the target. Default: `-100`
+
+    Shape:
+        - Input: :math:`(BT, V)`, where B is batch size, T is sequence length, V is vocab size.
+        - Target: :math:`(BT, V)`, same shape as the input.
+        - shift_labels (Optional): :math:`(BT,)`
+        - Output: a scalar.
+
+    Examples:
+    ```python
+    >>> (B, T, V) = (2, 2, 5)
+    >>> jsd = LigerJSD(beta=0.1)
+    >>> # input should be a distribution in the log space
+    >>> input = torch.randn(B * T, V, requires_grad=True).log_softmax(dim=-1)
+    >>> target = torch.randn(B * T, V).log_softmax(dim=-1)
+    >>> output = jsd(input, target)
+    >>>
+    >>> # Example with labels for supervised fine-tuning (SFT) context
+    >>> # Assume logits and corresponding labels are given
+    >>> student_logits = torch.randn(B * T, V, requires_grad=True).log_softmax(dim=-1)
+    >>> teacher_logits = torch.randn(B * T, V).log_softmax(dim=-1)
+    >>> labels = torch.randint(0, V, (B * T,), torch.long)
+    >>> # Shift so that tokens < n predict n
+    >>> shift_student_logits = student_logits[..., :-1, :].contiguous()
+    >>> shift_teacher_logits = teacher_logits[..., :-1, :].contiguous()
+    >>> shift_labels = labels[..., 1:].contiguous()
+    >>> # Flatten tokens
+    >>> shift_student_logits = shift_student_logits.view(-1, V)
+    >>> shift_teacher_logits = shift_teacher_logits.view(-1, V)
+    >>> shift_labels = shift_labels.view(-1)
+    >>> # Calculate loss
+    >>> loss_fct = LigerJSD(beta=0.1)
+    >>> loss = loss_fct(shift_studetn_logits, shift_teacher_logits, shift_labels)
+
+    ```
+    """
+
+    def __init__(self, beta: float = 0.5, ignore_index: int = -100):
+        super().__init__()
+        assert (
+            beta > 0 and beta < 1
+        ), f"beta must be greater than 0 and less than 1. Got: {beta}"
+        self.beta = beta
+        self.ignore_index = ignore_index
+
+    def forward(
+        self,
+        log_q: torch.Tensor,
+        log_p: torch.Tensor,
+        shift_labels: Optional[torch.LongTensor] = None,
+    ):
+        return LigerJSDFunction.apply(
+            log_q, log_p, shift_labels, self.beta, self.ignore_index
+        )

liger_kernel/transformers/kl_div.py

@@ -4,10 +4,11 @@ from liger_kernel.ops.kl_div import LigerKLDivLossFunction
 
 
 class LigerKLDIVLoss(nn.KLDivLoss):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, eps: float = 1e-10, *args, **kwargs):
         super(LigerKLDIVLoss, self).__init__(*args, **kwargs)
+        self.eps = eps
 
     def forward(self, y_pred, y_true):
         return LigerKLDivLossFunction.apply(
-            y_pred, y_true, self.reduction, self.log_target
+            y_pred, y_true, self.reduction, self.log_target, self.eps
         )