liger-kernel-nightly 0.5.4.dev20250305231637__tar.gz → 0.5.4.dev20250307064336__tar.gz

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/PKG-INFO

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

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.4.dev20250305231637"
+version = "0.5.4.dev20250307064336"
 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.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/cpo_loss.py

@@ -53,7 +53,27 @@ class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
         label_smoothing=0.0,
         compute_nll_loss=True,
         compiled=True,
+        average_log_prob=False,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with CPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            alpha (float): Weight for the alpha parameter
+            label_smoothing (float): Label smoothing factor
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            compiled (bool): Whether to use torch compile
+            average_log_prob (bool): Whether to average the log probability per non-masked token
+            chunk_size (int): Size of chunks for processing.
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
             cls=cls,
             ctx=ctx,
@@ -66,14 +86,15 @@ class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
             beta=beta,
             label_smoothing=label_smoothing,
             compute_nll_loss=compute_nll_loss,
-            average_log_prob=False,
+            average_log_prob=average_log_prob,
             compiled=compiled,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearCPOLoss(torch.nn.Module):
@@ -89,11 +110,19 @@ class LigerFusedLinearCPOLoss(torch.nn.Module):
         label_smoothing: float = 0.0,
         compute_nll_loss: bool = True,
         compiled: bool = True,
+        average_log_prob: bool = False,
+        chunk_size: int = 1,
     ):
         """
         Args:
             ignore_index (int): Index to ignore in the loss.
             beta (float): Weight for the odds ratio loss.
+            alpha (float): Weight for the alpha parameter.
+            label_smoothing (float): Label smoothing factor.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            average_log_prob (bool): Whether to average the log probability per non-masked token.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
@@ -102,8 +131,16 @@ class LigerFusedLinearCPOLoss(torch.nn.Module):
         self.label_smoothing = label_smoothing
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
+        self.average_log_prob = average_log_prob
+        self.chunk_size = chunk_size
 
-    def forward(self, lin_weight, _input, target, bias=None):
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+    ):
         return LigerFusedLinearCPOFunction.apply(
             _input,
             lin_weight,
@@ -115,4 +152,6 @@ class LigerFusedLinearCPOLoss(torch.nn.Module):
             self.label_smoothing,
             self.compute_nll_loss,
             self.compiled,
+            self.average_log_prob,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/dpo_loss.py

@@ -68,7 +68,27 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=False,
         compiled=True,
         use_ref_model=True,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with DPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size * seq_len, hidden_size)
+            ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size)
+            ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            compiled (bool): Whether to use torch compile
+            use_ref_model (bool): Whether to use a reference model
+            chunk_size (int): Size of chunks for processing.
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
             cls=cls,
             ctx=ctx,
@@ -84,12 +104,13 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
             ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearDPOLoss(torch.nn.Module):
@@ -104,6 +125,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
         compute_nll_loss: bool = False,
         compiled: bool = True,
         use_ref_model: bool = True,
+        chunk_size: int = 1,
     ):
         """
         Args:
@@ -112,6 +134,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
             compute_nll_loss (bool): Whether to compute the NLL loss.
             compiled (bool): Whether to use the torch compiled kernel.
             use_ref_model (bool): Whether to use a reference model for the DPO loss.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
@@ -119,6 +142,7 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
         self.use_ref_model = use_ref_model
+        self.chunk_size = chunk_size
 
     def forward(
         self,
@@ -143,4 +167,5 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             self.compiled,
             self.use_ref_model,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/fused_linear_rlhf.py

@@ -29,8 +29,27 @@ class LigerFusedLinearRLHFBase(torch.autograd.Function):
         ref_input=None,
         ref_weight=None,
         ref_bias=None,
+        chunk_size=1,
     ):
-        """Chunked forward pass for RLHF loss computation."""
+        """Chunked forward pass for RLHF loss computation.
+
+        Args:
+            cls: The class
+            ctx: Context for backward
+            _input: Input tensor
+            weight: Weight tensor
+            attention_mask: Attention mask tensor
+            rewards: Rewards tensor
+            bias: Bias tensor
+            num_generations: Number of generations per prompt
+            beta: Weight for the KL penalty
+            compiled: Whether to use torch compile
+            use_ref_model: Whether to use a reference model
+            ref_input: Reference model input tensor
+            ref_weight: Reference model weight tensor
+            ref_bias: Reference model bias tensor
+            chunk_size: Size of chunks for processing in other loss modules
+        """
         # Save for backward
         ctx.beta = beta
         ctx.rewards = rewards
@@ -106,7 +125,7 @@ class LigerFusedLinearRLHFBase(torch.autograd.Function):
         if compiled:
             accumulate_chunk = torch.compile(accumulate_chunk)
 
