liger-kernel-nightly 0.5.5.dev20250320214749__tar.gz → 0.5.5.dev20250324181221__tar.gz

{liger_kernel_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.5.dev20250320214749
+Version: 0.5.5.dev20250324181221
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -176,7 +176,7 @@ y = orpo_loss(lm_head.weight, x, target)
 - **Exact:** Computation is exact—no approximations! Both forward and backward passes are implemented with rigorous unit tests and undergo convergence testing against training runs without Liger Kernel to ensure accuracy.
 - **Lightweight:** Liger Kernel has 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, DeepSpeed, DDP, etc.).
-- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
+- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift), [oumi](https://github.com/oumi-ai/oumi/tree/main)
 
 ## Installation
 
@@ -386,8 +386,8 @@ loss.backward()
 ## Contact
 
 - For issues, create a Github ticket in this repository
-- For open discussion, join [our discord channel](https://discord.gg/gpumode)
-- For formal collaboration, send an email to yannchen@linkedin.com
+- For open discussion, join [our discord channel on GPUMode](https://discord.com/channels/1189498204333543425/1275130785933951039)
+- For formal collaboration, send an email to yannchen@linkedin.com and hning@linkedin.com
 
 ## Cite this work
 

{liger_kernel_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/README.md

@@ -128,7 +128,7 @@ y = orpo_loss(lm_head.weight, x, target)
 - **Exact:** Computation is exact—no approximations! Both forward and backward passes are implemented with rigorous unit tests and undergo convergence testing against training runs without Liger Kernel to ensure accuracy.
 - **Lightweight:** Liger Kernel has 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, DeepSpeed, DDP, etc.).
-- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
+- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift), [oumi](https://github.com/oumi-ai/oumi/tree/main)
 
 ## Installation
 
@@ -338,8 +338,8 @@ loss.backward()
 ## Contact
 
 - For issues, create a Github ticket in this repository
-- For open discussion, join [our discord channel](https://discord.gg/gpumode)
-- For formal collaboration, send an email to yannchen@linkedin.com
+- For open discussion, join [our discord channel on GPUMode](https://discord.com/channels/1189498204333543425/1275130785933951039)
+- For formal collaboration, send an email to yannchen@linkedin.com and hning@linkedin.com
 
 ## Cite this work
 

{liger_kernel_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/benchmark/benchmarks_visualizer.py

@@ -8,8 +8,8 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn as sns
 
-DATA_PATH = "data/all_benchmark_data.csv"
-VISUALIZATIONS_PATH = "visualizations/"
+DATA_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "data/all_benchmark_data.csv"))
+VISUALIZATIONS_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "visualizations/"))
 
 
 @dataclass

liger_kernel_nightly-0.5.5.dev20250324181221/benchmark/scripts/benchmark_dyt.py

@@ -0,0 +1,139 @@
+import os
+import sys
+
+import torch
+import triton
+
+from utils import QUANTILES
+from utils import SingleBenchmarkRunInput
+from utils import SingleBenchmarkRunOutput
+from utils import _test_memory
+from utils import parse_benchmark_script_args
+from utils import run_benchmarks
+
+from liger_kernel.utils import infer_device
+
+device = infer_device()
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+
+def bench_speed_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    from test.transformers.test_dyt import LigerDyT
+    from test.transformers.test_dyt import TorchDyT
+
+    BT = input.x
+    provider = input.kernel_provider
+    mode = input.kernel_operation_mode
+    extra_benchmark_config = input.extra_benchmark_config
+    hidden_size = extra_benchmark_config["hidden_size"]
+    dtype = extra_benchmark_config["dtype"]
+
+    x_shape = (BT, hidden_size)
+    torch_dyt = TorchDyT(hidden_size=hidden_size).to(device)
+    torch_compile_dyt = torch.compile(TorchDyT(hidden_size=hidden_size).to(device))
+    triton_dyt = LigerDyT(hidden_size=hidden_size).to(device)
+
+    x = torch.randn(x_shape, dtype=dtype, device=device)
+    dy = torch.randn_like(x)
+    x.requires_grad_(True)
+
+    def fwd():
+        if provider == "liger":
+            return triton_dyt(x)
+        elif provider == "torch":
+            return torch_dyt(x)
+        elif provider == "torch_compile":
+            return torch_compile_dyt(x)
+
+    if mode == "forward":
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, grad_to_none=[x], rep=500)
+    elif mode == "backward":
+        y = fwd()
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(
+            lambda: y.backward(dy, retain_graph=True),
+            quantiles=QUANTILES,
+            grad_to_none=[x],
+            rep=500,
+        )
+    elif mode == "full":
+
+        def full():
+            y = fwd()
+            y.backward(dy)
+
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, grad_to_none=[x], rep=500)
+
+    return SingleBenchmarkRunOutput(
+        y_20=ms_20,
+        y_50=ms_50,
+        y_80=ms_80,
+    )
+
+
+def bench_memory_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    from test.transformers.test_dyt import LigerDyT
+    from test.transformers.test_dyt import TorchDyT
+
+    BT = input.x
+    provider = input.kernel_provider
+    extra_benchmark_config = input.extra_benchmark_config
+    hidden_size = extra_benchmark_config["hidden_size"]
+    dtype = extra_benchmark_config["dtype"]
+
+    x_shape = (BT, hidden_size)
+    torch_dyt = TorchDyT(hidden_size=hidden_size).to(device)
+    torch_compile_dyt = torch.compile(TorchDyT(hidden_size=hidden_size).to(device))
+    triton_dyt = LigerDyT(hidden_size=hidden_size).to(device)
+
+    x = torch.randn(x_shape, dtype=dtype, device=device)
+    dy = torch.randn_like(x)
+    x.requires_grad_(True)
+
+    def fwd():
+        if provider == "liger":
+            return triton_dyt(x)
+        elif provider == "torch":
+            return torch_dyt(x)
+        elif provider == "torch_compile":
+            return torch_compile_dyt(x)
+
+    def full():
+        y = fwd()
+        y.backward(dy, retain_graph=True)
+
+    mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES)
+    return SingleBenchmarkRunOutput(
+        y_20=mem_20,
+        y_50=mem_50,
+        y_80=mem_80,
+    )
+
+
+if __name__ == "__main__":
+    args = parse_benchmark_script_args()
+
+    common_configs = {
+        "kernel_name": "dyt",
+        "x_name": "BT",
+        "x_label": "batch_size * seq_len",
+        "x_values": [2**i for i in range(10, 15)],
+        "kernel_providers": ["liger", "torch", "torch_compile"],
+        "extra_benchmark_configs": [{"hidden_size": 4096, "dtype": torch.float32}],
+        "overwrite": args.overwrite,
+    }
+
+    run_benchmarks(
+        bench_test_fn=bench_speed_dyt,
+        kernel_operation_modes=["forward", "backward", "full"],
+        metric_name="speed",
+        metric_unit="ms",
+        **common_configs,
+    )
+    run_benchmarks(
+        bench_test_fn=bench_memory_dyt,
+        kernel_operation_modes=["full"],
+        metric_name="memory",
+        metric_unit="MB",
+        **common_configs,
+    )

