liger-kernel-nightly 0.5.9.dev20250519015630__tar.gz → 0.5.9.dev20250519025610__tar.gz

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/PKG-INFO

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

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.9.dev20250519015630"
+version = "0.5.9.dev20250519025610"
 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.9.dev20250519025610/src/liger_kernel/transformers/fsdp.py

@@ -0,0 +1,55 @@
+from typing import Any
+from typing import Callable
+
+from torch.distributed.fsdp import FullyShardedDataParallel
+
+
+class _FSDPForwardRedirection:
+    """
+    Modified based on
+    https://github.com/Lightning-AI/pytorch-lightning/blob/d3f9c83d6efa4f1def36aa6c199600946cdb9117/src/lightning/pytorch/strategies/strategy.py#L601-L648
+    Redirect a method call through FullyShardedDataParallel.forward so that the FSDP module's root pre-forward and
+    post-forward can be properly executed around the method call.
+    This is needed in cases where we call a submodule of a FSDP module. For instance, when we want to call only
+    the `LlamaModel` part out of a FSDP-wrapped `LlamaForCausalLM` to get the hidden states without involving
+    GPU-memory-heavy `lm_head` and cross entropy computation, doing this directly (i.e. `model.model.forward()`)
+    will not work because the first `nn.Embedding` layer is not independently wrapped as a FSDP module (because of
+    the transformer-based wrapping policy), and not calling it through FSDP root module forward will not all-gather
+    its parameter, thus resulting in "RuntimeError: 'weight' must be 2-D" error. Similarly, if we want to call just
+    the `lm_head` part of a model, we need this trick too to properly get its params all-gathered.
+    """
+
+    def __call__(
+        self,
+        wrapper_module: FullyShardedDataParallel,
+        method: Callable,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """Reroutes a method call through the `wrapper_module`'s `forward` method.
+        Args:
+            wrapper_module: The module that has `original_module` wrapped.
+            original_module: The module that was wrapped inside `wrapper_module`.
+            method_name: The name of the method that should be called on the `original_module` after inputs get
+                redirected through the `wrapper_module`'s `forward` method.
+            *args: The positional arguments to the method `method_name`. They will get passed to a patched
+                `forward` method instead.
+            **kwargs: The keyword arguments to the method `method_name`. They will get passed to a patched
+                `forward` method instead.
+        """
+        assert isinstance(wrapper_module, FullyShardedDataParallel)
+        original_module = wrapper_module._fsdp_wrapped_module
+        original_forward = original_module.forward
+
+        def wrapped_forward(*_args: Any, **_kwargs: Any) -> Any:
+            # Unpatch ourselves immediately before calling the method `method_name`
+            # because itself may want to call the real `forward`
+            original_module.forward = original_forward  # type: ignore[method-assign]
+            # Call the actual method e.g. `.training_step(...)`
+            out = method(*_args, **_kwargs)
+            return out
+
+        # Patch the original_module's forward so we can redirect the arguments back to the real method
+        original_module.forward = wrapped_forward  # type: ignore[method-assign]
+        wrapper_output = wrapper_module(*args, **kwargs)
+        return wrapper_output

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/src/liger_kernel/transformers/model/llama.py

@@ -7,16 +7,22 @@ from typing import Union
 import torch
 import torch.nn.functional as F
 
+from torch.distributed.fsdp import FullyShardedDataParallel
 from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.utils.deprecation import deprecate_kwarg
 
+from liger_kernel.transformers.fsdp import _FSDPForwardRedirection
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
 from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.utils import PEFT_AVAILABLE
 
 if TYPE_CHECKING:
     from transformers.cache_utils import Cache
 
+if PEFT_AVAILABLE:
+    from peft.utils.other import ModulesToSaveWrapper
+
 
 def lce_forward_deprecated(
     self,
@@ -213,12 +219,12 @@ def lce_forward(
     loss = None
     # if in training mode, don't materialize logits
     if self.training and (labels is not None or shift_labels is not None):
-        loss = LigerForCausalLMLoss(
+        loss = lce_maybe_trainable_lm_head(
+            self,
             hidden_states=kept_hidden_states,
-            lm_head_weight=self.lm_head.weight,
+            hidden_size=self.config.hidden_size,
             labels=labels,
             shift_labels=shift_labels,
-            hidden_size=self.config.hidden_size,
             **loss_kwargs,
         )
 
@@ -243,3 +249,50 @@ def lce_forward(
         hidden_states=outputs.hidden_states,
         attentions=outputs.attentions,
     )
+
+
+def lce_maybe_trainable_lm_head(self, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs):
+    lm_head = self.lm_head
+
+    # Unwrap the module if lm_head has been added as trainable module in PEFT LoRA configuration,
+    # i.e. listed in the modules_to_save field of LoraConfig, so the lm_head weights are read
+    # from the unwrapped module.
+    # See https://huggingface.co/docs/peft/package_reference/lora for reference.
+    if PEFT_AVAILABLE and isinstance(lm_head, ModulesToSaveWrapper):
+        lm_head = lm_head.modules_to_save.default
+
+    # If FSDP is used and lm_head is trainable, e.g., during full fine-tuning or with LoRA,
+    # reading the lm_head module weights and calling the kernel must be done within FSDP forward pass
+    # so the module entire parameters are summoned and kept in memory during the kernel execution.
+    if isinstance(lm_head, FullyShardedDataParallel):
+        return _FSDPForwardRedirection()(
+            lm_head,
+            _liger_for_causal_lm_loss,
+            lm_head.module,
+            hidden_states,
+            hidden_size,
+            labels,
+            shift_labels,
+            **loss_kwargs,
+        )
+
+    # FSDP is not used so we can read the lm_head weights and call the kernel directly
+    return _liger_for_causal_lm_loss(
+        lm_head=self.lm_head,
+        hidden_states=hidden_states,
+        hidden_size=hidden_size,
+        labels=labels,
+        shift_labels=shift_labels,
+        **loss_kwargs,
+    )
+
+
+def _liger_for_causal_lm_loss(lm_head, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs):
+    return LigerForCausalLMLoss(
+        hidden_states=hidden_states,
+        lm_head_weight=lm_head.weight,
+        labels=labels,
+        hidden_size=hidden_size,
+        shift_labels=shift_labels,
+        **loss_kwargs,
+    )

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/src/liger_kernel/transformers/trainer/orpo_trainer.py

@@ -1,5 +1,3 @@
-from typing import Any
-from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Literal
@@ -13,57 +11,7 @@ from torch.distributed.fsdp import FullyShardedDataParallel
 from trl.trainer import ORPOTrainer
 
 from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss
-
-
-class _FSDPForwardRedirection:
-    """
-    Modified based on
-    https://github.com/Lightning-AI/pytorch-lightning/blob/d3f9c83d6efa4f1def36aa6c199600946cdb9117/src/lightning/pytorch/strategies/strategy.py#L601-L648
-    Redirect a method call through FullyShardedDataParallel.forward so that the FSDP module's root pre-forward and
-    post-forward can be properly executed around the method call.
-    This is needed in cases where we call a submodule of a FSDP module. For instance, when we want to call only
-    the `LlamaModel` part out of a FSDP-wrapped `LlamaForCausalLM` to get the hidden states without involving
-    GPU-memory-heavy `lm_head` and cross entropy computation, doing this directly (i.e. `model.model.forward()`)
-    will not work because the first `nn.Embedding` layer is not independently wrapped as a FSDP module (because of
-    the transformer-based wrapping policy), and not calling it through FSDP root module forward will not all-gather
-    its parameter, thus resulting in "RuntimeError: 'weight' must be 2-D" error. Similarly, if we want to call just
-    the `lm_head` part of a model, we need this trick too to properly get its params all-gathered.
-    """
-
-    def __call__(
-        self,
-        wrapper_module: FullyShardedDataParallel,
-        method: Callable,
-        *args: Any,
-        **kwargs: Any,
-    ):
-        """Reroutes a method call through the `wrapper_module`'s `forward` method.
-        Args:
-            wrapper_module: The module that has `original_module` wrapped.
-            original_module: The module that was wrapped inside `wrapper_module`.
-            method_name: The name of the method that should be called on the `original_module` after inputs get
-                redirected through the `wrapper_module`'s `forward` method.
-            *args: The positional arguments to the method `method_name`. They will get passed to a patched
-                `forward` method instead.
-            **kwargs: The keyword arguments to the method `method_name`. They will get passed to a patched
-                `forward` method instead.
-        """
-        assert isinstance(wrapper_module, FullyShardedDataParallel)
-        original_module = wrapper_module._fsdp_wrapped_module
-        original_forward = original_module.forward
-
-        def wrapped_forward(*_args: Any, **_kwargs: Any) -> Any:
-            # Unpatch ourselves immediately before calling the method `method_name`
-            # because itself may want to call the real `forward`
-            original_module.forward = original_forward  # type: ignore[method-assign]
-            # Call the actual method e.g. `.training_step(...)`
-            out = method(*_args, **_kwargs)
-            return out
-
-        # Patch the original_module's forward so we can redirect the arguments back to the real method
-        original_module.forward = wrapped_forward  # type: ignore[method-assign]
-        wrapper_output = wrapper_module(*args, **kwargs)
-        return wrapper_output
+from liger_kernel.transformers.fsdp import _FSDPForwardRedirection
 
 
 class LigerORPOTrainer(ORPOTrainer):

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/src/liger_kernel/utils.py

@@ -1,6 +1,17 @@
+try:
+    import peft  # noqa: F401
+
+    PEFT_AVAILABLE = True
+except ImportError:
+    PEFT_AVAILABLE = False
+
 import torch
 
 
+def is_peft_available():
+    return PEFT_AVAILABLE
+
+
 def infer_device():
     """
     Get current device name based on available devices

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.9.dev20250519015630 → liger_kernel_nightly-0.5.9.dev20250519025610}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -146,6 +146,7 @@ src/liger_kernel/transformers/__init__.py
 src/liger_kernel/transformers/auto_model.py
 src/liger_kernel/transformers/cross_entropy.py
 src/liger_kernel/transformers/dyt.py
+src/liger_kernel/transformers/fsdp.py
 src/liger_kernel/transformers/functional.py
 src/liger_kernel/transformers/fused_linear_cross_entropy.py
 src/liger_kernel/transformers/fused_linear_jsd.py