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

{liger_kernel_nightly-0.5.3.dev20250221233257 → 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.dev20250221233257
+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.dev20250221233257 → 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.dev20250221233257"
+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.dev20250221233257 → 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.dev20250221233257 → 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.dev20250221233257
+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.dev20250221233257 → 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)