liger-kernel-nightly 0.5.9.dev20250517045713__tar.gz → 0.5.9.dev20250519011716__tar.gz

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.9.dev20250517045713
+Version: 0.5.9.dev20250519011716
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -322,7 +322,8 @@ loss.backward()
 | SwiGLU                          | `liger_kernel.transformers.LigerSwiGLUMLP`                  |
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
-| Fused Linear CrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 
 ### Alignment Kernels

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/README.md

@@ -274,7 +274,8 @@ loss.backward()
 | SwiGLU                          | `liger_kernel.transformers.LigerSwiGLUMLP`                  |
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
-| Fused Linear CrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 
 ### Alignment Kernels

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/benchmark/scripts/benchmark_dyt.py

@@ -22,17 +22,18 @@ def bench_speed_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
     from test.transformers.test_dyt import LigerDyT
     from test.transformers.test_dyt import TorchDyT
 
-    BT = input.x
+    hidden_size = 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"]
+    BT = extra_benchmark_config["BT"]
+    beta = extra_benchmark_config["beta"]
     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)
+    torch_dyt = TorchDyT(hidden_size=hidden_size, beta=beta).to(device)
+    torch_compile_dyt = torch.compile(TorchDyT(hidden_size=hidden_size, beta=beta).to(device))
+    triton_dyt = LigerDyT(hidden_size=hidden_size, beta=beta).to(device)
 
     x = torch.randn(x_shape, dtype=dtype, device=device)
     dy = torch.randn_like(x)
@@ -75,16 +76,17 @@ def bench_memory_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput
     from test.transformers.test_dyt import LigerDyT
     from test.transformers.test_dyt import TorchDyT
 
-    BT = input.x
+    hidden_size = input.x
     provider = input.kernel_provider
     extra_benchmark_config = input.extra_benchmark_config
-    hidden_size = extra_benchmark_config["hidden_size"]
+    BT = extra_benchmark_config["BT"]
+    beta = extra_benchmark_config["beta"]
     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)
+    torch_dyt = TorchDyT(hidden_size=hidden_size, beta=beta).to(device)
+    torch_compile_dyt = torch.compile(TorchDyT(hidden_size=hidden_size, beta=beta).to(device))
+    triton_dyt = LigerDyT(hidden_size=hidden_size, beta=beta).to(device)
 
     x = torch.randn(x_shape, dtype=dtype, device=device)
     dy = torch.randn_like(x)
@@ -113,27 +115,28 @@ def bench_memory_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput
 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,
