liger-kernel-nightly 0.5.10.dev20250526154218__tar.gz → 0.5.10.dev20250528223223__tar.gz

{liger_kernel_nightly-0.5.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.10.dev20250526154218"
+version = "0.5.10.dev20250528223223"
 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.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/src/liger_kernel/transformers/model/llava.py

@@ -5,6 +5,7 @@ from typing import Union
 
 import torch
 
+from torch.nn import CrossEntropyLoss
 from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast
 from transformers.utils import is_torchdynamo_compiling
 from transformers.utils.deprecation import deprecate_kwarg
@@ -189,7 +190,20 @@ def lce_forward_deprecated(
 
         lce = LigerFusedLinearCrossEntropyLoss()
         loss = lce(self.language_model.lm_head.weight, shift_hidden_states, shift_labels)
-
+    else:
+        logits = self.language_model.lm_head(hidden_states)
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            if attention_mask is not None:
+                shift_attention_mask = attention_mask[..., 1:]
+                shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device))
     if not return_dict:
         # NOTE: This part has not been tested.
         output = outputs[1:]
@@ -349,6 +363,28 @@ def lce_forward(
             shift_hidden_states.view(-1, shift_hidden_states.size(-1)),
             shift_labels.view(-1).to(shift_hidden_states.device),
         )
+    else:
+        logits = self.language_model.lm_head(hidden_states)
+        if 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, :]
+            shift_labels = labels[..., 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)
 
     if not return_dict:
         # NOTE: This part has not been tested.

{liger_kernel_nightly-0.5.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/src/liger_kernel/transformers/monkey_patch.py

@@ -776,7 +776,7 @@ def apply_liger_kernel_to_gemma3_text(
 
     from transformers.models.gemma3 import modeling_gemma3
     from transformers.models.gemma3.modeling_gemma3 import Gemma3DecoderLayer
-    from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM
+    from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM, Gemma3TextModel
 
     from liger_kernel.transformers.gema3_rms import LigerRMSNormForGemma3
     from liger_kernel.transformers.model.gemma3 import causal_forward
@@ -807,9 +807,9 @@ def apply_liger_kernel_to_gemma3_text(
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
 
-        if isinstance(model, Gemma3ForCausalLM):
+        if isinstance(model, Gemma3ForCausalLM) or isinstance(model, Gemma3TextModel):
             # get the base model from the model instance
-            base_model = model.model
+            base_model = model.model if isinstance(model, Gemma3ForCausalLM) else model
 
             if rms_norm:
                 _patch_rms_norm_module_for_gemma3(base_model.norm)
@@ -1625,7 +1625,6 @@ def _apply_liger_kernel_to_instance(model: PreTrainedModel, **kwargs) -> None:
         return
 
     apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type]
-
     apply_fn_signature = inspect.signature(apply_fn)
 
     # Filter out the keyword arguments that are not supported by the apply function

{liger_kernel_nightly-0.5.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250526154218 → liger_kernel_nightly-0.5.10.dev20250528223223}/test/transformers/test_monkey_patch.py

@@ -667,7 +667,7 @@ def test_apply_liger_kernel_to_instance_for_gemma3_text():
 
 
 @pytest.mark.skipif(not is_gemma3_available(), reason="gemma3 module not available")
-def test_apply_liger_kernel_to_instance_for_gemma3():
+def test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation():
     # Ensure any monkey patching is cleaned up for subsequent tests
 
     with patch("transformers.models.gemma3.modeling_gemma3"):
@@ -687,8 +687,8 @@ def test_apply_liger_kernel_to_instance_for_gemma3():
             intermediate_size=64,
         )
         config = transformers.models.gemma3.configuration_gemma3.Gemma3Config(text_config, vision_config)
-        dummy_model_instance = Gemma3ForConditionalGeneration._from_config(config)
 
+        dummy_model_instance = Gemma3ForConditionalGeneration._from_config(config)
         assert isinstance(dummy_model_instance, Gemma3ForConditionalGeneration)
 
         # Check that model instance variables are not yet patched with Liger modules
@@ -704,11 +704,11 @@ def test_apply_liger_kernel_to_instance_for_gemma3():
             dummy_model_instance.multi_modal_projector.mm_soft_emb_norm.forward
         ) != inspect.getsource(LigerRMSNorm.forward)
 
-        assert inspect.getsource(dummy_model_instance.language_model.model.norm.forward) != inspect.getsource(
+        assert inspect.getsource(dummy_model_instance.language_model.norm.forward) != inspect.getsource(
             LigerRMSNorm.forward
         )
 
-        for layer in dummy_model_instance.language_model.model.layers:
+        for layer in dummy_model_instance.language_model.layers:
             assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerGEGLUMLP.forward)
             assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
             assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
@@ -736,10 +736,10 @@ def test_apply_liger_kernel_to_instance_for_gemma3():
             dummy_model_instance.multi_modal_projector.mm_soft_emb_norm.forward
         ) == inspect.getsource(LigerRMSNorm.forward)
 
-        assert inspect.getsource(dummy_model_instance.language_model.model.norm.forward) == inspect.getsource(
+        assert inspect.getsource(dummy_model_instance.language_model.norm.forward) == inspect.getsource(
             LigerRMSNorm.forward
         )
-        for layer in dummy_model_instance.language_model.model.layers:
+        for layer in dummy_model_instance.language_model.layers:
             assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerGEGLUMLP.forward)
             assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
             assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)