liger-kernel-nightly 0.5.8.dev20250422013410__tar.gz → 0.5.8.dev20250425182211__tar.gz

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.8.dev20250422013410"
+version = "0.5.8.dev20250425182211"
 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.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/chunked_loss/dpo_loss.py

@@ -68,6 +68,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=False,
         compiled=True,
         use_ref_model=True,
+        average_log_prob=False,
         chunk_size=1,
     ):
         """
@@ -85,6 +86,7 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
             compute_nll_loss (bool): Whether to compute the NLL loss
             compiled (bool): Whether to use torch compile
             use_ref_model (bool): Whether to use a reference model
+            average_log_prob (bool): Whether to average the log probability per non-masked token
             chunk_size (int): Size of chunks for processing.
         Returns:
             torch.Tensor: Computed loss
@@ -104,13 +106,14 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
             ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
+            average_log_prob=average_log_prob,
             chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearDPOLoss(torch.nn.Module):
@@ -125,6 +128,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
         compute_nll_loss: bool = False,
         compiled: bool = True,
         use_ref_model: bool = True,
+        average_log_prob: bool = True,
         chunk_size: int = 1,
     ):
         """
@@ -134,6 +138,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
             compute_nll_loss (bool): Whether to compute the NLL loss.
             compiled (bool): Whether to use the torch compiled kernel.
             use_ref_model (bool): Whether to use a reference model for the DPO loss.
+            average_log_prob (bool): Whether to average the log probability per non-masked token.
             chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
@@ -142,6 +147,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
         self.use_ref_model = use_ref_model
+        self.average_log_prob = average_log_prob
         self.chunk_size = chunk_size
 
     def forward(
@@ -167,5 +173,6 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             self.compiled,
             self.use_ref_model,
+            self.average_log_prob,
             self.chunk_size,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/gemma.py

@@ -201,14 +201,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/gemma2.py

@@ -213,14 +213,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             final_logit_softcapping=self.config.final_logit_softcapping,
             **loss_kwargs,

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/gemma3.py

@@ -104,13 +104,15 @@ def causal_forward(
     # 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
     logits = None
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             final_logit_softcapping=self.config.final_logit_softcapping,
             **loss_kwargs,

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/llama.py

@@ -213,14 +213,16 @@ def lce_forward(
     if self.config.pretraining_tp > 1:
         raise Exception("Liger Kernel does not support pretraining_tp!!")
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/mistral.py

@@ -92,14 +92,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/mixtral.py

@@ -226,14 +226,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/mllama.py

@@ -216,14 +216,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/olmo2.py

@@ -89,14 +89,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/phi3.py

@@ -214,14 +214,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/qwen2.py

@@ -200,14 +200,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/qwen2_5_vl.py

@@ -163,14 +163,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel/transformers/model/qwen2_vl.py

@@ -167,14 +167,16 @@ def lce_forward(
 
     hidden_states = outputs[0]
 
+    shift_labels = loss_kwargs.pop("shift_labels", None)
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if self.training and (labels is not None or shift_labels is not None):
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
+            shift_labels=shift_labels,
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.8.dev20250422013410 → liger_kernel_nightly-0.5.8.dev20250425182211}/test/convergence/bf16/test_mini_models_multimodal.py

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