-    )
+    for beta in [False, True]:
+        common_configs = {
+            "kernel_name": f"dyt_beta={beta}",
+            "x_name": "hidden_size",
+            "x_label": "hidden_size",
+            "x_values": [1024 * i for i in range(1, 17)],
+            "kernel_providers": ["liger", "torch", "torch_compile"],
+            "extra_benchmark_configs": [{"BT": 4096, "dtype": torch.bfloat16, "beta": beta}],
+            "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.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/docs/Low-Level-APIs.md

@@ -8,7 +8,9 @@
 | SwiGLU                          | `liger_kernel.transformers.LigerSwiGLUMLP`                  |
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
-| Fused Linear CrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
+
 
 ### RMS Norm
 
@@ -49,6 +51,12 @@ This kernel combines linear transformations with cross-entropy loss calculations
 !!! Example "Try it out"
     You can experiment as shown in this example [here](https://colab.research.google.com/drive/1Z2QtvaIiLm5MWOs7X6ZPS1MN3hcIJFbj?usp=sharing)
 
+### Sparsemax
+
+Sparsemax is a sparse alternative to softmax that produces sparse probability distributions. This kernel implements an efficient version of the sparsemax operation that can be used as a drop-in replacement for softmax in attention mechanisms or classification tasks.
+
+The implementation achieves significant speed improvements and memory savings compared to standard PyTorch implementations, particularly for large input tensors.
+
 ## Alignment Kernels
 
 | **Kernel**                      | **API**                                                     |

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.9.dev20250517045713"
+version = "0.5.9.dev20250519011716"
 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.9.dev20250519011716/src/liger_kernel/ops/dyt.py

@@ -0,0 +1,159 @@
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+from triton.language.extra.libdevice import tanh
+
+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.autotune([triton.Config({"BLOCK_N":bn}, num_stages=ns, num_warps=nw)
+#                   for bn in [1024, 2048, 4096]
+#                   for ns in [1,2,4]
+#                   for nw in [4, 8, 16, 32]
+#                   ],
+#                   key=['N'])
+@triton.jit
+def _dyt_fwd_kernel(X, Y, Alpha, Gamma, Beta, HAVE_BETA: tl.constexpr, N: tl.constexpr, BLOCK_N: tl.constexpr = 1024):
+    col = tl.cast(tl.program_id(0), tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = col < N
+    row_id = tl.cast(tl.program_id(1), tl.int64)
+
+    X += row_id * N
+    Y += row_id * N
+    alpha = tl.load(Alpha).to(tl.float32)
+
+    gamma = tl.load(Gamma + col, mask=mask, other=0.0).to(tl.float32)
+
+    x = tl.load(X + col, mask=mask, other=0.0).to(tl.float32)
+
+    tanh_x = tanh(alpha * x)
+    y = tanh_x * gamma
+    if HAVE_BETA:
+        beta = tl.load(Beta + col, mask=mask, other=0.0).to(tl.float32)
+        y += beta
+    tl.store(Y + col, y, mask=mask)
+
+
+# @triton.autotune([triton.Config({"BLOCK_N":bn}, num_stages=ns, num_warps=nw)
+#                   for bn in [1024, 2048, 4096]
+#                   for ns in [1,2,4]
+#                   for nw in [4, 8, 16]
+#                   ],
+#                   key=['N'])
+@triton.jit
+def _dyt_bwd_kernel(
+    DY, DX, DA, DG, DB, X, Alpha, Gamma, HAVE_BETA: tl.constexpr, M, N: tl.constexpr, BLOCK_N: tl.constexpr = 1024
+):
+    col = tl.cast(tl.program_id(0), tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = col < N
+    start_row_id = tl.cast(tl.program_id(1), tl.int64)
+
+    alpha = tl.load(Alpha).to(tl.float32)
+    da = 0.0
+    gamma = tl.load(Gamma + col, mask=mask, other=0.0).to(tl.float32)
+    dg = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    if HAVE_BETA:
+        db = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for row_id in range(start_row_id, M, tl.num_programs(1)):
+        x = tl.load(X + row_id * N + col, mask=mask, other=0.0).to(tl.float32)
+        dy = tl.load(DY + row_id * N + col, mask=mask, other=0.0).to(tl.float32)
+        tanh_x = tanh(alpha * x)
+        if HAVE_BETA:
+            db += dy
+        dg += dy * tanh_x
+        tmp = (1 - tanh_x * tanh_x) * dy * gamma
+        da += tl.sum(x * tmp, 0)
+        dx = alpha * tmp
+        tl.store(DX + row_id * N + col, dx, mask=mask)
+
+    tl.store(DG + start_row_id * N + col, dg, mask=mask)
+    if HAVE_BETA:
+        tl.store(DB + start_row_id * N + col, db, mask=mask)
+    tl.store(DA + start_row_id * tl.cdiv(N, 512) + tl.program_id(0), da)
+
+
+def liger_dyt_fwd(x, alpha, gamma, beta):
+    assert x.is_contiguous()
+    HAVE_BETA = True if beta is not None else False
+    input_shape = x.shape
+    x = x.view(-1, input_shape[-1])
+    M, N = x.shape
+
+    y = torch.empty_like(x)
+
+    if N >= 4096:
+        kwargs = {"BLOCK_N": min(triton.next_power_of_2(N), 2048), "num_warps": 4, "num_stages": 1}
+    else:
+        kwargs = {"BLOCK_N": min(triton.next_power_of_2(N), 1024), "num_warps": 4, "num_stages": 1}
+
+    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M)
+    _dyt_fwd_kernel[(grid)](
+        x,
+        y,
+        alpha,
+        gamma,
+        beta,
+        HAVE_BETA,
+        N,
+        **kwargs,
+    )
+    return y.view(input_shape)
+
+
+def liger_dyt_bwd(dy, x, alpha, gamma, beta):
+    assert dy.is_contiguous()
+    input_shape = x.shape
+    x = x.view(-1, input_shape[-1])
+    M, N = x.shape
+    HAVE_BETA = True if beta is not None else False
+
+    device = infer_device()
+    if device == "cuda":
+        NUM_SMS = torch.cuda.get_device_properties(x.device).multi_processor_count
+    elif device == "xpu":
+        NUM_SMS = torch.xpu.get_device_properties(x.device).gpu_subslice_count
+
+    da = torch.zeros(NUM_SMS, triton.cdiv(N, 512), dtype=torch.float32, device=x.device)
+    dg = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device)
+    db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None
+    dx = torch.empty_like(dy)
+
+    kwargs = {"BLOCK_N": min(triton.next_power_of_2(N), 1024), "num_warps": 8, "num_stages": 2}
+    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS)
+    _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N, **kwargs)
+    if HAVE_BETA:
+        db = db.sum(0).to(x.dtype)
+    dg = dg.sum(0).to(gamma.dtype)
+    da = da.sum().to(x.dtype).unsqueeze(0)
+    return dx.view(input_shape), da, dg, db
+
+
+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, beta)
+        return y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dy):
+        x, alpha, gamma, beta = ctx.saved_tensors
+        dx, dalpha, dgamma, dbeta = liger_dyt_bwd(dy, x, alpha, gamma, beta)
+        return dx, dalpha, dgamma, dbeta

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/src/liger_kernel/transformers/dyt.py

@@ -5,16 +5,18 @@ from liger_kernel.ops.dyt import LigerDyTFunction
 
 
 class LigerDyT(nn.Module):
