liger-kernel-nightly 0.5.5.dev20250315175408__tar.gz → 0.5.5.dev20250317215555__tar.gz

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/PKG-INFO

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

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/pyproject.toml

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

@@ -14,6 +14,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
@@ -200,22 +201,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
-
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
         if labels is not None:

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/gemma2.py

@@ -15,6 +15,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 logger = logging.getLogger(__name__)
 
@@ -212,25 +213,15 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
             softcap=self.config.final_logit_softcapping,
-            reduction=reduction,
+            **loss_kwargs,
         )
 
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
-
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
         if self.config.final_logit_softcapping is not None:

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/llama.py

@@ -15,6 +15,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 if TYPE_CHECKING:
     from transformers.cache_utils import Cache
@@ -212,21 +213,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

liger_kernel_nightly-0.5.5.dev20250317215555/src/liger_kernel/transformers/model/loss_utils.py

@@ -0,0 +1,57 @@
+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,
+    ignore_index: int = -100,
+    **kwargs,
+):
+    reduction = "sum" if num_items_in_batch is not None else "mean"
+    loss = F.liger_fused_linear_cross_entropy(
+        hidden_states,
+        lm_head_weight,
+        target,
+        reduction=reduction,
+        ignore_index=ignore_index,
+        **kwargs,
+    )
+    if reduction == "sum":
+        loss = loss / num_items_in_batch
+
+    return loss
+
+
+def LigerForCausalLMLoss(
+    hidden_states,
+    lm_head_weight,
+    labels,
+    hidden_size: int,
+    num_items_in_batch: int = None,
+    ignore_index: int = -100,
+    **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()
+
+    # Flatten the tokens
+    hidden_states = hidden_states.view(-1, hidden_size)
+    shift_labels = shift_labels.view(-1)
+    # Enable model parallelism
+    shift_labels = shift_labels.to(hidden_states.device)
+    loss = fixed_fused_linear_cross_entropy(
+        hidden_states,
+        lm_head_weight,
+        shift_labels,
+        num_items_in_batch,
+        ignore_index,
+        **kwargs,
+    )
+    return loss

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/mistral.py

@@ -13,7 +13,7 @@ from transformers.models.mistral.modeling_mistral import MISTRAL_INPUTS_DOCSTRIN
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
-from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
@@ -31,6 +31,7 @@ def lce_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste Mistral's forward but replace torch cross entropy with liger fused linear cross entropy
@@ -87,15 +88,13 @@ def lce_forward(
     logits = None
 
     if self.training and (labels is not None):
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
-        shift_labels = shift_labels.view(-1)
-
-        lce = LigerFusedLinearCrossEntropyLoss()
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:
         logits = self.lm_head(hidden_states)

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/mixtral.py

@@ -14,6 +14,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(MIXTRAL_INPUTS_DOCSTRING)
@@ -225,21 +226,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/mllama.py

@@ -13,6 +13,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(MLLAMA_INPUTS_DOCSTRING)
@@ -215,21 +216,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/olmo2.py

@@ -11,7 +11,7 @@ from transformers.models.olmo2.modeling_olmo2 import OLMO2_INPUTS_DOCSTRING
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
-from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(OLMO2_INPUTS_DOCSTRING)
@@ -89,21 +89,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/paligemma.py

@@ -172,7 +172,7 @@ def lce_forward(
             shift_labels = shift_labels.contiguous()
 
         # Flatten hidden state
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
+        shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size)
         shift_labels = shift_labels.view(-1).to(hidden_device)
 
         lce = LigerFusedLinearCrossEntropyLoss()

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/phi3.py

@@ -13,6 +13,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
@@ -213,21 +214,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/qwen2.py

@@ -13,6 +13,7 @@ from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
@@ -199,21 +200,13 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None):
-        # We do the same thing as ForCausalLMLoss but using Liger FLCE
-
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # flatten tokens
-        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 loss_kwargs else "mean"
-        lce = LigerFusedLinearCrossEntropyLoss(reduction=reduction)
-
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
-        if reduction == "sum":
-            loss /= loss_kwargs["num_items_in_batch"]
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
 
     else:  # if in inference mode materialize logits
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/qwen2_5_vl.py

@@ -12,7 +12,7 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLCausalL
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
-from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(QWEN2_5_VL_INPUTS_DOCSTRING)
@@ -36,6 +36,7 @@ def lce_forward(
     rope_deltas: Optional[torch.LongTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
     second_per_grid_ts: Optional[torch.Tensor] = None,
+    **loss_kwargs,
 ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
     r"""
     Copy paste Qwen2_5_VL's forward but replace torch cross entropy with liger fused linear cross entropy
@@ -166,15 +167,13 @@ def lce_forward(
     logits = None
 
     if self.training and (labels is not None):
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # Flatten tokens
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
-        shift_labels = shift_labels.view(-1)
-
-        lce = LigerFusedLinearCrossEntropyLoss()
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
     else:
         logits = self.lm_head(hidden_states)
         if labels is not None:

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel/transformers/model/qwen2_vl.py

@@ -14,7 +14,7 @@ from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLCausalLMOutput
 from transformers.utils import add_start_docstrings_to_model_forward
 from transformers.utils import replace_return_docstrings
 
-from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
 
 @add_start_docstrings_to_model_forward(QWEN2_VL_INPUTS_DOCSTRING)
@@ -37,6 +37,7 @@ def lce_forward(
     video_grid_thw: Optional[torch.LongTensor] = None,
     rope_deltas: Optional[torch.LongTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    **loss_kwargs,
 ) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]:
     r"""
     Copy paste Qwen2VL's forward but replace torch cross entropy with liger fused linear cross entropy
@@ -170,15 +171,13 @@ def lce_forward(
     logits = None
 
     if self.training and (labels is not None):
-        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        # Flatten tokens
-        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
-        shift_labels = shift_labels.view(-1)
-
-        lce = LigerFusedLinearCrossEntropyLoss()
-        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+        loss = LigerForCausalLMLoss(
+            hidden_states=hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            hidden_size=self.config.hidden_size,
+            **loss_kwargs,
+        )
     else:
         logits = self.lm_head(hidden_states)
         if labels is not None:

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -159,6 +159,7 @@ src/liger_kernel/transformers/model/__init__.py
 src/liger_kernel/transformers/model/gemma.py
 src/liger_kernel/transformers/model/gemma2.py
 src/liger_kernel/transformers/model/llama.py
+src/liger_kernel/transformers/model/loss_utils.py
 src/liger_kernel/transformers/model/mistral.py
 src/liger_kernel/transformers/model/mixtral.py
 src/liger_kernel/transformers/model/mllama.py

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/test/convergence/bf16/test_mini_models_multimodal.py

@@ -198,7 +198,6 @@ if PALIGEMMA_AVAILABLE:
             attn_implementation="eager",
             vocab_size=32000,
             projection_dim=1024,
-            hidden_size=1024,
         ),
     )
 

{liger_kernel_nightly-0.5.5.dev20250315175408 → liger_kernel_nightly-0.5.5.dev20250317215555}/test/convergence/fp32/test_mini_models_multimodal.py

@@ -197,7 +197,6 @@ if PALIGEMMA_AVAILABLE:
             attn_implementation="eager",
             vocab_size=32000,
             projection_dim=1024,
-            hidden_size=1024,
         ),
     )