PyPI - liger-kernel-nightly - Versions diffs - 0.5.6.dev20250411210855__tar.gz → 0.5.6.dev20250411224032__tar.gz - Mend

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.6.dev20250411210855"
+version = "0.5.6.dev20250411224032"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

@@ -222,7 +222,7 @@ def lce_forward(
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             hidden_size=self.config.hidden_size,
-            softcap=self.config.final_logit_softcapping,
+            final_logit_softcapping=self.config.final_logit_softcapping,
             **loss_kwargs,
         )

@@ -112,7 +112,7 @@ def causal_forward(
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             hidden_size=self.config.hidden_size,
-            softcap=self.config.final_logit_softcapping,
+            final_logit_softcapping=self.config.final_logit_softcapping,
             **loss_kwargs,
         )

@@ -1,14 +1,18 @@
+from typing import Optional
+import torch
 import torch.nn as nn
 import liger_kernel.transformers.functional as F
 def fixed_fused_linear_cross_entropy(
-    hidden_states,
-    lm_head_weight,
-    target,
-    num_items_in_batch: int = None,
+    hidden_states: torch.Tensor,
+    lm_head_weight: torch.Tensor,
+    target: torch.Tensor,
+    num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
+    final_logit_softcapping: Optional[float] = None,
     **kwargs,
 ):
     reduction = "sum" if num_items_in_batch is not None else "mean"
@@ -18,7 +22,7 @@ def fixed_fused_linear_cross_entropy(
         target,
         reduction=reduction,
         ignore_index=ignore_index,
-        **kwargs,
+        softcap=final_logit_softcapping,
     )
     if reduction == "sum":
         loss = loss / num_items_in_batch
@@ -31,15 +35,17 @@ def LigerForCausalLMLoss(
     lm_head_weight,
     labels,
     hidden_size: int,
-    num_items_in_batch: int = None,
+    num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
+    shift_labels: Optional[torch.Tensor] = None,
+    final_logit_softcapping: Optional[float] = None,
     **kwargs,
 ):
     # Skip upcast since intermediate values for the loss are all fp32 in kernel
-    labels = labels.to(hidden_states.device)
-    # Shift so that token < n predict n
-    labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
-    shift_labels = labels[..., 1:].contiguous()
+    if shift_labels is None:
+        # Shift so that token < n predict n
+        labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
+        shift_labels = labels[..., 1:].contiguous()
     # Flatten the tokens
     hidden_states = hidden_states.view(-1, hidden_size)
@@ -52,6 +58,7 @@ def LigerForCausalLMLoss(
         shift_labels,
         num_items_in_batch,
         ignore_index,
+        final_logit_softcapping,
         **kwargs,
     )
     return loss

liger-kernel-nightly 0.5.6.dev20250411210855__tar.gz → 0.5.6.dev20250411224032__tar.gz

liger-kernel-nightly 0.5.6.dev20250411210855tar.gz → 0.5.6.dev20250411224032tar.gz