liger-kernel-nightly 0.3.1.dev20241101044713__tar.gz → 0.3.1.dev20241102065152__tar.gz

{liger_kernel_nightly-0.3.1.dev20241101044713 → liger_kernel_nightly-0.3.1.dev20241102065152}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.3.1.dev20241101044713
+Version: 0.3.1.dev20241102065152
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -36,7 +36,7 @@ License-File: LICENSE-MIT-llmc
 License-File: LICENSE-MIT-triton
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2
-Requires-Dist: triton>=2.3.0
+Requires-Dist: triton>=2.3.1
 Provides-Extra: transformers
 Requires-Dist: transformers~=4.0; extra == "transformers"
 Provides-Extra: dev

{liger_kernel_nightly-0.3.1.dev20241101044713 → liger_kernel_nightly-0.3.1.dev20241102065152}/pyproject.toml

@@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.3.1.dev20241101044713"
+version = "0.3.1.dev20241102065152"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }
 dependencies = [
     "torch>=2.1.2",
-    "triton>=2.3.0",
+    "triton>=2.3.1",
 ]
 
 [project.optional-dependencies]

{liger_kernel_nightly-0.3.1.dev20241101044713 → liger_kernel_nightly-0.3.1.dev20241102065152}/src/liger_kernel/transformers/model/llama.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -18,6 +18,10 @@ from liger_kernel.transformers.fused_linear_cross_entropy import (
 )
 
 
+if TYPE_CHECKING:
+    from transformers.cache_utils import Cache
+
+
 @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
 @replace_return_docstrings(
     output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
@@ -27,7 +31,7 @@ def lce_forward_deprecated(
     input_ids: torch.LongTensor = None,
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    past_key_values: Optional[Union["Cache", List[torch.FloatTensor]]] = None,
     inputs_embeds: Optional[torch.FloatTensor] = None,
     labels: Optional[torch.LongTensor] = None,
     use_cache: Optional[bool] = None,
@@ -153,19 +157,19 @@ def lce_forward_deprecated(
 )
 def lce_forward(
     self,
-    input_ids=None,
-    attention_mask=None,
-    position_ids=None,
-    past_key_values=None,
-    inputs_embeds=None,
-    labels=None,
-    use_cache=None,
-    output_attentions=None,
-    output_hidden_states=None,
-    return_dict=None,
-    cache_position=None,
-    num_logits_to_keep=0,
-    **kwargs,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Union["Cache", List[torch.FloatTensor]]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    num_logits_to_keep: int = 0,
+    **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Args:
@@ -224,7 +228,6 @@ def lce_forward(
         output_hidden_states=output_hidden_states,
         return_dict=return_dict,
         cache_position=cache_position,
-        **kwargs,
     )
 
     hidden_states = outputs[0]
@@ -245,12 +248,12 @@ def lce_forward(
         shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
         shift_labels = shift_labels.view(-1)
 
-        reduction = "sum" if "num_items_in_batch" in kwargs else "mean"
+        reduction = "sum" if "num_items_in_batch" in loss_kwargs else "mean"
         lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
 
         loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
         if reduction == "sum":
-            loss /= kwargs["num_items_in_batch"]
+            loss /= loss_kwargs["num_items_in_batch"]
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
@@ -259,7 +262,7 @@ def lce_forward(
                 logits=logits,
                 labels=labels,
                 vocab_size=self.config.vocab_size,
-                **kwargs,
+                **loss_kwargs,
             )
 
     if not return_dict:

{liger_kernel_nightly-0.3.1.dev20241101044713 → liger_kernel_nightly-0.3.1.dev20241102065152}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.3.1.dev20241101044713
+Version: 0.3.1.dev20241102065152
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -36,7 +36,7 @@ License-File: LICENSE-MIT-llmc
 License-File: LICENSE-MIT-triton
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2
-Requires-Dist: triton>=2.3.0
+Requires-Dist: triton>=2.3.1
 Provides-Extra: transformers
 Requires-Dist: transformers~=4.0; extra == "transformers"
 Provides-Extra: dev

{liger_kernel_nightly-0.3.1.dev20241101044713 → liger_kernel_nightly-0.3.1.dev20241102065152}/src/liger_kernel_nightly.egg-info/requires.txt

@@ -1,5 +1,5 @@
 torch>=2.1.2
-triton>=2.3.0
+triton>=2.3.1
 
 [dev]
 transformers>=4.44.2