liger-kernel-nightly 0.5.8.dev20250429220905__tar.gz → 0.5.8.dev20250429233059__tar.gz

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.8.dev20250429220905"
+version = "0.5.8.dev20250429233059"
 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.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/src/liger_kernel/ops/cross_entropy.py

@@ -351,7 +351,10 @@ def cross_entropy_backward(_input, grad_output):
     # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
     if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
         pass
-
+    # If reduction is 'none'
+    elif grad_output.ndim > 0:
+        _input = _input * grad_output.unsqueeze(dim=1)
+    # If reduction is ['mean', 'sum'], grad_output is just a scalar
     # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
     # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
     else:

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -143,9 +143,10 @@ def fused_linear_cross_entropy_forward(
                 alpha=1.0,
             )
 
-    if reduction == "none":
-        loss = loss_1d
-        z_loss = z_loss_1d if return_z_loss else None
+    # Need extra calculations for backward if reduction=='none'. Not supporting reduction='none' now.
+    # if reduction == "none":
+    #     loss = loss_1d
+    #     z_loss = z_loss_1d if return_z_loss else None
 
     else:
         loss = torch.sum(loss_1d)

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/src/liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -23,8 +23,7 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         assert reduction in {
             "mean",
             "sum",
-            "none",
-        }, f"reduction must be one of 'mean', 'sum', or 'none'. Got: {reduction}"
+        }, f"reduction must be 'mean' or 'sum'. Got: {reduction}"
         assert softcap is None or softcap > 0, f"softcap must greater than 0.0 or None. Got: {softcap}"
         self.ce_weight = ce_weight
         self.ignore_index = ignore_index

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/test/transformers/test_cross_entropy.py

@@ -400,8 +400,58 @@ def _test_correctness_not_last_layer_once(target_ce, B, T, V, reduction, scalar,
     loss1 = output * 3
     loss2 = output2 * 3
 
-    loss1.backward(gradient=torch.ones_like(output))
-    loss2.backward(gradient=torch.ones_like(output))
+    grad_output = torch.rand_like(output)
+    loss1.backward(gradient=grad_output)
+    loss2.backward(gradient=grad_output)
+    assert torch.allclose(_input.grad, _input2.grad, atol=atol, rtol=rtol)
+
+
+def _test_correctness_not_last_layer_with_other_params_once(
+    target_ce,
+    B,
+    T,
+    V,
+    reduction,
+    ignore_index,
+    lse_square_scale,
+    label_smoothing,
+    softcap,
+    scalar,
+    dtype,
+    atol,
+    rtol,
+):
+    torch_ce = CrossEntropyWithZLoss(
+        reduction=reduction,
+        ignore_index=ignore_index,
+        lse_square_scale=lse_square_scale,
+        label_smoothing=label_smoothing,
+    )
+
+    _tensor = torch.randn(B * T, V, device=device, dtype=dtype) * scalar
+    _input = _tensor.detach().clone().requires_grad_(True)
+    _input2 = _tensor.detach().clone().requires_grad_(True)
+
+    target = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
+    # Assign some random number of elements as ignore_index
+    num_elements_to_assign = torch.randint(
+        1, B * T // 2, (1,)
+    ).item()  # Random number of elements to set to ignore_index
+    indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]  # Randomly select indices
+    target[indices_to_assign] = ignore_index
+
+    # upcasting to match liger's casting strategy
+    # and downcasting to original dtype
+    output = torch_ce(softcap * torch.tanh(_input.to(torch.float32) / softcap), target).to(dtype)
+    output2 = target_ce(_input2, target)
+    assert torch.allclose(output, output2, atol=atol, rtol=rtol)
+
+    loss1 = output * 3
+    loss2 = output2 * 3
+
+    grad_output = torch.rand_like(output)
+    loss1.backward(gradient=grad_output)
+    loss2.backward(gradient=grad_output)
     assert torch.allclose(_input.grad, _input2.grad, atol=atol, rtol=rtol)
 
 
@@ -864,6 +914,62 @@ def test_correctness_not_last_layer(B, T, V, reduction, scalar, dtype, atol, rto
     _test_correctness_not_last_layer_once(liger_ce, B, T, V, reduction, scalar, dtype, atol, rtol)
 
 
+@pytest.mark.parametrize(
+    "B, T, V",
+    [
+        (2, 1024, 32000),  # llama2, mistral
+        # # weird shapes
+        (3, 423, 32000),
+    ],
+)
+@pytest.mark.parametrize(
+    "ignore_index, lse_square_scale, label_smoothing, softcap",
+    [
+        (-100, 1e-4, 0.1, 30.0),
+        (42, 1e-5, 0.2, 40.0),
+    ],
+)
+@pytest.mark.parametrize("reduction", ["sum", "mean"])
+@pytest.mark.parametrize(
+    "scalar, dtype, atol, rtol",
+    [
+        pytest.param(
+            1.0,
+            torch.bfloat16,
+            1e-8,
+            5e-2,
+            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+        ),
+        (1.0, torch.float32, 1e-8, 1e-5),
+    ],
+)
+def test_correctness_not_last_layer_with_other_params(
+    B, T, V, reduction, ignore_index, lse_square_scale, label_smoothing, softcap, scalar, dtype, atol, rtol
+):
+    liger_ce = LigerCrossEntropyLoss(
+        reduction=reduction,
+        ignore_index=ignore_index,
+        lse_square_scale=lse_square_scale,
+        label_smoothing=label_smoothing,
+        softcap=softcap,
+    )
+    _test_correctness_not_last_layer_with_other_params_once(
+        liger_ce,
+        B,
+        T,
+        V,
+        reduction,
+        ignore_index,
+        lse_square_scale,
+        label_smoothing,
+        softcap,
+        scalar,
+        dtype,
+        atol,
+        rtol,
+    )
+
+
 def test_float32_internal():
     """
     This test validates that the internal softmax calculations occur in float32,

{liger_kernel_nightly-0.5.8.dev20250429220905 → liger_kernel_nightly-0.5.8.dev20250429233059}/test/transformers/test_fused_linear_cross_entropy.py

@@ -116,8 +116,6 @@ class LigerLMHeadCE(torch.nn.Module):
         ("mean", 1.0, torch.float32, 1e-5, 5e-4),
         ("sum", 1.0, torch.bfloat16, 5e-0, 5e1),
         ("sum", 1.0, torch.float32, 1e-3, 5e-2),
-        ("none", 1.0, torch.bfloat16, 5e-0, 5e1),
-        ("none", 1.0, torch.float32, 1e-3, 5e-2),
     ],
 )
 @pytest.mark.parametrize("bias", [True, False])
@@ -207,8 +205,9 @@ def test_correctness(
     if return_z_loss:
         assert_verbose_allclose(z_output1, z_output2, atol=atol, rtol=rtol)
 
-    output1.backward(gradient=torch.ones_like(output1))
-    output2.backward(gradient=torch.ones_like(output2))
+    grad_output = torch.ones_like(output1)
+    output1.backward(gradient=grad_output)
+    output2.backward(gradient=grad_output)
 
     assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)