-    def __init__(self, hidden_size, init_alpha=0.5):
+    def __init__(self, hidden_size, beta=True, 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))
+        self.beta = None
+        if beta:
+            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}"
+        return f"{self.hidden_size}, init_alpha={self.init_alpha}, beta={self.beta}"

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.9.dev20250517045713
+Version: 0.5.9.dev20250519011716
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -322,7 +322,8 @@ loss.backward()
 | SwiGLU                          | `liger_kernel.transformers.LigerSwiGLUMLP`                  |
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
-| Fused Linear CrossEntropy         | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 
 ### Alignment Kernels

{liger_kernel_nightly-0.5.9.dev20250517045713 → liger_kernel_nightly-0.5.9.dev20250519011716}/test/transformers/test_dyt.py

@@ -1,6 +1,5 @@
 import pytest
 import torch
-import torch.nn as nn
 
 from test.utils import assert_verbose_allclose
 from test.utils import infer_device
@@ -12,22 +11,37 @@ 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):
+# @torch.compile
+def torch_dyt_with_beta(x, alpha, gamma, beta):
+    return gamma * torch.tanh(x * alpha) + beta
+
+
+# @torch.compile
+def torch_dyt_without_beta(x, alpha, gamma):
+    return gamma * torch.tanh(x * alpha)
+
+
+class TorchDyT(torch.nn.Module):
+    def __init__(self, hidden_size, beta=True, 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))
+        self.alpha = torch.nn.Parameter(torch.ones(1) * init_alpha)
+        self.gamma = torch.nn.Parameter(torch.ones(hidden_size))
+        self.beta = None
+        if beta:
+            self.beta = torch.nn.Parameter(torch.zeros(hidden_size))
 
     def forward(self, x):
-        return self.gamma * torch.tanh(self.alpha * x) + self.beta
+        if self.beta is None:
+            return torch_dyt_without_beta(x, self.alpha, self.gamma)
+        return torch_dyt_with_beta(x, self.alpha, self.gamma, self.beta)
 
 
 set_seed(42)
 device = infer_device()
 
 
-@pytest.mark.parametrize("init_alpha", [0.5, 0.2, 1.0])
+@pytest.mark.parametrize("beta", [False, True])
+@pytest.mark.parametrize("init_alpha", [0.5])
 @pytest.mark.parametrize(
     "B, T, hidden_size",
     [
@@ -43,7 +57,7 @@ device = infer_device()
         (torch.float32, 1e-5, 1e-5),
     ],
 )
-def test_liger_dyt_correctness(B, T, hidden_size, init_alpha, dtype, atol, rtol):
+def test_liger_dyt_correctness(B, T, hidden_size, beta, init_alpha, dtype, atol, rtol):
     _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
 
     x1 = _input.clone().requires_grad_(True)
@@ -52,17 +66,19 @@ def test_liger_dyt_correctness(B, T, hidden_size, init_alpha, dtype, atol, rtol)
     # 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)
+    beta_data = 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 = TorchDyT(hidden_size=hidden_size, beta=beta, 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()
+    if beta:
+        torch_dyt.beta.data = beta_data.clone()
 
-    liger_dyt = LigerDyT(hidden_size=hidden_size, init_alpha=init_alpha).to(device).to(dtype)
+    liger_dyt = LigerDyT(hidden_size=hidden_size, beta=beta, 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()
+    if beta:
+        liger_dyt.beta.data = beta_data.clone()
 
     torch_output = torch_dyt(x1)
     liger_output = liger_dyt(x2)
@@ -76,9 +92,11 @@ def test_liger_dyt_correctness(B, T, hidden_size, init_alpha, dtype, atol, rtol)
     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)
+    if beta:
+        assert_verbose_allclose(torch_dyt.beta.grad, liger_dyt.beta.grad, rtol=rtol, atol=atol)
 
 
+@pytest.mark.parametrize("beta", [False, True])
 @pytest.mark.parametrize(
     "B, T, hidden_size",
     [
@@ -102,7 +120,7 @@ def test_liger_dyt_correctness(B, T, hidden_size, init_alpha, dtype, atol, rtol)
         ),
     ],
 )
-def test_liger_dyt_functional(B, T, hidden_size, dtype, atol, rtol):
+def test_liger_dyt_functional(B, T, hidden_size, beta, dtype, atol, rtol):
     _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
 
     x1 = _input.clone().requires_grad_(True)
@@ -111,15 +129,16 @@ def test_liger_dyt_functional(B, T, hidden_size, dtype, atol, rtol):
     # 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)
+    beta_data = 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)
+    beta1 = beta_data.clone().requires_grad_(True) if beta else None
 
     alpha2 = alpha.clone().requires_grad_(True)
     gamma2 = gamma.clone().requires_grad_(True)
-    beta2 = beta.clone().requires_grad_(True)
+
+    beta2 = beta_data.clone().requires_grad_(True) if beta else None
 
     output1 = liger_dyt(x1, alpha=alpha1, gamma=gamma1, beta=beta1)
     output2 = LigerDyTFunction.apply(x2, alpha2, gamma2, beta2)
@@ -133,4 +152,5 @@ def test_liger_dyt_functional(B, T, hidden_size, dtype, atol, rtol):
     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)
+    if beta:
+        assert_verbose_allclose(beta1.grad, beta2.grad, rtol=rtol, atol=atol)