liger-kernel-nightly 0.5.5.dev20250322021112__py3-none-any.whl → 0.5.5.dev20250324181221__py3-none-any.whl

liger_kernel/ops/dyt.py

@@ -0,0 +1,225 @@
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import infer_device
+
+if compare_version("triton", operator.ge, "3.0.0"):
+    try:
+        # typical import path with dispatch available
+        from triton.language.extra.libdevice import tanh
+    except ModuleNotFoundError:
+        # for working with NGC containers
+        from triton.language.extra.cuda.libdevice import tanh
+else:
+    from triton.language.math import tanh
+
+
+@triton.jit
+def _dyt_fwd_kernel(
+    x_ptr,
+    x_row_stride,
+    alpha_ptr,
+    gamma_ptr,
+    beta_ptr,
+    y_ptr,
+    y_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Reference:
+    https://arxiv.org/abs/2503.10622
+
+    Shapes:
+        - x: (BT, C)
+        - alpha: (1)
+        - gamma: (C)
+        - beta: (C)
+    """
+    row_idx = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_cols
+
+    x_ptr += row_idx * x_row_stride
+    y_ptr += row_idx * y_row_stride
+
+    alpha = tl.load(alpha_ptr)
+    gamma = tl.load(gamma_ptr + offsets, mask=mask)
+    beta = tl.load(beta_ptr + offsets, mask=mask)
+    x = tl.load(x_ptr + offsets, mask=mask)
+    y = gamma * tanh((alpha * x).cast(tl.float32)) + beta
+    tl.store(y_ptr + offsets, y, mask=mask)
+
+
+@triton.jit
+def _dyt_bwd_kernel(
+    x_ptr,
+    x_row_stride,
+    dy_ptr,
+    dy_row_stride,
+    dx_ptr,
+    dx_row_stride,
+    alpha_ptr,
+    dalpha_ptr,
+    gamma_ptr,
+    dgamma_ptr,
+    dgamma_row_stride,
+    n_cols,
+    n_rows,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Reference:
+    https://arxiv.org/abs/2503.10622
+
+    Shapes:
+        - x: (BT, C)
+        - alpha: (1)
+        - gamma: (C)
+        - dx: (BT, C)
+        - dy: (BT, C)
+        - dgamma: (sm_count, C)
+        - dalpha: (sm_count,)
+    """
+    # d(gamma * tanh(alpha * x) + beta) / dx
+    # = gamma * (1 - tanh^2(alpha * x)) * alpha
+    # d(gamma * tanh(alpha * x) + beta) / dalpha
+    # = gamma * (1 - tanh^2(alpha * x)) * x
+    # d(gamma * tanh(alpha * x) + beta) / dgamma
+    # = tanh(alpha * x)
+    # d(gamma * tanh(alpha * x)) / dbeta = 1
+    pid = tl.program_id(0)
+
+    row_start = pid * ROWS_PER_PROGRAM
+    row_end = min((pid + 1) * ROWS_PER_PROGRAM, n_rows)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_cols
+
+    dalpha = 0.0
+    dgamma = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+
+    x_ptr += row_start * x_row_stride
+    dx_ptr += row_start * dx_row_stride
+    dy_ptr += row_start * dy_row_stride
+    alpha = tl.load(alpha_ptr)
+    gamma = tl.load(gamma_ptr + offsets, mask=mask, other=0.0)
+
+    for _ in tl.range(row_start, row_end):
+        dy = tl.load(dy_ptr + offsets, mask=mask, other=0.0)
+        x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+        tanh_ax = tanh((alpha * x).cast(tl.float32))
+        sech2_ax = 1 - tanh_ax * tanh_ax
+
+        dx = dy * gamma * sech2_ax * alpha
+        dalpha += tl.sum(dy * gamma * sech2_ax * x)
+        dgamma += dy * tanh_ax
+        tl.store(dx_ptr + offsets, dx, mask=mask)
+
+        dy_ptr += dy_row_stride
+        x_ptr += x_row_stride
+        dx_ptr += dx_row_stride
+
+    tl.store(dgamma_ptr + pid * dgamma_row_stride + offsets, dgamma, mask=mask)
+    tl.store(dalpha_ptr + pid, dalpha)
+
+    pass
+
+
+def liger_dyt_fwd(x, alpha, gamma, beta):
+    shape = x.shape
+    dim = shape[-1]
+    x = x.view(-1, dim)
+    n_rows, n_cols = x.shape
+    y = torch.empty_like(x)
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    _dyt_fwd_kernel[(n_rows,)](
+        x_ptr=x,
+        alpha_ptr=alpha,
+        gamma_ptr=gamma,
+        beta_ptr=beta,
+        y_ptr=y,
+        x_row_stride=x.stride(0),
+        y_row_stride=y.stride(0),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return y.view(*shape)
+
+
+def liger_dyt_bwd(dy, x, alpha, gamma):
+    shape = dy.shape
+    dtype = x.dtype
+    dim = shape[-1]
+    dy = dy.view(-1, dim)
+    x = x.view(-1, dim)
+    n_rows, n_cols = dy.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    sm_count = 1
+    device = infer_device()
+    if device == "cuda":
+        sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count
+    elif device == "xpu":
+        sm_count = torch.xpu.get_device_properties(x.device).gpu_subslice_count
+    if n_cols > BLOCK_SIZE:
+        raise RuntimeError(
+            f"Feature dimension {dim} exceeds maximum supported size of {BLOCK_SIZE}. Consider using a smaller feature dimension."
+        )
+
+    dx = torch.empty_like(x, dtype=torch.float32)
+    _dalpha = torch.empty((sm_count,), dtype=torch.float32, device=x.device)
+    _dgamma = torch.empty((sm_count, n_cols), dtype=torch.float32, device=x.device)
+
+    grid = (sm_count,)
+    rows_per_program = triton.cdiv(n_rows, sm_count)
+    _dyt_bwd_kernel[grid](
+        x_ptr=x,
+        x_row_stride=x.stride(0),
+        dy_ptr=dy,
+        dy_row_stride=dy.stride(0),
+        dx_ptr=dx,
+        dx_row_stride=dx.stride(0),
+        alpha_ptr=alpha,
+        dalpha_ptr=_dalpha,
+        gamma_ptr=gamma,
+        dgamma_ptr=_dgamma,
+        dgamma_row_stride=_dgamma.stride(0),
+        n_cols=n_cols,
+        n_rows=n_rows,
+        ROWS_PER_PROGRAM=rows_per_program,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    dalpha = _dalpha.sum(dim=0, keepdim=True).to(dtype)
+    dgamma = _dgamma.sum(dim=0).to(dtype)
+    dbeta = dy.sum(dim=0).to(dtype)
+    return dx.view(*shape), dalpha, dgamma, dbeta
+
+
+class LigerDyTFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, x, alpha, gamma, beta):
+        y = liger_dyt_fwd(x, alpha, gamma, beta)
+        ctx.save_for_backward(x, alpha, gamma)
+        return y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        x, alpha, gamma = ctx.saved_tensors
+        dx, dalpha, dgamma, dbeta = liger_dyt_bwd(
+            grad_output,
+            x,
+            alpha,
+            gamma,
+        )
+
+        return (dx, dalpha, dgamma, dbeta)

liger_kernel/transformers/__init__.py

@@ -1,5 +1,6 @@
 from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noqa: F401
+from liger_kernel.transformers.dyt import LigerDyT  # noqa: F401
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD  # noqa: F401
 from liger_kernel.transformers.geglu import LigerGEGLUMLP  # noqa: F401

liger_kernel/transformers/dyt.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.dyt import LigerDyTFunction
+
+
+class LigerDyT(nn.Module):
+    def __init__(self, hidden_size, init_alpha=0.5):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.init_alpha = init_alpha
+        self.alpha = nn.Parameter(torch.ones(1) * init_alpha)
+        self.gamma = nn.Parameter(torch.ones(hidden_size))
+        self.beta = nn.Parameter(torch.zeros(hidden_size))
+
+    def forward(self, x):
+        return LigerDyTFunction.apply(x, self.alpha, self.gamma, self.beta)
+
+    def extra_repr(self):
+        return f"{self.hidden_size}, init_alpha={self.init_alpha}"

liger_kernel/transformers/functional.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
+from liger_kernel.ops.dyt import LigerDyTFunction
 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
@@ -192,3 +193,7 @@ def liger_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 
 def liger_swiglu(a, b):
     return LigerSiLUMulFunction.apply(a, b)
+
+
+def liger_dyt(x, alpha, gamma, beta):
+    return LigerDyTFunction.apply(x, alpha, gamma, beta)

{liger_kernel_nightly-0.5.5.dev20250322021112.dist-info → liger_kernel_nightly-0.5.5.dev20250324181221.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.5.dev20250322021112
+Version: 0.5.5.dev20250324181221
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.5.5.dev20250322021112.dist-info → liger_kernel_nightly-0.5.5.dev20250324181221.dist-info}/RECORD

@@ -17,6 +17,7 @@ liger_kernel/chunked_loss/orpo_loss.py,sha256=nu9UYG16dcMw93lvHi4_hYs3Q0FK1KnlmM
 liger_kernel/chunked_loss/simpo_loss.py,sha256=fy2w8KbhMrBv7b1jdIeH3bBFxY52bPQPZb3KwBvmurM,5385
 liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 liger_kernel/ops/cross_entropy.py,sha256=yKKhN63I7r9NxJye4wTLBvvKAyrXQt6jf4nBo3lJyVg,18860
+liger_kernel/ops/dyt.py,sha256=YD1-buHz9VmIX838VKzLc-lm5CeUQ4LAskGDWBUMQHA,6187
 liger_kernel/ops/fused_linear_cross_entropy.py,sha256=1Y3Uk_TCSjqKgoG2eot1ptnWXJXXQESqGvOmqAW1gsM,10912
 liger_kernel/ops/fused_linear_jsd.py,sha256=Seshez2qaM6HiTQ8_HEqSwhaeVruNT1SvIM4ZrAPBEU,9602
 liger_kernel/ops/geglu.py,sha256=axGvCIvlBzuluoAIrWTsp2iZM4BFKNInkPov8YVvH9E,4126
@@ -32,10 +33,11 @@ liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
 liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
 liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
 liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
-liger_kernel/transformers/__init__.py,sha256=34zWr2C9sg0H5ok6l-pNB6eZQr5W1w8Xl_hMbQWzEUY,2460
+liger_kernel/transformers/__init__.py,sha256=eGCDpnvIBX7bhE_jGo5RRBipwT62WE_obzlniedNzt8,2525
 liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
 liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
-liger_kernel/transformers/functional.py,sha256=ShLD3eb--XKNtllznCrOYTbo4f-1KVwzi0KLMICdrn4,4942
+liger_kernel/transformers/dyt.py,sha256=QMqqc14pkE0WhpRZvapfnNAun-6C0C_tHExL2ZJuCUA,648
+liger_kernel/transformers/functional.py,sha256=4h9Pdx_iINBqfv2Zod_c27qOpYXDDwbdVgatQ9_XBmI,5089
 liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=09Rt7FZzLH42VOcIbQ4dlQd0o3Rlb4vk6fqiOQ7WTD8,1778
 liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
 liger_kernel/transformers/geglu.py,sha256=mrgqzIUVd6lN7fkDKLkw5YaESDxDtFgbot430WwPVOQ,1107
@@ -69,9 +71,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=pdekW7l6Qg_aqa5SYKYlSWUF8m3lkOFvFLcIMEHrz9s,8338
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/METADATA,sha256=t43xPpqj832O9qh6D1IfCqO06dNiO8VVJEUMvDzgs6U,22959
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.5.dev20250322021112.dist-info/RECORD,,
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/METADATA,sha256=NyKmdw6KevABFKKrqEdmIf8agklqARr8azTzS4RRx0k,22959
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.5.dev20250324181221.dist-info/RECORD,,