liger-kernel-nightly 0.5.2.dev20241223042135__tar.gz → 0.5.2.dev20241229035411__tar.gz

{liger_kernel_nightly-0.5.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/PKG-INFO

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

{liger_kernel_nightly-0.5.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.2.dev20241223042135"
+version = "0.5.2.dev20241229035411"
 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.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/src/liger_kernel/chunked_loss/cpo_loss.py

@@ -33,7 +33,6 @@ class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
         loss = (-F.logsigmoid(logits) * (1 - label_smoothing) - F.logsigmoid(-logits) * label_smoothing).sum() / (
             full_target.shape[0] // 2
         )
-
         return loss
 
     @staticmethod

{liger_kernel_nightly-0.5.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -59,6 +59,7 @@ def fused_linear_cross_entropy_forward(
         logits_chunk = _input_chunk @ weight.t()  # chunk_size x V
         if bias is not None:
             logits_chunk = logits_chunk + bias
+
         target_chunk = target[start_idx:end_idx]  # chunk_size,
 
         n_rows = logits_chunk.shape[0]
@@ -112,7 +113,9 @@ def fused_linear_cross_entropy_forward(
         if grad_weight is not None:
             torch.addmm(
                 input=grad_weight,
-                mat1=logits_chunk.t(),
+                mat1=logits_chunk.t().to(
+                    _input_chunk.dtype
+                ),  # In an autocast scenario without bias, differing logits_chunk data types will cause an addmm operation error.
                 mat2=_input_chunk,
                 out=grad_weight,
                 alpha=alpha,
@@ -127,13 +130,16 @@ def fused_linear_cross_entropy_forward(
                 alpha=alpha,
             )
 
-    loss = torch.sum(loss_1d)
+    if reduction == "none":
+        loss = loss_1d
+    else:
+        loss = torch.sum(loss_1d)
     return loss, grad_input, grad_weight, grad_bias
 
 
 def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias):
     # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
-    if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
+    if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
         # 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.
         BT, H = grad_input.shape

{liger_kernel_nightly-0.5.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.2.dev20241223042135 → liger_kernel_nightly-0.5.2.dev20241229035411}/test/transformers/test_fused_linear_cross_entropy.py

@@ -108,6 +108,8 @@ 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])
@@ -185,8 +187,8 @@ def test_correctness(
 
     assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
 
-    output1.backward()
-    output2.backward()
+    output1.backward(gradient=torch.ones_like(output1))
+    output2.backward(gradient=torch.ones_like(output2))
 
     assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)
 
@@ -258,18 +260,20 @@ def test_correctness_functional(B, T, H, V, scalar, dtype, bias, atol, rtol):
     ],
 )
 @pytest.mark.parametrize(
-    "cast_dtype, atol, rtol",
+    "bias, cast_dtype, atol, rtol",
     [
-        (torch.bfloat16, 5e-3, 5e-2),
-        (torch.float16, 5e-3, 5e-2),
+        (True, torch.bfloat16, 5e-3, 5e-2),
+        (True, torch.float16, 5e-3, 5e-2),
+        (False, torch.bfloat16, 5e-3, 5e-2),
+        (False, torch.float16, 5e-3, 5e-2),
     ],
 )
-def test_amp(B, T, H, V, cast_dtype, atol, rtol):
+def test_amp(B, T, H, V, bias, cast_dtype, atol, rtol):
     dtype = torch.float32
     torch_lm_head_ce = TorchLMHeadCE(
         H=H,
         V=V,
-        bias=True,
+        bias=bias,
         label_smoothing=0.0,
         reduction="mean",
         dtype=dtype,
@@ -277,7 +281,7 @@ def test_amp(B, T, H, V, cast_dtype, atol, rtol):
     liger_lm_head_ce = LigerLMHeadCE(
         H=H,
         V=V,
-        bias=True,
+        bias=bias,
         label_smoothing=0.0,
         reduction="mean",
         dtype=dtype,