-        # Process input in chunks
+        # Process input in chunks based on num_generations
         chunks = max(1, _input.shape[0] // num_generations)
         _input_chunks = torch.chunk(_input, chunks=chunks, dim=0)
         _attention_mask_chunks = torch.chunk(attention_mask, chunks=chunks, dim=0)
@@ -210,11 +229,12 @@ class LigerFusedLinearRLHFBase(torch.autograd.Function):
             None,  # grad_attention_mask
             None,  # grad_rewards
             grad_bias,
-            None,  # grad_chunk_size
+            None,  # grad_num_generations
             None,  # grad_beta
             None,  # grad_compiled
             None,  # grad_use_ref_model
             None,  # grad_ref_input
             None,  # grad_ref_weight
             None,  # grad_ref_bias
+            None,  # grad_chunk_size
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/grpo_loss.py

@@ -79,7 +79,27 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
         compiled=True,
         use_ref_model=True,
         num_generations=1,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with GRPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            attention_mask (torch.Tensor): Attention mask tensor. Shape: (batch_size, seq_len)
+            rewards (torch.Tensor): Rewards tensor. Shape: (batch_size,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size * seq_len, hidden_size)
+            ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size)
+            ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,)
+            beta (float): Weight for the KL penalty
+            compiled (bool): Whether to use torch compile
+            use_ref_model (bool): Whether to use a reference model
+            num_generations (int): Number of generations per prompt
+            chunk_size (int): Size of chunks for processing.
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
             cls=cls,
             ctx=ctx,
@@ -95,6 +115,7 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             compiled=compiled,
             use_ref_model=use_ref_model,
             num_generations=num_generations,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
@@ -115,6 +136,7 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             None,  # grad_compiled
             None,  # grad_use_ref_model
             None,  # grad_num_generations
+            None,  # grad_chunk_size
         )
 
 
@@ -127,12 +149,22 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
         compiled: bool = True,
         use_ref_model: bool = True,
         num_generations: int = 1,
+        chunk_size: int = 1,
     ):
+        """
+        Args:
+            beta (float): Weight for the KL penalty.
+            compiled (bool): Whether to use torch compile.
+            use_ref_model (bool): Whether to use a reference model.
+            num_generations (int): Number of generations per prompt.
+            chunk_size (int): Size of chunks for processing.
+        """
         super().__init__()
         self.beta = beta
         self.compiled = compiled
         self.use_ref_model = use_ref_model
         self.num_generations = num_generations
+        self.chunk_size = chunk_size
 
     def forward(
         self,
@@ -158,4 +190,5 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
             self.compiled,
             self.use_ref_model,
             self.num_generations,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/jsd_loss.py

@@ -47,6 +47,7 @@ class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
         ignore_index: int = -100,
         temperature: float = 1.0,
         compiled: bool = True,
+        chunk_size: int = 1024,
     ):
         """
         Fused linear layer with JSD distillation loss.
@@ -62,6 +63,7 @@ class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
             ignore_index (int): Index to ignore in loss computation
             temperature (float): Temperature for softening/sharpening distributions
             compiled (bool): Whether to use torch compile
+            chunk_size (int): Size of chunks for processing.
         Returns:
             torch.Tensor: Computed loss
         """
@@ -75,7 +77,7 @@ class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
             target=true_labels,
             student_bias=student_bias,
             teacher_bias=teacher_bias,
-            chunk_size=1024,
+            chunk_size=chunk_size,
             weight_hard_loss=weight_hard_loss,
             weight_soft_loss=weight_soft_loss,
             beta=beta,
@@ -97,6 +99,7 @@ class LigerFusedLinearJSDFunction(LigerFusedLinearDistillationBase):
             None,  # ignore_index
             None,  # temperature
             None,  # compiled
+            None,  # chunk_size
         )
 
 
@@ -113,6 +116,7 @@ class LigerFusedLinearJSDLoss(torch.nn.Module):
         ignore_index: int = -100,
         temperature: float = 1.0,
         compiled: bool = True,
+        chunk_size: int = 1024,
     ):
         """
         Args:
@@ -122,6 +126,7 @@ class LigerFusedLinearJSDLoss(torch.nn.Module):
             temperature (float): Temperature for softening distributions
             compiled (bool): Whether to use torch compile
             beta (float): Coefficient beta of generalized JSD in the interval [0, 1]. Default: `0.5`.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         assert temperature != 0, "Temperature cannot be 0."
@@ -131,6 +136,7 @@ class LigerFusedLinearJSDLoss(torch.nn.Module):
         self.temperature = temperature
         self.compiled = compiled
         self.beta = beta
+        self.chunk_size = chunk_size
 
     def forward(
         self,
@@ -169,4 +175,5 @@ class LigerFusedLinearJSDLoss(torch.nn.Module):
             self.ignore_index,
             self.temperature,
             self.compiled,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/kto_loss.py

@@ -86,7 +86,29 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
         compiled=True,
         use_ref_model=True,
         average_log_prob=False,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with KTO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            preference_labels (torch.Tensor): Preference labels tensor. Shape: (batch_size,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size * seq_len, hidden_size)
+            ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size)
+            ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,)
+            kl (torch.Tensor, optional): KL divergence tensor. Shape: (batch_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Temperature parameter for the KTO loss
+            compiled (bool): Whether to use torch compile
+            use_ref_model (bool): Whether to use a reference model
+            average_log_prob (bool): Whether to average the log probability per non-masked token
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
             cls=cls,
             ctx=ctx,
@@ -104,6 +126,7 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
             ref_bias=ref_bias,
             average_log_prob=average_log_prob,
             kl=kl,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
@@ -121,6 +144,7 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -136,6 +160,7 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
         compiled: bool = True,
         use_ref_model: bool = False,
         average_log_prob: bool = False,
+        chunk_size: int = 1,
     ):
         """
         Args:
