liger-kernel-nightly 0.5.10.dev20250702150221__tar.gz → 0.5.10.dev20250704061237__tar.gz

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/.gitignore

@@ -6,6 +6,7 @@ site/
 venv/
 .ipynb_checkpoints/
 .vscode/
+.idea/
 
 # Misc
 .DS_Store

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.10.dev20250702150221"
+version = "0.5.10.dev20250704061237"
 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.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/ops/geglu.py

@@ -40,7 +40,7 @@ def _geglu_tanh_forward_kernel(a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE
     tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
     tanh_result = tanh(tanh_arg)
     geglu_a = 0.5 * a_row * (1 + tanh_result)
-    c_row = geglu_a * b_row
+    c_row = geglu_a.cast(b_row.dtype) * b_row
     tl.store(c + col_offsets, c_row, mask=mask)
 
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/ops/swiglu.py

@@ -26,7 +26,7 @@ def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BL
     # sigmoid requires type float32
     a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
     b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
-    c_row = silu(a_row) * b_row
+    c_row = silu(a_row).cast(b_row.dtype) * b_row
     tl.store(c_ptr + col_offsets, c_row, mask=mask)
 
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/gemma.py

@@ -27,6 +27,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
 
@@ -81,7 +82,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/gemma2.py

@@ -30,6 +30,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
     **kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -85,7 +86,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/llama.py

@@ -37,6 +37,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste llama forward but replace torch cross entropy with liger fused linear cross entropy
@@ -91,7 +92,15 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    # if in training mode, don't materialize logits
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/mistral.py

@@ -133,6 +133,3 @@ def lce_forward(
         hidden_states=outputs.hidden_states,
         attentions=outputs.attentions,
     )
-
-
-# Note: Grad Acc is not fixed in mistral at transformer 4.46.1

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/phi3.py

@@ -26,6 +26,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste phi3 forward from transfomers v4.44.2 but replace torch cross entropy with liger fused linear cross entropy
@@ -80,7 +81,14 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
-    if self.training and labels is not None:
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/model/qwen2.py

@@ -26,6 +26,7 @@ def lce_forward_deprecated(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
     Copy paste Qwen2's forward but replace torch cross entropy with liger fused linear cross entropy
@@ -80,6 +81,13 @@ def lce_forward_deprecated(
     loss = None
     logits = None
 
+    if skip_logits and labels is None:
+        raise ValueError("skip_logits is True, but labels is None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and labels is not None
+
     if self.training and (labels is not None):
         shift_hidden_states = hidden_states[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel/transformers/monkey_patch.py

@@ -611,10 +611,17 @@ def apply_liger_kernel_to_mistral(
     if cross_entropy:
         modeling_mistral.CrossEntropyLoss = LigerCrossEntropyLoss
     if fused_linear_cross_entropy:
-        if model is not None:
-            model.forward = MethodType(mistral_lce_forward, model)
+        if transformer_version >= version.parse("4.49.0"):
+            if model is not None:
+                model.forward = MethodType(mistral_lce_forward, model)
+            else:
+                modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
         else:
-            modeling_mistral.MistralForCausalLM.forward = mistral_lce_forward
+            logger.warning(
+                "The latest version of Liger does not support transformers < 4.49.0 for llava. Please downgrade your liger version or upgrade your transformer version."
+            )
+            logger.warning("LigerFusedLinearCrossEntropy patch is not applied.")
+
     if swiglu:
         modeling_mistral.MistralMLP = LigerSwiGLUMLP
 

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -16,7 +16,6 @@ setup.py
 .github/workflows/nvi-ci.yml
 .github/workflows/publish-nightly.yml
 .github/workflows/publish-release.yml
-.idea/workspace.xml
 benchmark/README.md
 benchmark/__init__.py
 benchmark/benchmarks_visualizer.py

{liger_kernel_nightly-0.5.10.dev20250702150221 → liger_kernel_nightly-0.5.10.dev20250704061237}/test/convergence/bf16/test_mini_models.py

@@ -9,8 +9,6 @@ from transformers.models.gemma2 import Gemma2Config
 from transformers.models.gemma2 import Gemma2ForCausalLM
 from transformers.models.llama import LlamaConfig
 from transformers.models.llama import LlamaForCausalLM
-from transformers.models.llama4 import Llama4ForCausalLM
-from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
 from transformers.models.mistral import MistralConfig
 from transformers.models.mistral import MistralForCausalLM
 from transformers.models.mixtral import MixtralConfig
@@ -65,6 +63,14 @@ from test.utils import set_seed
 from test.utils import simple_collate_fn
 from test.utils import supports_bfloat16
 
+try:
+    from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
+    from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
+
+    LLAMA4_AVAILABLE = True
+except ImportError:
+    LLAMA4_AVAILABLE = False
+
 try:
     # Mllama is only available in transformers>=4.45.0
     from transformers.models.mllama.configuration_mllama import MllamaTextConfig
@@ -156,35 +162,6 @@ from liger_kernel.utils import infer_device
 device = infer_device()
 
 MINI_MODEL_SETUPS = {
-    "mini_llama4": MiniModelConfig(
-        liger_kernel_patch_func=apply_liger_kernel_to_llama4,
-        liger_kernel_patch_revert_func=revert_liger_kernel_to_llama4,
-        model_class=Llama4ForCausalLM,
-        mini_model_config=Llama4TextConfig(
-            bos_token_id=1,  # None
-            eos_token_id=2,  # 151329, 151336, 151338
-            pad_token_id=2,  # 151329
-            partial_rotary_factor=1.0,
-            cross_attention_layers=None,
-            dropout=0,
-            hidden_act="silu",
-            hidden_size=1024,  # 6144
-            initializer_range=0.02,
-            intermediate_size=2048,  # 14336
-            max_position_embeddings=4096,  # 32768
-            num_attention_heads=8,  # 48
-            num_hidden_layers=4,  # 61
-            num_key_value_heads=2,
-            rms_norm_eps=1e-5,
-            rope_scaling=None,
-            rope_theta=10000.0,
-            tie_word_embeddings=False,
-            use_cache=True,
-            vocab_size=32000,  # 151552
-            attention_bias=True,
-            attn_implementation="sdpa",  # default value, pytorch native attention
-        ),
-    ),
     "mini_llama3": MiniModelConfig(
         liger_kernel_patch_func=apply_liger_kernel_to_llama,
         liger_kernel_patch_revert_func=revert_liger_kernel_to_llama,
@@ -413,6 +390,37 @@ MINI_MODEL_SETUPS = {
     ),
 }
 
+if LLAMA4_AVAILABLE:
+    MINI_MODEL_SETUPS["mini_llama4"] = MiniModelConfig(
+        liger_kernel_patch_func=apply_liger_kernel_to_llama4,
+        liger_kernel_patch_revert_func=revert_liger_kernel_to_llama4,
+        model_class=Llama4ForCausalLM,
+        mini_model_config=Llama4TextConfig(
+            bos_token_id=1,  # None
+            eos_token_id=2,  # 151329, 151336, 151338
+            pad_token_id=2,  # 151329
+            partial_rotary_factor=1.0,
+            cross_attention_layers=None,
+            dropout=0,
+            hidden_act="silu",
+            hidden_size=1024,  # 6144
+            initializer_range=0.02,
+            intermediate_size=2048,  # 14336
+            max_position_embeddings=4096,  # 32768
+            num_attention_heads=8,  # 48
+            num_hidden_layers=4,  # 61
+            num_key_value_heads=2,
+            rms_norm_eps=1e-5,
+            rope_scaling=None,
+            rope_theta=10000.0,
+            tie_word_embeddings=False,
+            use_cache=True,
+            vocab_size=32000,  # 151552
+            attention_bias=True,
+            attn_implementation="sdpa",  # default value, pytorch native attention
+        ),
+    )
+
 
 if QWEN3_AVAILABLE:
     MINI_MODEL_SETUPS["mini_qwen3"] = MiniModelConfig(
@@ -902,23 +910,29 @@ def run_mini_model(
         pytest.param(
             "mini_llama4",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,
             1e-1,
             1e-2,
             1e-2,
-            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+            marks=[
+                pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+                pytest.mark.skipif(
+                    not LLAMA4_AVAILABLE,
+                    reason="Llama not available in this version of transformers",
+                ),
+            ],
         ),
         pytest.param(
             "mini_llama3",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,
             1e-2,
             1e-2,
@@ -928,10 +942,10 @@ def run_mini_model(
         pytest.param(
             "mini_llava",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,
             1e-1,
             1e-2,
@@ -942,17 +956,21 @@ def run_mini_model(
                     not LLAVA_AVAILABLE,
                     reason="LLaVa not available in this version of transformers",
                 ),
+                pytest.mark.skipif(
+                    version.parse(transformers.__version__) < version.parse("4.49.0"),
+                    reason="Mistral not available in transformers<=4.49.0",
+                ),
             ],
         ),
         pytest.param(
             "mini_granite3",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,  # 1e-1
-            1e-1,  # 1e-2
+            1e-2,  # 1e-2
             1e-2,
             1e-2,
             marks=[
@@ -966,9 +984,9 @@ def run_mini_model(
         pytest.param(
             "mini_mllama",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             1e-2,
             1e-1,
             1e-2,
@@ -985,10 +1003,10 @@ def run_mini_model(
         pytest.param(
             "mini_qwen2",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,
             1e-2,
             1e-2,
@@ -998,10 +1016,10 @@ def run_mini_model(
         pytest.param(
             "mini_qwen3",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
+            5e-2,
             1e-1,
             1e-2,
             1e-2,
@@ -1014,13 +1032,16 @@ def run_mini_model(
                 ),
             ],
         ),
+        # TODO(tcc): Investigate qwen3_moe on different machines.
+        # The loss diverges on ci test (A10G), but it never diverges on my local machine (3080).
+        # Qwen3_moe can pass float32 tests.
         pytest.param(
             "mini_qwen3_moe",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
-            1e-2,
+            5e-2,
+            5e-2,
             1e-1,  # 1e-1
             1e-1,  # 1e-2
             1e-2,
@@ -1036,12 +1057,12 @@ def run_mini_model(
         pytest.param(
             "mini_qwen2_vl",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             5e-2,
-            1,  # 1e-1
-            1e-1,  # 1e-2
+            1e-1,  # 1e-1
+            1e-2,  # 1e-2
             1e-2,
             1e-2,
             marks=[
@@ -1052,16 +1073,15 @@ def run_mini_model(
                 ),
             ],
         ),
-        # TODO: logits tolerances are significantly larger than the other tests, need to investigate
         pytest.param(
             "mini_qwen2_5_vl",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             5e-2,
-            3,  # 1e-1
-            1e-1,  # 1e-2
+            1e-1,  # 1e-1
+            1e-2,  # 1e-2
             1e-2,
             1e-2,
             marks=[
@@ -1075,9 +1095,9 @@ def run_mini_model(
         pytest.param(
             "mini_phi3",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             1e-2,
             1e-1,
             1e-2,
@@ -1088,22 +1108,28 @@ def run_mini_model(
         pytest.param(
             "mini_mistral",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
-            1e-2,
+            5e-2,
+            5e-2,
             1e-1,
             1e-2,
             1e-2,
             1e-2,
-            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+            marks=[
+                pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+                pytest.mark.skipif(
+                    version.parse(transformers.__version__) < version.parse("4.49.0"),
+                    reason="Mistral not available in transformers<=4.49.0",
+                ),
+            ],
         ),
         pytest.param(
             "mini_olmo2",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             1e-2,
             1e-1,
             1e-2,
@@ -1120,9 +1146,9 @@ def run_mini_model(
         pytest.param(
             "mini_glm4",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
+            1e-2,
             1e-2,
             1e-1,
             1e-2,
@@ -1156,27 +1182,27 @@ def run_mini_model(
         pytest.param(
             "mini_gemma1",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
             1e-2,
             1e-1,
             1e-2,
             1e-2,
+            1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
         ),
         pytest.param(
             "mini_gemma1.1",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
             1e-2,
             1e-1,
             1e-2,
             1e-2,
+            1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
         ),
         # TODO: Gemma2 test for bf16 is not passing within the tolerance range, might be casting issue, need to investigate
@@ -1198,12 +1224,12 @@ def run_mini_model(
         pytest.param(
             "mini_gemma3_text",
             32,
-            1e-4,
+            1e-5,
             torch.bfloat16,
-            1e-3,
             1e-2,
-            3e-1,
-            4e-1,
+            1e-2,
+            1e-1,
+            1e-2,
             1e-2,
             1e-2,
             marks=[
@@ -1240,6 +1266,7 @@ def test_mini_model(
         torch.tensor([actual_output["loss"]]),
         atol=loss_atol,
         rtol=loss_rtol,
+        extra_info="[Loss]",
     )
 
     # Compare the topk logprobs from evaluation step
@@ -1249,6 +1276,7 @@ def test_mini_model(
             actual_output["topk_logprobs"],
             atol=logprobs_atol,
             rtol=logprobs_rtol,
+            extra_info="[Top k logprobs]",
         )
 
     # Compare the params from the last step
@@ -1257,4 +1285,6 @@ def test_mini_model(
         expected_output["model"].named_parameters(),
         actual_output["model"].named_parameters(),
     ):
-        assert_verbose_allclose(expected_param[1], actual_param[1], atol=param_atol, rtol=param_rtol)
+        assert_verbose_allclose(
+            expected_param[1], actual_param[1], atol=param_atol, rtol=param_rtol, extra_info="[Model parameters]"
+        )