liger-kernel 0.5.2__py3-none-any.whl → 0.5.3__py3-none-any.whl

liger_kernel/transformers/monkey_patch.py

@@ -1,9 +1,11 @@
 import inspect
 import logging
+
 from functools import partial
 from typing import Callable
 
 import transformers
+
 from packaging import version
 from transformers import PreTrainedModel
 
@@ -12,38 +14,24 @@ from liger_kernel.transformers.functional import liger_cross_entropy
 from liger_kernel.transformers.geglu import LigerGEGLUMLP
 from liger_kernel.transformers.layer_norm import LigerLayerNorm
 from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forward
-from liger_kernel.transformers.model.gemma import (
-    lce_forward_deprecated as gemma_lce_forward_deprecated,
-)
+from liger_kernel.transformers.model.gemma import lce_forward_deprecated as gemma_lce_forward_deprecated
 from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_forward
-from liger_kernel.transformers.model.gemma2 import (
-    lce_forward_deprecated as gemma2_lce_forward_deprected,
-)
+from liger_kernel.transformers.model.gemma2 import lce_forward_deprecated as gemma2_lce_forward_deprected
 from liger_kernel.transformers.model.llama import lce_forward as llama_lce_forward
-from liger_kernel.transformers.model.llama import (
-    lce_forward_deprecated as llama_lce_forward_deprecated,
-)
+from liger_kernel.transformers.model.llama import lce_forward_deprecated as llama_lce_forward_deprecated
 from liger_kernel.transformers.model.mistral import lce_forward as mistral_lce_forward
 from liger_kernel.transformers.model.mixtral import lce_forward as mixtral_lce_forward
-from liger_kernel.transformers.model.mixtral import (
-    lce_forward_deprecated as mixtral_lce_forward_deprecated,
-)
+from liger_kernel.transformers.model.mixtral import lce_forward_deprecated as mixtral_lce_forward_deprecated
 from liger_kernel.transformers.model.phi3 import lce_forward as phi3_lce_forward
-from liger_kernel.transformers.model.phi3 import (
-    lce_forward_deprecated as phi3_lce_forward_deprecated,
-)
+from liger_kernel.transformers.model.phi3 import lce_forward_deprecated as phi3_lce_forward_deprecated
 from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward
-from liger_kernel.transformers.model.qwen2 import (
-    lce_forward_deprecated as qwen2_lce_forward_deprecated,
-)
+from liger_kernel.transformers.model.qwen2 import lce_forward_deprecated as qwen2_lce_forward_deprecated
 from liger_kernel.transformers.qwen2vl_mrope import liger_multimodal_rotary_pos_emb
 from liger_kernel.transformers.rms_norm import LigerRMSNorm
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
-from liger_kernel.transformers.swiglu import (
-    LigerBlockSparseTop2MLP,
-    LigerPhi3SwiGLUMLP,
-    LigerSwiGLUMLP,
-)
+from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP
+from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP
+from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
 
 transformer_version = version.parse(transformers.__version__)
 
@@ -57,23 +45,17 @@ def _bind_method_to_module(module, method_name: str, new_method: Callable):
     module.__dict__[method_name] = new_method.__get__(module, module.__class__)
 
 
-def _patch_rms_norm_module(
-    module, offset=0.0, eps=1e-6, casting_mode="llama", in_place=True
-):
+def _patch_rms_norm_module(module, offset=0.0, eps=1e-6, casting_mode="llama", in_place=True):
     module.offset = offset
     module.casting_mode = casting_mode
-    module.variance_epsilon = (
-        getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
-    )
+    module.variance_epsilon = getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
     module.in_place = in_place
     _bind_method_to_module(module, "forward", LigerRMSNorm.forward)
     _bind_method_to_module(module, "extra_repr", LigerRMSNorm.extra_repr)
 
 
 def _patch_layer_norm_module(module, eps=1e-6):
-    module.variance_epsilon = (
-        getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
-    )
+    module.variance_epsilon = getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
     module.hidden_size = module.normalized_shape
     _bind_method_to_module(module, "forward", LigerLayerNorm.forward)
     _bind_method_to_module(module, "extra_repr", LigerLayerNorm.extra_repr)