@@ -143,7 +168,8 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
             beta (float): Temperature parameter for the KTO loss
             compiled (bool): Whether to use compiled operations
             use_ref_model (bool): Whether to use a reference model for the DPO loss.
-            average_log_prob (bool): Whether to average the log probability per non-masked token.
+            average_log_prob (bool): Whether to average the log probability per non-masked token
+            chunk_size (int): Size of chunks for processing
         """
         super().__init__()
         self.ignore_index = ignore_index
@@ -151,6 +177,7 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
         self.compiled = compiled
         self.use_ref_model = use_ref_model
         self.average_log_prob = average_log_prob
+        self.chunk_size = chunk_size
 
     def forward(
         self,
@@ -179,4 +206,5 @@ class LigerFusedLinearKTOLoss(torch.nn.Module):
             self.compiled,
             self.use_ref_model,
             self.average_log_prob,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/orpo_loss.py

@@ -55,7 +55,24 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=True,
         nll_target=None,
         compiled=True,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with ORPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            nll_target (torch.LongTensor, optional): Target tensor for NLL loss. Shape: (batch_size * seq_len,)
+            compiled (bool): Whether to use torch compile
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
             cls=cls,
             ctx=ctx,
@@ -68,12 +85,13 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
             compute_nll_loss=compute_nll_loss,
             nll_target=nll_target,
             compiled=compiled,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None
+        return *grads, None, None, None, None, None, None
 
 
 class LigerFusedLinearORPOLoss(torch.nn.Module):
@@ -87,19 +105,31 @@ class LigerFusedLinearORPOLoss(torch.nn.Module):
         beta: float = 0.1,
         compute_nll_loss: bool = True,
         compiled: bool = True,
+        chunk_size: int = 1,
     ):
         """
         Args:
             ignore_index (int): Index to ignore in the loss.
             beta (float): Weight for the odds ratio loss.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
         self.beta = beta
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
+        self.chunk_size = chunk_size
 
-    def forward(self, lin_weight, _input, target, bias=None, nll_target=None):
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+        nll_target=None,
+    ):
         return LigerFusedLinearORPOFunction.apply(
             _input,
             lin_weight,
@@ -110,4 +140,5 @@ class LigerFusedLinearORPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             nll_target,
             self.compiled,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel/chunked_loss/simpo_loss.py

@@ -62,27 +62,47 @@ class LigerFusedLinearSimPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=False,
         compiled=True,
         gamma=0.5,
+        chunk_size=1,
     ):
+        """
+        Fused linear layer with SimPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            alpha (float): Weight for the alpha parameter
+            label_smoothing (float): Label smoothing factor
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            compiled (bool): Whether to use torch compile
+            gamma (float): Weight for the gamma parameter
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
         return super().forward(
-            cls,
-            ctx,
-            _input,
-            weight,
-            target,
-            bias,
-            compute_nll_loss=compute_nll_loss,
+            cls=cls,
+            ctx=ctx,
+            _input=_input,
+            weight=weight,
+            target=target,
+            bias=bias,
             ignore_index=ignore_index,
             alpha=alpha,
             beta=beta,
             label_smoothing=label_smoothing,
+            compute_nll_loss=compute_nll_loss,
             compiled=compiled,
             gamma=gamma,
+            chunk_size=chunk_size,
         )
 
     @staticmethod
     def backward(ctx, *grad_output):
         grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
-        return *grads, None, None, None, None, None, None, None
+        return *grads, None, None, None, None, None, None, None, None
 
 
 class LigerFusedLinearSimPOLoss(torch.nn.Module):
@@ -99,11 +119,18 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
         compute_nll_loss: bool = True,
         compiled: bool = True,
         gamma: float = 0.5,
+        chunk_size: int = 1,
     ):
         """
         Args:
             ignore_index (int): Index to ignore in the loss.
             beta (float): Weight for the odds ratio loss.
+            alpha (float): Weight for the alpha parameter.
+            label_smoothing (float): Label smoothing factor.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            gamma (float): Weight for the gamma parameter.
+            chunk_size (int): Size of chunks for processing.
         """
         super().__init__()
         self.ignore_index = ignore_index
@@ -113,8 +140,15 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
         self.compute_nll_loss = compute_nll_loss
         self.compiled = compiled
         self.gamma = gamma
+        self.chunk_size = chunk_size
 
-    def forward(self, lin_weight, _input, target, bias=None):
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+    ):
         return LigerFusedLinearSimPOFunction.apply(
             _input,
             lin_weight,
@@ -127,4 +161,5 @@ class LigerFusedLinearSimPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             self.compiled,
             self.gamma,
+            self.chunk_size,
         )

{liger_kernel_nightly-0.5.4.dev20250305231637 → liger_kernel_nightly-0.5.4.dev20250307064336}/src/liger_kernel_nightly.egg-info/PKG-INFO

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