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

liger_kernel/chunked_loss/__init__.py

@@ -1,5 +1,6 @@
 from liger_kernel.chunked_loss.cpo_loss import LigerFusedLinearCPOLoss  # noqa: F401
 from liger_kernel.chunked_loss.dpo_loss import LigerFusedLinearDPOLoss  # noqa: F401
+from liger_kernel.chunked_loss.grpo_loss import LigerFusedLinearGRPOLoss  # noqa: F401
 from liger_kernel.chunked_loss.jsd_loss import LigerFusedLinearJSDLoss  # noqa: F401
 from liger_kernel.chunked_loss.kto_loss import LigerFusedLinearKTOLoss  # noqa: F401
 from liger_kernel.chunked_loss.orpo_loss import LigerFusedLinearORPOLoss  # noqa: F401

liger_kernel/chunked_loss/cpo_loss.py

@@ -39,8 +39,9 @@ class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
 
         return loss, chosen_rewards, rejected_rewards
 
-    @staticmethod
+    @classmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
@@ -52,27 +53,48 @@ class LigerFusedLinearCPOFunction(LigerFusedLinearPreferenceBase):
         label_smoothing=0.0,
         compute_nll_loss=True,
         compiled=True,
+        average_log_prob=False,
+        chunk_size=1,
     ):
-        return LigerFusedLinearPreferenceBase.forward(
-            ctx,
-            _input,
-            weight,
-            target,
-            bias,
-            loss_fn=LigerFusedLinearCPOFunction.preference_loss_fn,
+        """
+        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,
+            _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,
-            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):
@@ -88,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
@@ -101,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,
@@ -114,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/chunked_loss/dpo_loss.py

@@ -52,8 +52,9 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
         loss = -F.logsigmoid(logits_diff).sum() / (full_target.shape[0] // 2)
         return loss, chosen_rewards, rejected_rewards
 
-    @staticmethod
+    @classmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
@@ -67,14 +68,34 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
         compute_nll_loss=False,
         compiled=True,
         use_ref_model=True,
+        chunk_size=1,
     ):
-        return LigerFusedLinearPreferenceBase.forward(
+        """
+        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,
             _input=_input,
             weight=weight,
             target=target,
             bias=bias,
-            loss_fn=LigerFusedLinearDPOFunction.preference_loss_fn,
             ignore_index=ignore_index,
             beta=beta,
             compute_nll_loss=compute_nll_loss,
@@ -83,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):
@@ -103,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:
@@ -111,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
@@ -118,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,
@@ -142,4 +167,5 @@ class LigerFusedLinearDPOLoss(torch.nn.Module):
             self.compute_nll_loss,
             self.compiled,
             self.use_ref_model,
+            self.chunk_size,
         )

liger_kernel/chunked_loss/fused_linear_distillation.py

