liger-kernel-nightly 0.5.10.dev20250522174514__tar.gz → 0.5.10.dev20250523162037__tar.gz

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/.gitignore

@@ -14,6 +14,9 @@ venv/
 build/
 dist/
 
+# Doc Build
+doc_site/
+
 # Lockfiles
 uv.lock
 

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/Makefile

@@ -48,13 +48,19 @@ run-benchmarks:
 # MkDocs Configuration
 MKDOCS = mkdocs
 CONFIG_FILE = mkdocs.yml
+SITE_DIR = doc_site
 
 # MkDocs targets
+
+# Serve the documentation 
 serve:
 	$(MKDOCS) serve -f $(CONFIG_FILE)
 
+# Build the documentation into the specified site directory
 build:
-	$(MKDOCS) build -f $(CONFIG_FILE)
+	$(MKDOCS) build -f $(CONFIG_FILE) --site-dir $(SITE_DIR)
 
+# Clean the output directory
 clean:
-	rm -rf site/
+	rm -rf $(SITE_DIR)/
+

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/mkdocs.yml

@@ -1,5 +1,6 @@
 site_name: Liger-Kernel Docs
-# site_url: ...
+site_dir: './doc_site'
+# site_url: https://linkedin.github.io/Liger-Kernel/
 # site_author: LinkedIn
 site_description: Efficient Triton Kernels for LLM Training
 theme:

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/pyproject.toml

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

@@ -137,6 +137,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -199,8 +200,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -209,7 +217,7 @@ def lce_forward(
             hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/gemma2.py

@@ -146,6 +146,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -213,8 +214,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -225,7 +233,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if self.config.final_logit_softcapping is not None:
             logits = logits / self.config.final_logit_softcapping

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/gemma3.py

@@ -35,6 +35,7 @@ def causal_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -101,7 +102,11 @@ def causal_forward(
     shift_labels = loss_kwargs.pop("shift_labels", None)
     loss = None
     logits = None
-    if self.training and (labels is not None or shift_labels is not None):
+
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -151,6 +156,7 @@ def multimodal_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **lm_kwargs,
 ) -> Union[Tuple, Gemma3CausalLMOutputWithPast]:
     r"""
@@ -272,7 +278,13 @@ def multimodal_forward(
     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:
+        skip_logits = self.training and (labels is not None)
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :]
         shift_labels = labels[..., 1:]
 

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/glm4.py

@@ -26,6 +26,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -89,8 +90,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -100,7 +108,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/llama.py

@@ -151,6 +151,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -218,7 +219,14 @@ def lce_forward(
     logits = None
     loss = None
     # if in training mode, don't materialize logits
-    if self.training and (labels is not None or shift_labels is not None):
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = lce_maybe_trainable_lm_head(
             self,
             hidden_states=kept_hidden_states,
@@ -228,7 +236,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/llava.py

@@ -223,6 +223,7 @@ def lce_forward(
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
     image_sizes: torch.Tensor = None,
+    skip_logits: Optional[bool] = None,
     **lm_kwargs,
 ) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
     r"""
@@ -325,7 +326,10 @@ def lce_forward(
     loss = None
     logits = None
 
-    if self.training and (labels is not None):
+    # Overwrite skip_logits, since llava never materializes logits
+    skip_logits = labels is not None
+
+    if skip_logits:
         # Shift so that tokens < n predict n
         if attention_mask is not None:
             # we use the input attention mask to shift the logits and labels, because it is 2D.

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/mistral.py

@@ -27,6 +27,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -93,7 +94,13 @@ def lce_forward(
     loss = None
     logits = None
 
-    if self.training and (labels is not None or shift_labels is not None):
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/mixtral.py

@@ -156,6 +156,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
     r"""
@@ -224,8 +225,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -235,7 +243,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
 
         loss = None

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/mllama.py

@@ -147,6 +147,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -215,8 +216,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -226,7 +234,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/olmo2.py

@@ -26,6 +26,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -89,8 +90,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -100,7 +108,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/paligemma.py

@@ -216,6 +216,7 @@ def lce_forward(
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **lm_kwargs,
 ) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]:
     r"""
@@ -331,7 +332,13 @@ def lce_forward(
     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:
+        skip_logits = self.training and (labels is not None)
+
+    if skip_logits:
         shift_hidden_states = hidden_states[..., :-1, :]
         shift_labels = labels[..., 1:]
 

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/phi3.py

@@ -136,6 +136,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -212,8 +213,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -223,7 +231,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/qwen2.py

@@ -135,6 +135,7 @@ def lce_forward(
     return_dict: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
     r"""
@@ -198,8 +199,15 @@ def lce_forward(
     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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -209,7 +217,7 @@ def lce_forward(
             **loss_kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/qwen2_5_vl.py

@@ -30,6 +30,7 @@ def lce_forward(
     rope_deltas: Optional[torch.LongTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
     second_per_grid_ts: Optional[torch.Tensor] = None,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
     r"""
@@ -161,7 +162,13 @@ def lce_forward(
     loss = None
     logits = None
 
-    if self.training and (labels is not None or shift_labels is not None):
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/qwen2_vl.py

@@ -31,6 +31,7 @@ def lce_forward(
     video_grid_thw: Optional[torch.LongTensor] = None,
     rope_deltas: Optional[torch.LongTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]:
     r"""
@@ -165,7 +166,13 @@ def lce_forward(
     loss = None
     logits = None
 
-    if self.training and (labels is not None or shift_labels is not None):
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=hidden_states,
             lm_head_weight=self.lm_head.weight,

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/qwen3.py

@@ -22,6 +22,7 @@ def lce_forward(
     output_hidden_states: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **kwargs,
 ) -> CausalLMOutputWithPast:
     r"""
@@ -82,8 +83,15 @@ def lce_forward(
     shift_labels = 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 or shift_labels is not None):
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are 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 or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
@@ -93,7 +101,7 @@ def lce_forward(
             **kwargs,
         )
 
-    else:  # if in inference mode materialize logits
+    else:
         logits = self.lm_head(kept_hidden_states)
         if labels is not None:
             loss = self.loss_function(

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel/transformers/model/qwen3_moe.py

@@ -25,6 +25,7 @@ def lce_forward(
     output_router_logits: Optional[bool] = None,
     cache_position: Optional[torch.LongTensor] = None,
     logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
     **loss_kwargs,
 ) -> MoeCausalLMOutputWithPast:
     r"""
@@ -91,8 +92,10 @@ def lce_forward(
     logits = None
     loss = None
 
-    # if in training mode, do not materialize logits
-    if self.training and (labels is not None or shift_labels is not None):
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
         loss = LigerForCausalLMLoss(
             hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,

{liger_kernel_nightly-0.5.10.dev20250522174514 → liger_kernel_nightly-0.5.10.dev20250523162037}/src/liger_kernel_nightly.egg-info/PKG-INFO

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