liger-kernel-nightly 0.5.4.dev20250225015144__tar.gz → 0.5.4.dev20250225020243__tar.gz

{liger_kernel_nightly-0.5.4.dev20250225015144 → liger_kernel_nightly-0.5.4.dev20250225020243}/PKG-INFO

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

{liger_kernel_nightly-0.5.4.dev20250225015144 → liger_kernel_nightly-0.5.4.dev20250225020243}/pyproject.toml

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

@@ -285,6 +285,10 @@ def cross_entropy_forward(
 
     target_mask = target != ignore_index
     n_non_ignore = target_mask.sum().item()
+    assert (target * target_mask).max() < _input.shape[-1], (
+        f"Target {target.max()} is out of bounds. Expected < {_input.shape[-1]}"
+    )
+    assert (target * target_mask).min() >= 0, f"Target {target.min()} is out of bounds. Expected >= 0"
     sum_non_ignore_weight = n_non_ignore
     weight_sum = 0.0
     if weight is not None:

{liger_kernel_nightly-0.5.4.dev20250225015144 → liger_kernel_nightly-0.5.4.dev20250225020243}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.4.dev20250225015144 → liger_kernel_nightly-0.5.4.dev20250225020243}/test/transformers/test_cross_entropy.py

@@ -290,6 +290,32 @@ def _test_correctness_with_z_loss_with_other_params_once(
     assert_verbose_allclose(_input.grad, _input2.grad, atol=atol, rtol=rtol)
 
 
+def _test_correctness_with_out_of_bounds_target_once(target_ce, B, T, V, ignore_index):
+    torch.manual_seed(0)
+
+    _tensor = torch.randn(B * T, V, device=device, dtype=torch.bfloat16)
+    _input = _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
+
+    # Assign out of bounds target
+    num_out_of_bounds = torch.randint(1, B * T // 2, (1,)).item()
+    indices_to_assign = torch.randperm(B * T)[:num_out_of_bounds]  # Randomly select indices
+    target[indices_to_assign] = torch.randint(V, 2 * V, (num_out_of_bounds,)).to(device)
+
+    try:
+        _ = target_ce(_input, target)
+        assert False, "Should have thrown an error"
+    except AssertionError as e:
+        assert "out of bounds" in str(e)
+
+
 def _test_correctness_with_weight_once(target_ce, B, T, V, reduction, weight, scalar, dtype, atol, rtol):
     torch.manual_seed(0)
     torch_ce = CrossEntropyLoss(weight=weight, reduction=reduction)
@@ -916,3 +942,16 @@ def test_float32_internal():
 
     torch.allclose(X_bf16, X_fp32.bfloat16())
     torch.allclose(loss_bf16, loss_fp32)
+
+
+@pytest.mark.parametrize(
+    "B, T, V, ignore_index",
+    [
+        (2, 4096, 32000, 2),
+        # weird shapes
+        (3, 423, 32000, -123),
+    ],
+)
+def test_correctness_with_out_of_bounds_target_once(B, T, V, ignore_index):
+    liger_ce = LigerCrossEntropyLoss(ignore_index=ignore_index)
+    _test_correctness_with_out_of_bounds_target_once(liger_ce, B, T, V, ignore_index)