liger-kernel-nightly 0.5.3.dev20250221230243__tar.gz → 0.5.3.dev20250224175624__tar.gz

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/PKG-INFO

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

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/benchmark/scripts/benchmark_rope.py

@@ -1,7 +1,6 @@
 import torch
 import triton
 
-from test.utils import transformers_version_dispatch
 from transformers.models.llama.configuration_llama import LlamaConfig
 from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
 from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
@@ -14,6 +13,7 @@ from utils import run_benchmarks
 
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
 from liger_kernel.utils import infer_device
+from liger_kernel.utils import transformers_version_dispatch
 
 device = infer_device()
 

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.3.dev20250221230243"
+version = "0.5.3.dev20250224175624"
 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.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/src/liger_kernel/ops/layer_norm.py

@@ -57,13 +57,14 @@ def _layer_norm_forward_kernel(
     B_row = tl.load(B_ptr + col_offsets, mask=mask, other=0)
 
     mean = tl.sum(X_row, axis=0) / n_cols
-    var = tl.sum((X_row - mean) * (X_row - mean), axis=0) / n_cols
+    Xmm = tl.where(mask, X_row - mean, 0)
+    var = tl.sum(Xmm * Xmm, axis=0) / n_cols
     rstd = rsqrt(var + eps)
 
     tl.store(Mean_ptr, mean)
     tl.store(RSTD_ptr, rstd)
 
-    Y_row = (X_row - mean) * rstd * W_row + B_row
+    Y_row = Xmm * rstd * W_row + B_row
 
     tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
 
@@ -147,9 +148,11 @@ def layer_norm_forward(X, W, B, eps):
     Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
     RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
-    assert X.shape[1] == W.shape[0], (
-        f"Incompatible hidden size dimension between input tensor with shape[1] = {X.shape[1]} and weight tensor with shape[0] = {W.shape[0]}"
-    )
+    if X.shape[1] != W.shape[0]:
+        raise ValueError(
+            f"Incompatible dimensions: input feature size (X.shape[1]={X.shape[1]}) "
+            f"must match weight size (W.shape[0]={W.shape[0]})"
+        )
 
     _layer_norm_forward_kernel[(n_rows,)](
         Y,
@@ -190,11 +193,21 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
 
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
     if n_cols > BLOCK_SIZE:
-        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+        raise RuntimeError(
+            f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}. Consider using a smaller feature dimension."
+        )
 
     rows_per_program = math.ceil(n_rows / sm_count)
     grid = (sm_count,)
-    triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16
+    triton_dtype = (
+        tl.float32
+        if X.dtype == torch.float32
+        else tl.bfloat16
+        if X.dtype == torch.bfloat16
+        else tl.float16
+        if X.dtype == torch.float16
+        else tl.float32  # fallback to float32 for other types
+    )
     _layer_norm_backward_kernel[grid](
         X,
         W,

liger_kernel_nightly-0.5.3.dev20250224175624/src/liger_kernel/utils.py

@@ -0,0 +1,62 @@
+import torch
+
+
+def infer_device():
+    """
+    Get current device name based on available devices
+    """
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.xpu.is_available():
+        return "xpu"
+    elif torch.hip.is_available():
+        return "hip"
+    else:
+        return "cpu"
+
+
+def transformers_version_dispatch(
+    required_version: str,
+    before_fn,
+    after_fn,
+    before_args: tuple = (),
+    after_args: tuple = (),
+    before_kwargs: dict = None,
+    after_kwargs: dict = None,
+):
+    """
+    Dispatches to different functions based on package version comparison.
+
+    Args:
+        required_version: Version to compare against (e.g. "4.48.0")
+        before_fn: Function to call if package_version < required_version
+        after_fn: Function to call if package_version >= required_version
+        before_args: Positional arguments for before_fn
+        after_args: Positional arguments for after_fn
+        before_kwargs: Keyword arguments for before_fn
+        after_kwargs: Keyword arguments for after_fn
+
+    Returns:
+        Result from either before_fn or after_fn
+
+    Example:
+        >>> rotary_emb = transformers_version_dispatch(
+        ...     "4.48.0",
+        ...     LlamaRotaryEmbedding,
+        ...     LlamaRotaryEmbedding,
+        ...     before_args=(head_dim,),
+        ...     after_args=(LlamaConfig(head_dim=head_dim),),
+        ...     before_kwargs={'device': device},
+        ...     after_kwargs={'device': device}
+        ... )
+    """
+    from packaging import version
+    from transformers import __version__ as transformers_version
+
+    before_kwargs = before_kwargs or {}
+    after_kwargs = after_kwargs or {}
+
+    if version.parse(transformers_version) < version.parse(required_version):
+        return before_fn(*before_args, **before_kwargs)
+    else:
+        return after_fn(*after_args, **after_kwargs)

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/test/transformers/test_layer_norm.py

@@ -14,6 +14,8 @@ device = infer_device()
     [
         (2, 8, 64),
         (4, 16, 128),
+        (1, 1, 1023),  # Minimal batch/seq with near power-of-2 hidden
+        (3, 7, 256),  # Prime numbers for batch/seq
     ],
 )
 @pytest.mark.parametrize(
@@ -22,7 +24,15 @@ device = infer_device()
         (torch.float32, 1e-5, 1e-5),
     ],
 )
