liger-kernel-nightly 0.5.8.dev20250503021755__tar.gz → 0.5.8.dev20250504194120__tar.gz

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.8.dev20250503021755"
+version = "0.5.8.dev20250504194120"
 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.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/gemma.py

@@ -200,6 +200,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -207,7 +210,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -215,8 +218,7 @@ def lce_forward(
             **loss_kwargs,
         )
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/gemma2.py

@@ -212,6 +212,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -219,7 +222,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -229,8 +232,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if self.config.final_logit_softcapping is not None:
             logits = logits / self.config.final_logit_softcapping
             logits = torch.tanh(logits)

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/glm4.py

@@ -88,6 +88,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -95,7 +98,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -104,8 +107,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/llama.py

@@ -209,6 +209,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     if self.config.pretraining_tp > 1:
         raise Exception("Liger Kernel does not support pretraining_tp!!")
@@ -219,7 +222,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -228,8 +231,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/mistral.py

@@ -91,6 +91,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     loss = None
@@ -98,7 +101,7 @@ def lce_forward(
 
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -107,8 +110,7 @@ def lce_forward(
         )
 
     else:
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
 
         loss = None
         if labels is not None:

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/mixtral.py

@@ -225,6 +225,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -232,7 +235,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -241,8 +244,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
 
         loss = None
         if labels is not None:

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/mllama.py

@@ -215,6 +215,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -222,7 +225,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -231,8 +234,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/olmo2.py

@@ -88,6 +88,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -95,7 +98,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -104,8 +107,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/phi3.py

@@ -213,6 +213,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -220,7 +223,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -229,8 +232,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel/transformers/model/qwen2.py

@@ -199,6 +199,9 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
@@ -206,7 +209,7 @@ def lce_forward(
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -215,8 +218,7 @@ def lce_forward(
         )
 
     else:  # if in inference mode materialize logits
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250503021755 → liger_kernel_nightly-0.5.8.dev20250504194120}/test/convergence/bf16/test_mini_models_multimodal.py

@@ -872,7 +872,7 @@ def run_mini_model_multimodal(
             32,
             1e-4,
             torch.bfloat16,
-            1e-3,
+            3e-3,
             1e-2,
             0.4,  # Increase the absolute tolerance for the logits of Gemma-3.
             1e-1,