@@ -125,6 +125,7 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
 
     @staticmethod
     def forward(
+        cls,
         ctx,
         student_input,
         student_weight,
@@ -133,7 +134,6 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
         target,
         student_bias=None,
         teacher_bias=None,
-        loss_fn=None,
         chunk_size=1024,
         ignore_index=-100,
         weight_hard_loss=0.5,
@@ -175,7 +175,7 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
 
         loss_func_to_call = partial(
             LigerFusedLinearDistillationBase._compute_loss,
-            distillation_loss_fn=loss_fn,
+            distillation_loss_fn=cls.distillation_loss_fn,
             full_target=target,
             ignore_index=ignore_index,
             weight_hard_loss=weight_hard_loss,
@@ -263,4 +263,4 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
             grad_weight = grad_weight * grad_output
             grad_bias = grad_bias * grad_output if grad_bias is not None else None
 
-        return grad_input, grad_weight, None, grad_bias
+        return grad_input, grad_weight, None, None, None, grad_bias

liger_kernel/chunked_loss/fused_linear_preference.py

@@ -16,12 +16,12 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
 
     @staticmethod
     def forward(
+        cls,
         ctx,
         _input,
         weight,
         target,
         bias=None,
-        loss_fn=None,
         chunk_size=1,
         ignore_index=-100,
         alpha=1.0,
@@ -89,7 +89,7 @@ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
 
         compute_loss = partial(
             LigerFusedLinearPreferenceBase._compute_loss,
-            preference_loss_fn=loss_fn,
+            preference_loss_fn=cls.preference_loss_fn,
             ignore_index=ignore_index,
             alpha=alpha,
             beta=beta,

liger_kernel/chunked_loss/fused_linear_rlhf.py

@@ -0,0 +1,240 @@
+from abc import abstractmethod
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+
+class LigerFusedLinearRLHFBase(torch.autograd.Function):
+    @abstractmethod
+    def rlhf_loss_fn(*args, **kwargs):
+        """
+        To be extended by subclasses.
+        """
+        raise NotImplementedError("RLHF loss function must be implemented.")
+
+    @staticmethod
+    def forward(
+        cls,
+        ctx,
+        _input,
+        weight,
+        attention_mask,
+        rewards,
+        bias=None,
+        num_generations=4,
+        beta=0.1,
+        compiled=True,
+        use_ref_model=False,
+        ref_input=None,
+        ref_weight=None,
+        ref_bias=None,
+        chunk_size=1,
+    ):
+        """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
+
+        # Initialize accumulators
+        loss_acc = torch.zeros((), device=_input.device)
+        grad_weight = torch.zeros_like(weight)  # [V, H]
+        grad_inputs = []
+        grad_bias = torch.zeros_like(bias) if bias is not None else None  # [V]
+        aggregated_metrics = []
+
+        # Create a partial function with fixed arguments
+        compute_loss = partial(
+            LigerFusedLinearRLHFBase._compute_chunk_loss,
+            beta=beta,
+            use_ref_model=use_ref_model,
+            ref_weight=ref_weight,
+            ref_bias=ref_bias,
+            rlhf_loss_fn=cls.rlhf_loss_fn,
+        )
+
+        def fused_fwd_bwd(input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk):
+            """Fused forward and backward for a chunk."""
+            if bias is not None:
+                return torch.func.grad_and_value(compute_loss, argnums=(0, 1, 5), has_aux=True)(
+                    input_chunk,  # arg 0
+                    weight,  # arg 1
+                    attention_mask_chunk,  # arg 2
+                    rewards_chunk,  # arg 3
+                    ref_input_chunk,  # arg 4
+                    bias,  # arg 5
+                )
+            else:
+                return torch.func.grad_and_value(compute_loss, argnums=(0, 1), has_aux=True)(
+                    input_chunk,  # arg 0
+                    weight,  # arg 1
+                    attention_mask_chunk,  # arg 2
+                    rewards_chunk,  # arg 3
+                    ref_input_chunk,  # arg 4
+                )
+
+        def accumulate_chunk(input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk=None):
+            if bias is not None:
+                (chunk_grad_input, chunk_grad_weight, chunk_grad_bias), (chunk_loss, chunk_metrics) = fused_fwd_bwd(
+                    input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk
+                )
+                grad_bias.add_(chunk_grad_bias)
+            else:
+                (chunk_grad_input, chunk_grad_weight), (chunk_loss, chunk_metrics) = fused_fwd_bwd(
+                    input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk
+                )
+
+            # Accumulate gradients and loss
+            grad_weight.add_(chunk_grad_weight)
+            grad_inputs.append(chunk_grad_input)
+            loss_acc.add_(chunk_loss)
+
+            # Initialize storage for metrics on first chunk
+            if len(aggregated_metrics) == 0:
+                for metric in chunk_metrics:
+                    if metric.ndim == 0:
+                        aggregated_metrics.append(torch.zeros((), device=metric.device))
+                    else:
+                        aggregated_metrics.append([])
+
+            # Accumulate metrics
+            for i, metric in enumerate(chunk_metrics):
+                if metric.ndim == 0:
+                    aggregated_metrics[i].add_(metric)
+                else:
+                    aggregated_metrics[i].append(metric)
+
+        if compiled:
+            accumulate_chunk = torch.compile(accumulate_chunk)
+
+        # 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)
+        _rewards_chunks = torch.chunk(rewards, chunks=chunks, dim=0)
+        _ref_input_chunks = torch.chunk(ref_input, chunks=chunks, dim=0) if use_ref_model else [None] * chunks
+
+        for input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk in zip(
+            _input_chunks, _attention_mask_chunks, _rewards_chunks, _ref_input_chunks
+        ):
+            # Mark dynamic dimensions
+            torch._dynamo.mark_dynamic(input_chunk, 1)
+            torch._dynamo.mark_dynamic(attention_mask_chunk, 1)
+            if ref_input_chunk is not None:
+                torch._dynamo.mark_dynamic(ref_input_chunk, 1)
+
+            accumulate_chunk(input_chunk, attention_mask_chunk, rewards_chunk, ref_input_chunk)
+
+        # Scale accumulated loss by number of chunks since we're averaging
+        loss_acc = loss_acc / chunks
+
+        # Combine gradients
+        grad_input = torch.cat(grad_inputs, dim=0)
+
+        # Save for backward
+        ctx.save_for_backward(grad_input, grad_weight, grad_bias)
+
+        # Finalize metrics
+        final_metrics = []
+        for metric in aggregated_metrics:
+            if isinstance(metric, list):
+                final_metrics.append(torch.cat(metric, dim=0))
+            else:
+                final_metrics.append(metric / chunks)
+
+        return loss_acc, tuple(final_metrics)
+
+    @staticmethod
+    def _compute_chunk_loss(
+        input_chunk,
+        weight,
+        attention_mask_chunk,
+        rewards_chunk,
+        ref_input_chunk=None,
+        bias=None,
+        beta=0.1,
+        use_ref_model=False,
+        ref_weight=None,
+        ref_bias=None,
+        rlhf_loss_fn=None,
+    ):
+        """Compute loss for a single chunk."""
+        # Get policy log probabilities using chunk_forward
+        log_probs, _, logits_mean = LigerFusedLinearRLHFBase.chunk_forward(input_chunk, weight, bias=bias)
+
+        # Get reference log probabilities if needed
+        ref_log_probs = None
+        if use_ref_model and ref_input_chunk is not None:
+            with torch.no_grad():
+                ref_log_probs, _, _ = LigerFusedLinearRLHFBase.chunk_forward(ref_input_chunk, ref_weight, bias=ref_bias)
+
+        # Compute chunk loss and metrics using the provided loss function
+        chunk_loss, chunk_metrics = rlhf_loss_fn(
+            log_probs=log_probs,
+            attention_mask=attention_mask_chunk,
+            rewards=rewards_chunk,
+            ref_log_probs=ref_log_probs,
+            beta=beta,
+        )
+
+        return chunk_loss, (logits_mean, *chunk_metrics)
+
+    @staticmethod
+    def chunk_forward(input_chunk, weight, bias=None):
+        """Forward pass computation for a single chunk without explicit reshaping."""
+        # Directly compute logits via batched matrix multiplication: [B, T, H] @ [H, V] -> [B, T, V]
+        logits = torch.matmul(input_chunk, weight.t())
+        if bias is not None:
+            logits = logits + bias  # Broadcasts bias to [B, T, V]
+
+        # Compute log probabilities using softmax over the last dimension
+        log_probs = F.log_softmax(logits.float(), dim=-1)
+
+        # Monitoring: compute mean of logits
+        batch_size, seq_len, _ = input_chunk.shape
+        logits_mean = logits.sum() / (batch_size * seq_len * weight.shape[0])
+        return log_probs, logits, logits_mean
+
+    @staticmethod
+    def backward(ctx, grad_output, *grad_metrics):
+        """Backward pass for RLHF loss."""
+        grad_input, grad_weight, grad_bias = ctx.saved_tensors
+        if grad_output != 1.0:
+            grad_input = grad_input * grad_output
+            grad_weight = grad_weight * grad_output
+            if grad_bias is not None:
+                grad_bias = grad_bias * grad_output
+
+        return (
+            grad_input,
+            grad_weight,
+            None,  # grad_attention_mask
+            None,  # grad_rewards
+            grad_bias,
+            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
+        )