-def test_liger_layer_norm(batch_size, seq_len, hidden_size, dtype, atol, rtol):
+def test_liger_layer_norm(
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+) -> None:
+    """Test basic layer norm functionality against PyTorch implementation."""
     torch.manual_seed(0)
 
     x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device)
@@ -64,7 +74,15 @@ def test_liger_layer_norm(batch_size, seq_len, hidden_size, dtype, atol, rtol):
         (torch.float32, 1e-5, 1e-5),
     ],
 )
-def test_liger_layer_norm_functional(hidden_size, batch_size, seq_len, dtype, atol, rtol):
+def test_liger_layer_norm_functional(
+    hidden_size: int,
+    batch_size: int,
+    seq_len: int,
+    dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+) -> None:
+    """Test functional layer norm interface against autograd function."""
     torch.manual_seed(0)
 
     input = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device)
@@ -73,12 +91,10 @@ def test_liger_layer_norm_functional(hidden_size, batch_size, seq_len, dtype, at
     x2 = input.clone().requires_grad_(True)
 
     w = torch.randn(hidden_size, device=device, dtype=dtype)
-
     w1 = w.clone().requires_grad_(True)
     w2 = w.clone().requires_grad_(True)
 
     b = torch.randn(hidden_size, device=device, dtype=dtype)
-
     b1 = b.clone().requires_grad_(True)
     b2 = b.clone().requires_grad_(True)
 
@@ -88,7 +104,6 @@ def test_liger_layer_norm_functional(hidden_size, batch_size, seq_len, dtype, at
     assert torch.allclose(y1, y2, atol=atol, rtol=rtol)
 
     grad_output = torch.randn_like(y2)
-
     y1.backward(grad_output, retain_graph=True)
     y2.backward(grad_output, retain_graph=True)
 

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/test/transformers/test_rope.py

@@ -2,7 +2,6 @@ import pytest
 import torch
 
 from test.utils import supports_bfloat16
-from test.utils import transformers_version_dispatch
 from transformers.models.llama.configuration_llama import LlamaConfig
 from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
 from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
@@ -11,6 +10,7 @@ from liger_kernel.ops.rope import LigerRopeFunction
 from liger_kernel.transformers.functional import liger_rope
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
 from liger_kernel.utils import infer_device
+from liger_kernel.utils import transformers_version_dispatch
 
 device = infer_device()
 

{liger_kernel_nightly-0.5.3.dev20250221230243 → liger_kernel_nightly-0.5.3.dev20250224175624}/test/utils.py

@@ -214,53 +214,6 @@ def supports_bfloat16():
         return False
 
 
-def transformers_version_dispatch(
-    required_version: str,
-    before_fn,
-    after_fn,
-    before_args: tuple = (),
-    after_args: tuple = (),
-    before_kwargs: dict = None,
-    after_kwargs: dict = None,
-):
-    """
-    Dispatches to different functions based on package version comparison.
-
-    Args:
-        required_version: Version to compare against (e.g. "4.48.0")
-        before_fn: Function to call if package_version < required_version
-        after_fn: Function to call if package_version >= required_version
-        before_args: Positional arguments for before_fn
-        after_args: Positional arguments for after_fn
-        before_kwargs: Keyword arguments for before_fn
-        after_kwargs: Keyword arguments for after_fn
-
-    Returns:
-        Result from either before_fn or after_fn
-
-    Example:
-        >>> rotary_emb = transformers_version_dispatch(
-        ...     "4.48.0",
-        ...     LlamaRotaryEmbedding,
-        ...     LlamaRotaryEmbedding,
-        ...     before_args=(head_dim,),
-        ...     after_args=(LlamaConfig(head_dim=head_dim),),
-        ...     before_kwargs={'device': device},
-        ...     after_kwargs={'device': device}
-        ... )
-    """
-    from packaging import version
-    from transformers import __version__ as transformers_version
-
-    before_kwargs = before_kwargs or {}
-    after_kwargs = after_kwargs or {}
-
-    if version.parse(transformers_version) < version.parse(required_version):
-        return before_fn(*before_args, **before_kwargs)
-    else:
-        return after_fn(*after_args, **after_kwargs)
-
-
 def revert_liger_kernel_to_granite(model_config: MiniModelConfig):
     """
     Revert all Liger kernel patches applied to Granite.

liger_kernel_nightly-0.5.3.dev20250221230243/src/liger_kernel/utils.py

@@ -1,15 +0,0 @@
-import torch
-
-
-def infer_device():
-    """
-    Get current device name based on available devices
-    """
-    if torch.cuda.is_available():
-        return "cuda"
-    elif torch.xpu.is_available():
-        return "xpu"
-    elif torch.hip.is_available():
-        return "hip"
-    else:
-        return "cpu"