@@ -145,9 +127,7 @@ def apply_liger_kernel_to_llama(
 
         for decoder_layer in base_model.layers:
             if swiglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -184,17 +164,13 @@ def apply_liger_kernel_to_mllama(
     ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
 
     from transformers.models.mllama import modeling_mllama
-    from transformers.models.mllama.modeling_mllama import (
-        MllamaForCausalLM,
-        MllamaForConditionalGeneration,
-        MllamaTextModel,
-        MllamaVisionModel,
-    )
+    from transformers.models.mllama.modeling_mllama import MllamaForCausalLM
+    from transformers.models.mllama.modeling_mllama import MllamaForConditionalGeneration
+    from transformers.models.mllama.modeling_mllama import MllamaTextModel
+    from transformers.models.mllama.modeling_mllama import MllamaVisionModel
 
     from liger_kernel.transformers.model.mllama import lce_forward as mllama_lce_forward
-    from liger_kernel.transformers.model.mllama import (
-        lce_forward_deprecated as mllama_lce_forward_deprecated,
-    )
+    from liger_kernel.transformers.model.mllama import lce_forward_deprecated as mllama_lce_forward_deprecated
 
     if rope:
         modeling_mllama.apply_rotary_pos_emb = liger_rotary_pos_emb
@@ -241,9 +217,7 @@ def apply_liger_kernel_to_mllama(
                 _patch_rms_norm_module(text_model.norm)
             for decoder_layer in text_model.layers:
                 if swiglu:
-                    _bind_method_to_module(
-                        decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
-                    )
+                    _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
                 if rms_norm:
                     _patch_rms_norm_module(decoder_layer.input_layernorm)
                     _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -317,9 +291,7 @@ def apply_liger_kernel_to_mistral(
 
         for decoder_layer in base_model.layers:
             if swiglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -391,9 +363,7 @@ def apply_liger_kernel_to_mixtral(
         for decoder_layer in base_model.layers:
             if swiglu:
                 for expert in decoder_layer.block_sparse_moe.experts:
-                    _bind_method_to_module(
-                        expert, "forward", LigerBlockSparseTop2MLP.forward
-                    )
+                    _bind_method_to_module(expert, "forward", LigerBlockSparseTop2MLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -431,12 +401,8 @@ def apply_liger_kernel_to_gemma(
     from transformers.models.gemma.modeling_gemma import GemmaModel
 
     # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/gemma/modeling_gemma.py#L109
-    LigerRMSNormForGemma = partial(
-        LigerRMSNorm, offset=1.0, init_fn="zeros", casting_mode="gemma"
-    )
-    _patch_rms_norm_module_for_gemma = partial(
-        _patch_rms_norm_module, casting_mode="gemma", offset=1.0
-    )
+    LigerRMSNormForGemma = partial(LigerRMSNorm, offset=1.0, init_fn="zeros", casting_mode="gemma")
+    _patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode="gemma", offset=1.0)
 
     if rope:
         modeling_gemma.apply_rotary_pos_emb = liger_rotary_pos_emb
@@ -471,9 +437,7 @@ def apply_liger_kernel_to_gemma(
 
         for decoder_layer in base_model.layers:
             if geglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerGEGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerGEGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module_for_gemma(decoder_layer.input_layernorm)
                 _patch_rms_norm_module_for_gemma(decoder_layer.post_attention_layernorm)
@@ -510,9 +474,7 @@ def apply_liger_kernel_to_gemma2(
     from transformers.models.gemma2 import modeling_gemma2
     from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
 
-    LigerRMSNormForGemma2 = partial(
-        LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False
-    )
+    LigerRMSNormForGemma2 = partial(LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False)
     _patch_rms_norm_module_for_gemma2 = partial(
         _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
     )
@@ -551,20 +513,12 @@ def apply_liger_kernel_to_gemma2(
 
         for decoder_layer in base_model.layers:
             if geglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerGEGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerGEGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module_for_gemma2(decoder_layer.input_layernorm)
-                _patch_rms_norm_module_for_gemma2(
-                    decoder_layer.post_attention_layernorm
-                )
-                _patch_rms_norm_module_for_gemma2(
-                    decoder_layer.pre_feedforward_layernorm
-                )
-                _patch_rms_norm_module_for_gemma2(
-                    decoder_layer.post_feedforward_layernorm
-                )
+                _patch_rms_norm_module_for_gemma2(decoder_layer.post_attention_layernorm)
+                _patch_rms_norm_module_for_gemma2(decoder_layer.pre_feedforward_layernorm)
+                _patch_rms_norm_module_for_gemma2(decoder_layer.post_feedforward_layernorm)
 
 
 def apply_liger_kernel_to_qwen2(
@@ -633,9 +587,7 @@ def apply_liger_kernel_to_qwen2(
 
         for decoder_layer in base_model.layers:
             if swiglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -674,14 +626,10 @@ def apply_liger_kernel_to_qwen2_vl(
     from transformers.models.qwen2_vl import modeling_qwen2_vl
     from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLModel
 
-    from liger_kernel.transformers.model.qwen2_vl import (
-        lce_forward as qwen2_vl_lce_forward,
-    )
+    from liger_kernel.transformers.model.qwen2_vl import lce_forward as qwen2_vl_lce_forward
 
     if rope:
-        modeling_qwen2_vl.apply_multimodal_rotary_pos_emb = (
-            liger_multimodal_rotary_pos_emb
-        )
+        modeling_qwen2_vl.apply_multimodal_rotary_pos_emb = liger_multimodal_rotary_pos_emb
     if rms_norm:
         # https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L439
         modeling_qwen2_vl.Qwen2RMSNorm = LigerRMSNorm
@@ -712,9 +660,7 @@ def apply_liger_kernel_to_qwen2_vl(
             _patch_rms_norm_module(base_model.norm)
         for decoder_layer in base_model.layers:
             if swiglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -783,9 +729,7 @@ def apply_liger_kernel_to_phi3(
 
         for decoder_layer in base_model.layers:
             if swiglu:
-                _bind_method_to_module(
-                    decoder_layer.mlp, "forward", LigerPhi3SwiGLUMLP.forward
-                )
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerPhi3SwiGLUMLP.forward)
             if rms_norm:
                 _patch_rms_norm_module(decoder_layer.input_layernorm)
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
@@ -826,24 +770,16 @@ def _apply_liger_kernel(model_type: str, **kwargs) -> None:
         return
 
     if model_type not in MODEL_TYPE_TO_APPLY_LIGER_FN.keys():
-        logger.info(
-            f"There are currently no Liger kernels supported for model type: {model_type}."
-        )
+        logger.info(f"There are currently no Liger kernels supported for model type: {model_type}.")
         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
-    applicable_kwargs = {
-        key: value
-        for key, value in kwargs.items()
-        if key in apply_fn_signature.parameters
-    }
+    applicable_kwargs = {key: value for key, value in kwargs.items() if key in apply_fn_signature.parameters}
 
-    logger.info(
-        f"Applying Liger kernels for model type: {model_type} with kwargs: {applicable_kwargs}"
-    )
+    logger.info(f"Applying Liger kernels for model type: {model_type} with kwargs: {applicable_kwargs}")
 
     # Assume this is invoked pre-model initialization, so we only need to patch transformers code
     apply_fn(**applicable_kwargs)
@@ -857,20 +793,14 @@ def _apply_liger_kernel_to_instance(model: PreTrainedModel, **kwargs) -> None:
         - model: the model instance to apply Liger kernels to
         - kwargs: keyword arguments that are passed to the corresponding apply_liger_kernel_to_* function.
     """
-    model_type = getattr(model, "config", None) and getattr(
-        model.config, "model_type", None
-    )
+    model_type = getattr(model, "config", None) and getattr(model.config, "model_type", None)
 
     if not model_type:
-        logger.info(
-            "Model type could not be determined from model config. No Liger kernels will be applied."
-        )
+        logger.info("Model type could not be determined from model config. No Liger kernels will be applied.")
         return
 
     if model_type not in MODEL_TYPE_TO_APPLY_LIGER_FN.keys():
-        logger.info(
-            f"There are currently no Liger kernels supported for model type: {model_type}."
-        )
+        logger.info(f"There are currently no Liger kernels supported for model type: {model_type}.")
         return
 
     apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type]
@@ -878,11 +808,7 @@ def _apply_liger_kernel_to_instance(model: PreTrainedModel, **kwargs) -> None:
     apply_fn_signature = inspect.signature(apply_fn)
 
     # Filter out the keyword arguments that are not supported by the apply function
-    applicable_kwargs = {
-        key: value
-        for key, value in kwargs.items()
-        if key in apply_fn_signature.parameters
-    }
+    applicable_kwargs = {key: value for key, value in kwargs.items() if key in apply_fn_signature.parameters}
     logger.info(
         f"Applying Liger kernels to model instance with model type: {model_type} with kwargs: {applicable_kwargs}"
     )

liger_kernel/transformers/rms_norm.py

@@ -19,9 +19,7 @@ class LigerRMSNorm(nn.Module):
             "ones",
             "zeros",
         ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
-        self.weight = nn.Parameter(
-            torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size)
-        )
+        self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
         self.variance_epsilon, self.offset, self.casting_mode, self.in_place = (
             eps,
             offset,
@@ -40,4 +38,6 @@ class LigerRMSNorm(nn.Module):
         )
 
     def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
+        return (
+            f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
+        )

liger_kernel/transformers/rope.py

@@ -8,8 +8,8 @@ def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     Args:
         q (torch.Tensor): The query tensor of shape (bsz, n_q_head, seq_len, head_dim).
         k (torch.Tensor): The key tensor of shape (bsz, n_kv_head, seq_len, head_dim).
-        cos (torch.Tensor): The cosine tensor of shape (1, seq_len, head_dim).
-        sin (torch.Tensor): The sine tensor of shape (1, seq_len, head_dim).
+        cos (torch.Tensor): The cosine tensor of shape (1, seq_len, head_dim) or (bsz, seq_len, head_dim).
+        sin (torch.Tensor): The sine tensor of shape (1, seq_len, head_dim) or (bsz, seq_len, head_dim).
         position_ids (torch.Tensor, optional): The position ids tensor. Defaults to None.
         unsqueeze_dim (int, optional): The dimension to unsqueeze. Defaults to 1.
 

liger_kernel/transformers/swiglu.py

@@ -16,10 +16,7 @@ class LigerSwiGLUMLP(nn.Module):
             raise ValueError(f"Activation function {config.hidden_act} not supported.")
 
     def forward(self, x):
-
-        return self.down_proj(
-            LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x))
-        )
+        return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
 
 
 class LigerBlockSparseTop2MLP(nn.Module):
@@ -36,7 +33,6 @@ class LigerBlockSparseTop2MLP(nn.Module):
             raise ValueError(f"Activation function {config.hidden_act} not supported.")
 
     def forward(self, x):
-
         return self.w2(LigerSiLUMulFunction.apply(self.w1(x), self.w3(x)))
 
 
@@ -51,9 +47,7 @@ class LigerPhi3SwiGLUMLP(nn.Module):
         self.config = config
         self.hidden_size = config.hidden_size
         self.intermediate_size = config.intermediate_size
-        self.gate_up_proj = nn.Linear(
-            self.hidden_size, 2 * self.intermediate_size, bias=False
-        )
+        self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         if config.hidden_act not in ["silu", "swish"]:
             raise ValueError(f"Activation function {config.hidden_act} not supported.")

liger_kernel/transformers/trainer/__init__.py

@@ -1,6 +1,4 @@
 try:
-    from liger_kernel.transformers.trainer.orpo_trainer import (  # noqa: F401
-        LigerORPOTrainer,
-    )
+    from liger_kernel.transformers.trainer.orpo_trainer import LigerORPOTrainer  # noqa: F401
 except ImportError:
     raise ImportError("Please `pip install trl` to use LigerORPOTrainer")

liger_kernel/transformers/trainer/orpo_trainer.py

@@ -1,7 +1,14 @@
-from typing import Any, Callable, Dict, List, Literal, Tuple, Union
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Literal
+from typing import Tuple
+from typing import Union
 
 import torch
 import torch.nn as nn
+
 from torch.distributed.fsdp import FullyShardedDataParallel
 from trl.trainer import ORPOTrainer
 
@@ -17,7 +24,7 @@ class _FSDPForwardRedirection:
     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.Emebedding` layer is not independently wrapped as a FSDP module (because of
+    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.
@@ -62,9 +69,7 @@ class _FSDPForwardRedirection:
 class LigerORPOTrainer(ORPOTrainer):
     def concatenated_forward(
         self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]
-    ) -> Tuple[
-        torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor
-    ]:
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
         """
         Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
         We do this to avoid doing two forward passes, because it's faster for FSDP.
@@ -79,9 +84,7 @@ class LigerORPOTrainer(ORPOTrainer):
 
         model_kwargs = (
             {
-                "decoder_input_ids": self._shift_right(
-                    concatenated_batch["concatenated_labels"]
-                ),
+                "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]),
             }
             if self.is_encoder_decoder
             else {}
@@ -90,6 +93,13 @@ class LigerORPOTrainer(ORPOTrainer):
         if self.aux_loss_enabled:
             model_kwargs["output_router_logits"] = True
 
+        if self.is_encoder_decoder:
+            labels = concatenated_batch["concatenated_labels"].clone()
+        else:
+            labels = concatenated_batch["concatenated_input_ids"].clone()
+            attention_mask = concatenated_batch["concatenated_attention_mask"]
+            labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id)
+
         if isinstance(model, FullyShardedDataParallel):
             outputs = _FSDPForwardRedirection()(
                 model,
@@ -109,22 +119,27 @@ class LigerORPOTrainer(ORPOTrainer):
                 **model_kwargs,
             )
 
-        orpo_loss_fn = LigerFusedLinearORPOLoss(
-            ignore_index=self.label_pad_token_id, beta=self.beta
-        )
+        orpo_loss_fn = LigerFusedLinearORPOLoss(ignore_index=self.label_pad_token_id, beta=self.beta)
 
-        def orpo_partial(lm_head, last_hidden_state, concatenated_labels):
+        def orpo_partial(lm_head, last_hidden_state, concatenated_labels, nll_target):
             return orpo_loss_fn(
-                lm_head.weight, last_hidden_state, concatenated_labels, lm_head.bias
+                lm_head.weight, last_hidden_state, concatenated_labels, lm_head.bias, nll_target=nll_target
             )
 
         orpo_loss, aux_outputs = _FSDPForwardRedirection()(
             model,
             orpo_partial,
             model.lm_head,
-            outputs.last_hidden_state,
-            concatenated_batch["concatenated_labels"],
+            outputs.last_hidden_state[:, :-1] if not self.is_encoder_decoder else outputs.last_hidden_state,
+            concatenated_batch["concatenated_labels"][:, 1:]
+            if not self.is_encoder_decoder
+            else concatenated_batch["concatenated_labels"],
+            labels[:, 1:] if not self.is_encoder_decoder else labels,
         )
+        # if aux_loss_enabled, add the aux_loss to the orpo_loss
+        if self.aux_loss_enabled:
+            orpo_loss += self.aux_loss_coef * outputs.aux_loss
+
         return orpo_loss, aux_outputs
 
     def get_batch_loss_metrics(
@@ -145,9 +160,7 @@ class LigerORPOTrainer(ORPOTrainer):
         ) = aux_outputs[:5]
 
         # return loss, metrics
-        chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = aux_outputs[
-            5:
-        ]
+        chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = aux_outputs[5:]
 
         reward_accuracies = (chosen_rewards > rejected_rewards).float()
 

liger_kernel/triton/__init__.py

@@ -1,3 +1 @@
-from liger_kernel.triton.monkey_patch import (  # noqa: F401
-    apply_liger_triton_cache_manager,
-)
+from liger_kernel.triton.monkey_patch import apply_liger_triton_cache_manager  # noqa: F401

liger_kernel/triton/monkey_patch.py

@@ -37,6 +37,4 @@ def apply_liger_triton_cache_manager():
     Experimental feature to get around transient FileNotFoundError in triton compilation.
     For more details please see https://github.com/triton-lang/triton/pull/4295
     """
-    os.environ["TRITON_CACHE_MANAGER"] = (
-        "liger_kernel.triton.monkey_patch:LigerTritonFileCacheManager"
-    )
+    os.environ["TRITON_CACHE_MANAGER"] = "liger_kernel.triton.monkey_patch:LigerTritonFileCacheManager"

{liger_kernel-0.5.2.dist-info → liger_kernel-0.5.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: liger_kernel
-Version: 0.5.2
+Version: 0.5.3
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -32,10 +32,6 @@ License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: torch>=2.1.2
 Requires-Dist: triton>=2.3.1
-Provides-Extra: transformers
-Requires-Dist: transformers~=4.0; extra == "transformers"
-Provides-Extra: trl
-Requires-Dist: trl>=0.11.0; extra == "trl"
 Provides-Extra: dev
 Requires-Dist: transformers>=4.44.2; extra == "dev"
 Requires-Dist: matplotlib>=3.7.2; extra == "dev"
@@ -46,13 +42,11 @@ Requires-Dist: pytest>=7.1.2; extra == "dev"
 Requires-Dist: pytest-xdist; extra == "dev"
 Requires-Dist: pytest-rerunfailures; extra == "dev"
 Requires-Dist: datasets>=2.19.2; extra == "dev"
-Requires-Dist: torchvision>=0.16.2; extra == "dev"
 Requires-Dist: seaborn; extra == "dev"
-Provides-Extra: amd
-Requires-Dist: torch>=2.6.0.dev; extra == "amd"
-Requires-Dist: setuptools-scm>=8; extra == "amd"
-Requires-Dist: torchvision>=0.20.0.dev; extra == "amd"
-Requires-Dist: triton>=3.0.0; extra == "amd"
+Requires-Dist: mkdocs; extra == "dev"
+Requires-Dist: mkdocs-material; extra == "dev"
+Dynamic: provides-extra
+Dynamic: requires-dist
 
 <a name="readme-top"></a>
 
@@ -116,7 +110,8 @@ Requires-Dist: triton>=3.0.0; extra == "amd"
 <details>
   <summary>Latest News 🔥</summary>
 
-  - [2024/12/15] We release LinkedIn Engineering Blog - [Liger-Kernel: Empowering an open source ecosystem of Triton Kernels for Efficient LLM Training](https://www.linkedin.com/blog/engineering/open-source/liger-kernel-open-source-ecosystem-for-efficient-llm-training)
+  - [2024/12/11] We release [v0.5.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.5.0): 80% more memory efficient post training losses (DPO, ORPO, CPO, etc)!
+  - [2024/12/5] We release LinkedIn Engineering Blog - [Liger-Kernel: Empowering an open source ecosystem of Triton Kernels for Efficient LLM Training](https://www.linkedin.com/blog/engineering/open-source/liger-kernel-open-source-ecosystem-for-efficient-llm-training)
   - [2024/11/6] We release [v0.4.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.4.0): Full AMD support, Tech Report, Modal CI, Llama-3.2-Vision!
   - [2024/10/21] We have released the tech report of Liger Kernel on Arxiv: https://arxiv.org/pdf/2410.10989
   - [2024/9/6] We release v0.2.1 ([X post](https://x.com/liger_kernel/status/1832168197002510649)). 2500+ Stars, 10+ New Contributors, 50+ PRs, 50k Downloads in two weeks!
@@ -128,7 +123,7 @@ Requires-Dist: triton>=3.0.0; extra == "amd"
 
 **Liger Kernel** is a collection of Triton kernels designed specifically for LLM training. It can effectively increase multi-GPU **training throughput by 20%** and reduces **memory usage by 60%**. We have implemented **Hugging Face Compatible** `RMSNorm`, `RoPE`, `SwiGLU`, `CrossEntropy`, `FusedLinearCrossEntropy`, and more to come. The kernel works out of the box with [Flash Attention](https://github.com/Dao-AILab/flash-attention), [PyTorch FSDP](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html), and [Microsoft DeepSpeed](https://github.com/microsoft/DeepSpeed). We welcome contributions from the community to gather the best kernels for LLM training.
 
-We've also added optimized Post-Training kernels that deliver **up to 80% memory savings** for alignment and distillation tasks. We support losses like DPO, CPO, ORPO, SimPO, JSD, and many more. 
+We've also added optimized Post-Training kernels that deliver **up to 80% memory savings** for alignment and distillation tasks. We support losses like DPO, CPO, ORPO, SimPO, JSD, and many more. Check out [how we optimize the memory](https://x.com/hsu_byron/status/1866577403918917655).
 
 ## Supercharge Your Model with Liger Kernel
 
@@ -145,6 +140,21 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
 > - Benchmark conditions: LLaMA 3-8B, Batch Size = 8, Data Type = `bf16`, Optimizer = AdamW, Gradient Checkpointing = True, Distributed Strategy = FSDP1 on 8 A100s.
 > - Hugging Face models start to OOM at a 4K context length, whereas Hugging Face + Liger Kernel scales up to 16K.
 
+## Optimize Post Training with Liger Kernel
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/linkedin/Liger-Kernel/main/docs/images/post-training.png" width="50%" alt="Post Training">
+</p>
+
+We provide optimized post training kernels like DPO, ORPO, SimPO, and more which can reduce memory usage by up to 80%. You can easily use them as python modules.
+
+```python
+from liger_kernel.chunked_loss import LigerFusedLinearDPOLoss
+orpo_loss = LigerFusedLinearORPOLoss()
+y = orpo_loss(lm_head.weight, x, target)
+```
+
+
 ## Examples
 
 | **Use Case**                                    | **Description**                                                                                   |
@@ -202,11 +212,13 @@ To install from source:
 ```bash
 git clone https://github.com/linkedin/Liger-Kernel.git
 cd Liger-Kernel
+
+# Install Default Dependencies
+# Setup.py will detect whether you are using AMD or NVIDIA
 pip install -e .
-# or if installing on amd platform
-pip install -e .[amd] --extra-index-url https://download.pytorch.org/whl/nightly/rocm6.2 # rocm6.2
-# or if using transformers
-pip install -e .[transformers]
+
+# Setup Development Dependencies
+pip install -e ".[dev]"
 ```
 
 
@@ -252,7 +264,7 @@ model = transformers.AutoModelForCausalLM("path/to/llama/model")
 
 ### 3. Compose Your Own Model
 
-You can take individual [kernels](#kernels) to compose your models.
+You can take individual [kernels](https://github.com/linkedin/Liger-Kernel?tab=readme-ov-file#model-kernels) to compose your models.
 
 ```python
 from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
@@ -291,7 +303,7 @@ loss.backward()
 | Gemma1      | `liger_kernel.transformers.apply_liger_kernel_to_gemma`    | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
 | Gemma2      | `liger_kernel.transformers.apply_liger_kernel_to_gemma2`   | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
 | Qwen2, Qwen2.5, & QwQ      | `liger_kernel.transformers.apply_liger_kernel_to_qwen2`    | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
-| Qwen2-VL       | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl`    | RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
+| Qwen2-VL, & QVQ       | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl`    | RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy        |
 | Phi3 & Phi3.5       | `liger_kernel.transformers.apply_liger_kernel_to_phi3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy         |
 
 
@@ -340,16 +352,17 @@ loss.backward()
 
 ## Contributing, Acknowledgements, and License
 
-- [Contributing Guidelines](https://github.com/linkedin/Liger-Kernel/blob/main/docs/CONTRIBUTING.md)
-- [Acknowledgements](https://github.com/linkedin/Liger-Kernel/blob/main/docs/Acknowledgement.md)
-- [License Information](https://github.com/linkedin/Liger-Kernel/blob/main/docs/License.md)
+- [Contributing Guidelines](https://github.com/linkedin/Liger-Kernel/blob/main/docs/contributing.md)
+- [Acknowledgements](https://github.com/linkedin/Liger-Kernel/blob/main/docs/acknowledgement.md)
+- [License Information](https://github.com/linkedin/Liger-Kernel/blob/main/docs/license.md)
 
 ## Sponsorship and Collaboration
- 
+
+- [Glows.ai](https://platform.glows.ai/): Sponsoring NVIDIA GPUs for our open source developers.
 - [AMD](https://www.amd.com/en.html): Providing AMD GPUs for our AMD CI.
 - [Intel](https://www.intel.com/): Providing Intel GPUs for our Intel CI.
 - [Modal](https://modal.com/): Free 3000 credits from GPU MODE IRL for our NVIDIA CI.
-- [EmbeddedLLM](https://embeddedllm.com/): Making Liger Kernel run fast and stable on AMD. 
+- [EmbeddedLLM](https://embeddedllm.com/): Making Liger Kernel run fast and stable on AMD.
 - [HuggingFace](https://huggingface.co/): Integrating Liger Kernel into Hugging Face Transformers and TRL.
 - [Lightning AI](https://lightning.ai/): Integrating Liger Kernel into Lightning Thunder.
 - [Axolotl](https://axolotl.ai/): Integrating Liger Kernel into Axolotl.