{liger_kernel_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.5.dev20250320214749"
+version = "0.5.5.dev20250324181221"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

liger_kernel_nightly-0.5.5.dev20250324181221/src/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_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/src/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_nightly-0.5.5.dev20250324181221/src/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_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/src/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.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.5.dev20250320214749
+Version: 0.5.5.dev20250324181221
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -176,7 +176,7 @@ y = orpo_loss(lm_head.weight, x, target)
 - **Exact:** Computation is exact—no approximations! Both forward and backward passes are implemented with rigorous unit tests and undergo convergence testing against training runs without Liger Kernel to ensure accuracy.
 - **Lightweight:** Liger Kernel has 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, DeepSpeed, DDP, etc.).
-- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
+- **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift), [oumi](https://github.com/oumi-ai/oumi/tree/main)
 
 ## Installation
 
@@ -386,8 +386,8 @@ loss.backward()
 ## Contact
 
 - For issues, create a Github ticket in this repository
-- For open discussion, join [our discord channel](https://discord.gg/gpumode)
-- For formal collaboration, send an email to yannchen@linkedin.com
+- For open discussion, join [our discord channel on GPUMode](https://discord.com/channels/1189498204333543425/1275130785933951039)
+- For formal collaboration, send an email to yannchen@linkedin.com and hning@linkedin.com
 
 ## Cite this work
 

{liger_kernel_nightly-0.5.5.dev20250320214749 → liger_kernel_nightly-0.5.5.dev20250324181221}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -24,6 +24,7 @@ benchmark/scripts/benchmark_cpo_loss.py
 benchmark/scripts/benchmark_cross_entropy.py
 benchmark/scripts/benchmark_distill_jsd_loss.py
 benchmark/scripts/benchmark_dpo_loss.py
+benchmark/scripts/benchmark_dyt.py
 benchmark/scripts/benchmark_embedding.py
 benchmark/scripts/benchmark_fused_linear_cross_entropy.py
 benchmark/scripts/benchmark_fused_linear_jsd.py
@@ -121,6 +122,7 @@ src/liger_kernel/chunked_loss/orpo_loss.py
 src/liger_kernel/chunked_loss/simpo_loss.py
 src/liger_kernel/ops/__init__.py
 src/liger_kernel/ops/cross_entropy.py
+src/liger_kernel/ops/dyt.py
 src/liger_kernel/ops/fused_linear_cross_entropy.py
 src/liger_kernel/ops/fused_linear_jsd.py
 src/liger_kernel/ops/geglu.py
@@ -139,6 +141,7 @@ src/liger_kernel/ops/experimental/mm_int8int2.py
 src/liger_kernel/transformers/__init__.py
 src/liger_kernel/transformers/auto_model.py
 src/liger_kernel/transformers/cross_entropy.py
+src/liger_kernel/transformers/dyt.py
 src/liger_kernel/transformers/functional.py
 src/liger_kernel/transformers/fused_linear_cross_entropy.py
 src/liger_kernel/transformers/fused_linear_jsd.py
@@ -209,6 +212,7 @@ test/resources/tiny_shakespeare_tokenized/dataset_info.json
 test/resources/tiny_shakespeare_tokenized/state.json
 test/transformers/test_auto_model.py
 test/transformers/test_cross_entropy.py
+test/transformers/test_dyt.py
 test/transformers/test_embedding.py
 test/transformers/test_flex_attention.py
 test/transformers/test_fused_linear_cross_entropy.py

liger_kernel_nightly-0.5.5.dev20250324181221/test/transformers/test_dyt.py

@@ -0,0 +1,136 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from test.utils import assert_verbose_allclose
+from test.utils import infer_device
+from test.utils import set_seed
+from test.utils import supports_bfloat16
+
+from liger_kernel.ops.dyt import LigerDyTFunction
+from liger_kernel.transformers.dyt import LigerDyT
+from liger_kernel.transformers.functional import liger_dyt
+
+
+class TorchDyT(nn.Module):
+    def __init__(self, hidden_size, init_alpha=0.5):
+        super().__init__()
+        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 self.gamma * torch.tanh(self.alpha * x) + self.beta
+
+
+set_seed(42)
+device = infer_device()
+
+
+@pytest.mark.parametrize("init_alpha", [0.5, 0.2, 1.0])
+@pytest.mark.parametrize(
+    "B, T, hidden_size",
+    [
+        (2, 8, 4096),
+        (4, 16, 2048),
+        (1, 1, 1023),  # Minimal batch/seq with near power-of-2 hidden
+        (3, 7, 256),  # Prime numbers for batch/seq
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-5, 1e-5),
+    ],
+)
+def test_liger_dyt_correctness(B, T, hidden_size, init_alpha, dtype, atol, rtol):
+    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
+
+    x1 = _input.clone().requires_grad_(True)
+    x2 = _input.clone().requires_grad_(True)
+
+    # initialize weights
+    alpha = torch.randn(1, device=device, dtype=dtype)
+    gamma = torch.randn(hidden_size, device=device, dtype=dtype)
+    beta = torch.randn(hidden_size, device=device, dtype=dtype)
+
+    torch_dyt = TorchDyT(hidden_size=hidden_size, init_alpha=init_alpha).to(device).to(dtype)
+    torch_dyt.alpha.data = alpha.clone()
+    torch_dyt.gamma.data = gamma.clone()
+    torch_dyt.beta.data = beta.clone()
+
+    liger_dyt = LigerDyT(hidden_size=hidden_size, init_alpha=init_alpha).to(device).to(dtype)
+    liger_dyt.alpha.data = alpha.clone()
+    liger_dyt.gamma.data = gamma.clone()
+    liger_dyt.beta.data = beta.clone()
+
+    torch_output = torch_dyt(x1)
+    liger_output = liger_dyt(x2)
+
+    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol)
+
+    grad_output = torch.randn_like(_input)
+    torch_output.backward(grad_output)
+    liger_output.backward(grad_output)
+
+    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(torch_dyt.alpha.grad, liger_dyt.alpha.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(torch_dyt.gamma.grad, liger_dyt.gamma.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(torch_dyt.beta.grad, liger_dyt.beta.grad, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize(
+    "B, T, hidden_size",
+    [
+        (2, 8, 4096),
+        (4, 16, 2048),
+        (1, 1, 1023),  # Minimal batch/seq with near power-of-2 hidden
+        (3, 7, 256),  # Prime numbers for batch/seq
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        # atol is for small values: they have more difference, so set atol higher
+        # rtol is for larger values: they are very close, so set rtol lower
+        (torch.float32, 1e-5, 1e-5),
+        pytest.param(
+            torch.bfloat16,
+            1e-8,
+            5e-2,
+            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+        ),
+    ],
+)
+def test_liger_dyt_functional(B, T, hidden_size, dtype, atol, rtol):
+    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
+
+    x1 = _input.clone().requires_grad_(True)
+    x2 = _input.clone().requires_grad_(True)
+
+    # initialize weights
+    alpha = torch.randn(1, device=device, dtype=dtype)
+    gamma = torch.randn(hidden_size, device=device, dtype=dtype)
+    beta = torch.randn(hidden_size, device=device, dtype=dtype)
+
+    alpha1 = alpha.clone().requires_grad_(True)
+    gamma1 = gamma.clone().requires_grad_(True)
+    beta1 = beta.clone().requires_grad_(True)
+
+    alpha2 = alpha.clone().requires_grad_(True)
+    gamma2 = gamma.clone().requires_grad_(True)
+    beta2 = beta.clone().requires_grad_(True)
+
+    output1 = liger_dyt(x1, alpha=alpha1, gamma=gamma1, beta=beta1)
+    output2 = LigerDyTFunction.apply(x2, alpha2, gamma2, beta2)
+
+    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)
+
+    grad_output = torch.randn_like(_input)
+    output1.backward(grad_output)
+    output2.backward(grad_output)
+
+    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(alpha1.grad, alpha2.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(gamma1.grad, gamma2.grad, rtol=rtol, atol=atol)
+    assert_verbose_allclose(beta1.grad, beta2.grad, rtol=rtol, atol=atol)