liger-kernel-nightly 0.6.2.dev20251005233549__tar.gz → 0.6.2.dev20251008084122__tar.gz

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/PKG-INFO

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

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/pyproject.toml

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

@@ -228,10 +228,11 @@ def lce_forward(
         )
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/gemma2.py

@@ -252,8 +252,14 @@ def lce_forward(
             logits = logits * self.config.final_logit_softcapping
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/gemma3.py

@@ -119,8 +119,14 @@ def causal_forward(
             logits = logits / self.config.final_logit_softcapping
             logits = torch.tanh(logits)
             logits = logits * self.config.final_logit_softcapping
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **loss_kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]
@@ -275,6 +281,25 @@ def multimodal_forward(
             # Flatten the tokens
             loss_fct = nn.CrossEntropyLoss()
 
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(flat_logits, flat_labels)
+        elif shift_labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            shift_logits = logits[..., :-1, :]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
+                shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+
             flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
             flat_labels = shift_labels.view(-1).to(shift_logits.device)
             loss = loss_fct(flat_logits, flat_labels)

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/glm4.py

@@ -111,10 +111,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/glm4v.py

@@ -133,10 +133,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/glm4v_moe.py

@@ -134,10 +134,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/llama.py

@@ -248,10 +248,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/llama4.py

@@ -91,10 +91,11 @@ def lce_forward(
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/llava.py

@@ -313,9 +313,13 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
-                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **lm_kwargs
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.text_config.vocab_size,
+                **lm_kwargs,
             )
 
     if not return_dict:

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/mistral.py

@@ -115,10 +115,11 @@ def lce_forward(
         logits = self.lm_head(kept_hidden_states)
 
         loss = None
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/mixtral.py

@@ -248,8 +248,14 @@ def lce_forward(
         logits = self.lm_head(kept_hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
     aux_loss = None
     if output_router_logits:
         aux_loss = load_balancing_loss_func(

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/mllama.py

@@ -239,10 +239,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/olmo2.py

@@ -111,10 +111,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/paligemma.py

@@ -379,6 +379,25 @@ def lce_forward(
             # Flatten the tokens
             loss_fct = CrossEntropyLoss()
 
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(flat_logits, flat_labels)
+        elif shift_labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            shift_logits = logits[..., :-1, :]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
+                shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+
             flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
             flat_labels = shift_labels.view(-1).to(shift_logits.device)
             loss = loss_fct(flat_logits, flat_labels)

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/phi3.py

@@ -91,10 +91,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/qwen2.py

@@ -228,10 +228,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/qwen2_5_vl.py

@@ -133,8 +133,13 @@ def lce_forward(
         logits = self.lm_head(hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.vocab_size,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/qwen2_vl.py

@@ -129,8 +129,13 @@ def lce_forward(
         logits = self.lm_head(hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.vocab_size,
+            )
 
     return Qwen2VLCausalLMOutputWithPast(
         loss=loss,

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/qwen3.py

@@ -103,10 +103,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/qwen3_moe.py

@@ -107,8 +107,14 @@ def lce_forward(
         )
     else:  # if in inference model materialize logits
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
 
     aux_loss = None
     if output_router_logits:

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel/transformers/model/smollm3.py

@@ -121,10 +121,11 @@ def lce_forward(
 
     else:
         logits = self.lm_head(kept_hidden_states)
-        if labels is not None:
+        if labels is not None or shift_labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
+                shift_labels=shift_labels,
                 vocab_size=self.config.vocab_size,
                 **kwargs,
             )

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.6.2.dev20251005233549 → liger_kernel_nightly-0.6.2.dev20251008084122}/test/utils.py

@@ -14,7 +14,9 @@ from typing import Tuple
 import numpy as np
 import torch
 import torch.nn as nn
+import transformers
 
+from packaging import version
 from tokenizers import AddedToken
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -149,18 +151,20 @@ class MiniModelConfig:
 
 def simple_collate_fn(data: List[Dict[str, Any]]):
     """A basic collate function to use for DataLoader"""
+    batch = {}
 
     input_ids = torch.stack([torch.tensor(item["input_ids"]) for item in data])
     attention_mask = torch.stack([torch.tensor(item["attention_mask"]) for item in data])
     labels = input_ids.clone()
+    batch["input_ids"] = input_ids
+    batch["attention_mask"] = attention_mask
+    batch["labels"] = labels
+    if version.parse("4.54.1") <= version.parse(transformers.__version__):
+        shift_labels = nn.functional.pad(labels, (0, 1), value=-100)
+        shift_labels = shift_labels[..., 1:].contiguous()
+        batch["shift_labels"] = shift_labels
 
-    return BatchEncoding(
-        {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "labels": labels,
-        }
-    )
+    return BatchEncoding(batch)
 
 
 def multimodal_collate_fn(data: List[Dict[str, Any]]):
@@ -174,6 +178,10 @@ def multimodal_collate_fn(data: List[Dict[str, Any]]):
 
     labels = input_ids.clone()
     batch["labels"] = labels
+    if version.parse("4.54.1") <= version.parse(transformers.__version__):
+        shift_labels = nn.functional.pad(labels, (0, 1), value=-100)
+        shift_labels = shift_labels[..., 1:].contiguous()
+        batch["shift_labels"] = shift_labels
 
     # Collate all other keys, e.g. pixel_values, attention_mask, image_grid_thw, etc
     for key in keys: