liger-kernel-nightly 0.5.5.dev20250331170510__py3-none-any.whl → 0.5.5.dev20250402185606__py3-none-any.whl

liger_kernel/chunked_loss/functional.py

@@ -1,5 +1,6 @@
 from liger_kernel.chunked_loss.cpo_loss import LigerFusedLinearCPOFunction
 from liger_kernel.chunked_loss.dpo_loss import LigerFusedLinearDPOFunction
+from liger_kernel.chunked_loss.grpo_loss import LigerFusedLinearGRPOFunction
 from liger_kernel.chunked_loss.jsd_loss import LigerFusedLinearJSDFunction
 from liger_kernel.chunked_loss.kto_loss import LigerFusedLinearKTOFunction
 from liger_kernel.chunked_loss.orpo_loss import LigerFusedLinearORPOFunction
@@ -11,3 +12,4 @@ liger_fused_linear_jsd = LigerFusedLinearJSDFunction.apply
 liger_fused_linear_cpo = LigerFusedLinearCPOFunction.apply
 liger_fused_linear_simpo = LigerFusedLinearSimPOFunction.apply
 liger_fused_linear_kto = LigerFusedLinearKTOFunction.apply
+liger_fused_linear_grpo = LigerFusedLinearGRPOFunction.apply

liger_kernel/chunked_loss/fused_linear_ppo.py

@@ -0,0 +1,330 @@
+from abc import abstractmethod
+from functools import partial
+
+import torch
+import torch._dynamo.config
+import torch.nn.functional as F
+
+
+class LigerFusedLinearPPOBase(torch.autograd.Function):
+    @abstractmethod
+    def ppo_loss_fn(*args, **kwargs):
+        """
+        To be extended by subclasses.
+        """
+        raise NotImplementedError("PPO loss function must be implemented.")
+
+    @staticmethod
+    def forward(
+        cls,
+        ctx,
+        _input,
+        weight,
+        selected_token_ids,
+        attention_mask,
+        advantages,
+        bias=None,
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
+        ref_input=None,
+        ref_weight=None,
+        ref_bias=None,
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        beta=0.04,
+        temperature=1.0,
+        compiled=True,
+        use_ref_model=False,
+        chunk_size=1,
+    ):
+        """Chunked forward pass for PPO loss computation.
+
+        Args:
+            cls: The class
+            ctx: Context for backward
+            _input: Input tensor
+            weight: Weight tensor
+            selected_token_ids: Selected token ids tensor
+            attention_mask: Attention mask tensor
+            advantages: Advantages tensor
+            bias: Bias tensor
+            ref_per_token_logps: Reference model log probs per token tensor
+            old_per_token_logps: Old per token log probabilities tensor
+            ref_input: Reference model input tensor
+            ref_weight: Reference model weight tensor
+            ref_bias: Reference model bias tensor
+            epsilon_low: Lower bound for clipping the importance sampling ratio
+            epsilon_high: Upper bound for clipping the importance sampling ratio
+            beta: Weight for the KL penalty
+            temperature: Temperature for the logits
+            compiled: Whether to use torch compile
+            use_ref_model: Whether to use a reference model
+            chunk_size: Size of chunks for processing in other loss modules
+        """
+        if use_ref_model:
+            assert ref_per_token_logps is not None or ref_input is not None, (
+                "If use_ref_model is True, ref_per_token_logps or ref_input must be provided"
+            )
+            if ref_per_token_logps is not None and ref_input is not None:
+                raise Warning("Both ref_per_token_logps and ref_input are provided. Using ref_per_token_logps.")
+        # Initialize accumulators
+        loss_acc = torch.zeros((), device=_input.device, dtype=torch.float32)
+        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(
+            LigerFusedLinearPPOBase._compute_chunk_loss,
+            ref_weight=ref_weight,
+            ref_bias=ref_bias,
+            full_attention_mask=attention_mask,
+            epsilon_low=epsilon_low,
+            epsilon_high=epsilon_high,
+            beta=beta,
+            temperature=temperature,
+            use_ref_model=use_ref_model,
+            ppo_loss_fn=cls.ppo_loss_fn,
+        )
+
+        def fused_fwd_bwd(
+            input_chunk,
+            selected_token_ids_chunk,
+            attention_mask_chunk,
+            advantages_chunk,
+            ref_per_token_logps_chunk,
+            old_per_token_logps_chunk,
+            ref_input_chunk,
+        ):
+            """Fused forward and backward for a chunk."""
+            argnums = (0, 1, 5) if bias is not None else (0, 1)
+            return torch.func.grad_and_value(compute_loss, argnums=argnums, has_aux=True)(
+                input_chunk,  # arg 0
+                weight,  # arg 1
+                selected_token_ids_chunk,  # arg 2
+                attention_mask_chunk,  # arg 3
+                advantages_chunk,  # arg 4
+                bias,  # arg 5
+                ref_per_token_logps_chunk=ref_per_token_logps_chunk,  # arg 6
+                old_per_token_logps_chunk=old_per_token_logps_chunk,  # arg 7
+                ref_input_chunk=ref_input_chunk,  # arg 8
+            )
+
+        def accumulate_chunk(
+            input_chunk,
+            selected_token_ids_chunk,
+            attention_mask_chunk,
+            advantages_chunk,
+            ref_per_token_logps_chunk=None,
+            old_per_token_logps_chunk=None,
+            ref_input_chunk=None,
+        ):
+            (chunk_grad_input, chunk_grad_weight, *chunk_grad_bias), (chunk_loss, chunk_metrics) = fused_fwd_bwd(
+                input_chunk,
+                selected_token_ids_chunk,
+                attention_mask_chunk,
+                advantages_chunk,
+                ref_per_token_logps_chunk,
+                old_per_token_logps_chunk,
+                ref_input_chunk,
+            )
+            if bias is not None:
+                grad_bias.add_(chunk_grad_bias[0])
+
+            # 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:
+            # TODO: Figure out what is better to compile here
+            # accumulate_chunk = torch.compile(accumulate_chunk)
+            fused_fwd_bwd = torch.compile(fused_fwd_bwd)
+
+        # Process input in chunks based on chunk_size
+        chunks = max(1, _input.shape[0] // chunk_size)
+        _input_chunks = torch.chunk(_input, chunks=chunks, dim=0)
+        _selected_token_ids_chunks = torch.chunk(selected_token_ids, chunks=chunks, dim=0)
+        _attention_mask_chunks = torch.chunk(attention_mask, chunks=chunks, dim=0)
+        _advantages_chunks = torch.chunk(advantages, chunks=chunks, dim=0)
+        _ref_per_token_logps_chunks = (
+            torch.chunk(ref_per_token_logps, chunks=chunks, dim=0)
+            if use_ref_model and ref_per_token_logps is not None
+            else [None] * chunks
+        )
+        _old_per_token_logps_chunks = (
+            torch.chunk(old_per_token_logps, chunks=chunks, dim=0)
+            if old_per_token_logps is not None
+            else [None] * chunks
+        )
+        # if ref_log_probs is not none, then we don't need ref_input to calculate the log probs
+        _ref_input_chunks = (
+            torch.chunk(ref_input, chunks=chunks, dim=0)
+            if use_ref_model and ref_per_token_logps is None
+            else [None] * chunks
+        )
+
+        for (
+            input_chunk,
+            selected_token_ids_chunk,
+            attention_mask_chunk,
+            advantages_chunk,
+            ref_per_token_logps_chunk,
+            old_per_token_logps_chunk,
+            ref_input_chunk,
+        ) in zip(
+            _input_chunks,
+            _selected_token_ids_chunks,
+            _attention_mask_chunks,
+            _advantages_chunks,
+            _ref_per_token_logps_chunks,
+            _old_per_token_logps_chunks,
+            _ref_input_chunks,
+        ):
+            # Mark dynamic dimensions
+            torch._dynamo.mark_dynamic(input_chunk, 1)
+            torch._dynamo.mark_dynamic(selected_token_ids_chunk, 1)
+            torch._dynamo.mark_dynamic(attention_mask_chunk, 1)
+            if ref_per_token_logps_chunk is not None:
+                torch._dynamo.mark_dynamic(ref_per_token_logps_chunk, 1)
+            if ref_input_chunk is not None:
+                torch._dynamo.mark_dynamic(ref_input_chunk, 1)
+            if old_per_token_logps_chunk is not None:
+                torch._dynamo.mark_dynamic(old_per_token_logps_chunk, 1)
+
+            accumulate_chunk(
+                input_chunk,
+                selected_token_ids_chunk,
+                attention_mask_chunk,
+                advantages_chunk,
+                ref_per_token_logps_chunk,
+                old_per_token_logps_chunk,
+                ref_input_chunk,
+            )
+
+        # 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)
+
+        return loss_acc, tuple(final_metrics)
+
+    @staticmethod
+    def _compute_chunk_loss(
+        input_chunk,
+        weight,
+        selected_token_ids_chunk,
+        attention_mask_chunk,
+        advantages_chunk,
+        bias=None,
+        ref_per_token_logps_chunk=None,
+        old_per_token_logps_chunk=None,
+        ref_input_chunk=None,
+        ref_weight=None,
+        ref_bias=None,
+        full_attention_mask=None,
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        beta=0.04,
+        temperature=1.0,
+        use_ref_model=False,
+        ppo_loss_fn=None,
+    ):
+        """Compute loss for a single chunk."""
+        # Get policy log probabilities using chunk_forward
+        log_probs, _ = LigerFusedLinearPPOBase.chunk_forward(input_chunk, weight, bias=bias, temperature=temperature)
+
+        # Get reference log probabilities if needed
+        ref_log_probs = None
+        if use_ref_model and ref_per_token_logps_chunk is None:
+            with torch.no_grad():
+                ref_log_probs, _ = LigerFusedLinearPPOBase.chunk_forward(
+                    ref_input_chunk, ref_weight, bias=ref_bias, temperature=temperature
+                )
+
+        # Compute chunk loss and metrics using the provided loss function
+        chunk_loss, chunk_metrics = ppo_loss_fn(
+            log_probs=log_probs,
+            selected_token_ids=selected_token_ids_chunk,
+            attention_mask=attention_mask_chunk,
+            advantages=advantages_chunk,
+            full_attention_mask=full_attention_mask,
+            ref_per_token_logps=ref_per_token_logps_chunk.float() if ref_per_token_logps_chunk is not None else None,
+            old_per_token_logps=old_per_token_logps_chunk.float() if old_per_token_logps_chunk is not None else None,
+            ref_log_probs=ref_log_probs,  # used when ref_per_token_logps is None
+            epsilon_low=epsilon_low,
+            epsilon_high=epsilon_high,
+            beta=beta,
+        )
+
+        return chunk_loss, chunk_metrics
+
+    @staticmethod
+    def chunk_forward(input_chunk, weight, bias=None, temperature=1.0):
+        """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]
+        if temperature != 1.0:
+            logits = logits / temperature
+
+        # Compute log probabilities using softmax over the last dimension
+        log_probs = F.log_softmax(logits.float(), dim=-1)
+
+        return log_probs, logits
+
+    @staticmethod
+    def backward(ctx, grad_output, *grad_metrics):
+        """Backward pass for PPO 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_selected_token_ids
+            None,  # grad_attention_mask
+            None,  # grad_advantages
+            grad_bias,
+            None,  # grad_ref_per_token_logps
+            None,  # grad_old_per_token_logps
+            None,  # grad_ref_input
+            None,  # grad_ref_weight
+            None,  # grad_ref_bias
+            None,  # grad_epsilon_low
+            None,  # grad_epsilon_high
+            None,  # grad_beta
+            None,  # grad_temperature
+            None,  # grad_compiled
+            None,  # grad_use_ref_model
+            None,  # grad_chunk_size
+        )

liger_kernel/chunked_loss/grpo_loss.py

@@ -1,66 +1,76 @@
 import torch
 
-from liger_kernel.chunked_loss.fused_linear_rlhf import LigerFusedLinearRLHFBase
+from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
 
 
-class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
+def k3_loss_fn(log_p, log_q):
+    # computes k3 estimate of KL[q, p]
+    # ref: http://joschu.net/blog/kl-approx.html
+    return torch.exp(log_p - log_q) - (log_p - log_q) - 1.0
+
+
+def clip_coef_fn(coef, epsilon_low, epsilon_high):
+    return torch.clamp(coef, 1 - epsilon_low, 1 + epsilon_high)
+
+
+class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
     @staticmethod
-    def rlhf_loss_fn(
+    def ppo_loss_fn(
         log_probs,
+        selected_token_ids,
         attention_mask,
-        rewards,
-        ref_log_probs=None,
-        beta=0.1,
+        advantages,
+        full_attention_mask,
+        ref_per_token_logps=None,  # shape: [chunk_size, seq_len]
+        old_per_token_logps=None,
+        ref_log_probs=None,  # used when ref_per_token_logps is None (shape: [chunk_size, seq_len, vocab_size])
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        beta=0.04,
         **kwargs,
     ):
         """GRPO Loss Function matching GRPOTrainer implementation."""
-        # Get chosen token probabilities
-        chosen_tokens = log_probs.argmax(dim=-1)  # (batch_size, seq_len)
-        chosen_token_logprobs = log_probs.gather(dim=-1, index=chosen_tokens.unsqueeze(-1)).squeeze(
+        per_token_logps = log_probs.gather(dim=-1, index=selected_token_ids.unsqueeze(-1)).squeeze(
             -1
         )  # (batch_size, seq_len)
 
         # Get reference model probabilities
-        if ref_log_probs is not None:
-            with torch.no_grad():
-                ref_token_logprobs = ref_log_probs.gather(dim=-1, index=chosen_tokens.unsqueeze(-1)).squeeze(-1)
-        else:
-            ref_token_logprobs = chosen_token_logprobs.detach()
-
-        # Compute advantages per batch entry in a grouped fashion
-        mean_grouped_rewards = rewards.mean()  # [batch_size,]
-        std_grouped_rewards = rewards.std()  # [batch_size,]
-
-        # Calculate advantages using the same epsilon as in GRPOTrainer
-        eps = 1e-4
-        advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + eps)
+        if ref_per_token_logps is None:
+            if ref_log_probs is not None:
+                with torch.no_grad():
+                    ref_per_token_logps = ref_log_probs.gather(dim=-1, index=selected_token_ids.unsqueeze(-1)).squeeze(
+                        -1
+                    )
+            else:
+                ref_per_token_logps = per_token_logps.detach()
 
         # Compute policy gradient loss with importance sampling ratio
-        ratio = torch.exp(chosen_token_logprobs - chosen_token_logprobs.detach())
-        policy_loss = -ratio * advantages.unsqueeze(1)
-
-        # Compute KL penalty
-        kl_div = (
-            torch.exp(ref_token_logprobs - chosen_token_logprobs) - (ref_token_logprobs - chosen_token_logprobs) - 1.0
-        )
-
-        # Combine losses
-        per_token_loss = policy_loss + beta * kl_div
-
-        # Apply masking and normalize
-        masked_loss = per_token_loss * attention_mask
-        seq_lengths = attention_mask.sum()
-        seq_lengths = torch.clamp(seq_lengths, min=1.0)
-        loss = masked_loss.sum() / seq_lengths
+        old_per_token_logps = old_per_token_logps if old_per_token_logps is not None else per_token_logps.detach()
+        coef_1 = torch.exp(per_token_logps - old_per_token_logps)
+        coef_2 = clip_coef_fn(coef_1, epsilon_low, epsilon_high)
+        per_token_loss1 = coef_1 * advantages.unsqueeze(1)
+        per_token_loss2 = coef_2 * advantages.unsqueeze(1)
+        per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+        if beta != 0.0:
+            # Compute KL penalty (approximates KL[per_token_logps, ref_per_token_logps])
+            kl_div = k3_loss_fn(ref_per_token_logps, per_token_logps)
+            # Combine losses
+            per_token_loss = per_token_loss + beta * kl_div
+
+        # Note: We normalize by the number of tokens in the batch (using full_attention_mask),
+        # which is consistent with the DAPO loss implementation (https://arxiv.org/html/2503.14476v1)
+        # and TRL GRPO implementation
+        # (https://github.com/huggingface/trl/blob/e751a16df56e70190fb94bed4a2035eec3303777/trl/trainer/grpo_trainer.py#L966)
+        loss = (per_token_loss * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0)
 
         # Calculate metrics
-        metrics = (
-            chosen_token_logprobs.mean(),  # mean log prob
-            chosen_token_logprobs.std(),  # std log prob
-            log_probs.mean(),  # mean all log probs
-            ((kl_div * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)).mean(),  # mean KL div
+        metrics = []
+        if beta != 0.0:
+            metrics.append(((kl_div * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0)))
+        is_clipped = ((coef_1 < 1 - epsilon_low) & (advantages.unsqueeze(1) < 0)) | (
+            (coef_1 > 1 + epsilon_high) & (advantages.unsqueeze(1) > 0)
         )
-
+        metrics.append((is_clipped * attention_mask).sum() / torch.clamp(full_attention_mask.sum(), min=1.0))
         return loss, metrics
 
     @classmethod
@@ -69,16 +79,21 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
         ctx,
         _input,
         weight,
+        selected_token_ids,
         attention_mask,
-        rewards,
+        advantages,
         bias=None,
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
         ref_input=None,
         ref_weight=None,
         ref_bias=None,
-        beta=0.1,
+        beta=0.04,
+        epsilon_low=0.2,
+        epsilon_high=0.2,
+        temperature=1.0,
         compiled=True,
         use_ref_model=True,
-        num_generations=1,
         chunk_size=1,
     ):
         """
@@ -86,16 +101,18 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
         Args:
             _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
             weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            selected_token_ids (torch.Tensor): Selected token ids tensor. Shape: (batch_size, seq_len)
             attention_mask (torch.Tensor): Attention mask tensor. Shape: (batch_size, seq_len)
-            rewards (torch.Tensor): Rewards tensor. Shape: (batch_size,)
+            advantages (torch.Tensor): Advantages tensor. Shape: (batch_size,)
             bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ref_per_token_logps:  Reference model log probs per token tensor. Shape:(batch_size, seq_len)
             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
+            temperature (float): Temperature for the logits
             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
@@ -105,16 +122,21 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             ctx=ctx,
             _input=_input,
             weight=weight,
+            selected_token_ids=selected_token_ids,
             attention_mask=attention_mask,
-            rewards=rewards,
+            advantages=advantages,
             bias=bias,
+            ref_per_token_logps=ref_per_token_logps,
+            old_per_token_logps=old_per_token_logps,
             ref_input=ref_input,
             ref_weight=ref_weight,
             ref_bias=ref_bias,
             beta=beta,
+            epsilon_low=epsilon_low,
+            epsilon_high=epsilon_high,
+            temperature=temperature,
             compiled=compiled,
             use_ref_model=use_ref_model,
-            num_generations=num_generations,
             chunk_size=chunk_size,
         )
 
@@ -126,16 +148,22 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearRLHFBase):
             grad_output: Gradient of the loss (scalar)
             grad_metrics: Gradients of the metrics (not used in backward computation)
         """
-        grads = LigerFusedLinearRLHFBase.backward(ctx, grad_output)
+        grads = LigerFusedLinearPPOBase.backward(ctx, grad_output)
         return (
-            *grads[:5],  # grad_input, grad_weight, grad_attention_mask, grad_rewards, grad_bias
+            *grads[
+                :6
+            ],  # grad_input, grad_weight, grad_selected_token_ids, grad_attention_mask, grad_advantages, grad_bias
+            None,  # grad_ref_per_token_logps
+            None,  # grad_old_per_token_logps
             None,  # grad_ref_input
             None,  # grad_ref_weight
             None,  # grad_ref_bias
             None,  # grad_beta
+            None,  # grad_epsilon_low
+            None,  # grad_epsilon_high
+            None,  # grad_temperature
             None,  # grad_compiled
             None,  # grad_use_ref_model
-            None,  # grad_num_generations
             None,  # grad_chunk_size
         )
 
@@ -145,34 +173,43 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
 
     def __init__(
         self,
-        beta: float = 0.1,
+        beta: float = 0.04,
         compiled: bool = True,
         use_ref_model: bool = True,
-        num_generations: int = 1,
         chunk_size: int = 1,
+        epsilon_low: float = 0.2,
+        epsilon_high: float = 0.2,
+        temperature: float = 1.0,
     ):
         """
         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.
+            epsilon_low (float): Lower bound for the importance sampling ratio.
+            epsilon_high (float): Upper bound for the importance sampling ratio.
+            temperature (float): Temperature for the logits.
         """
         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
+        self.epsilon_low = epsilon_low
+        self.epsilon_high = epsilon_high
+        self.temperature = temperature
 
     def forward(
         self,
         _input,
         lin_weight,
+        selected_token_ids,
         attention_mask,
-        rewards,
+        advantages,
         bias=None,
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
         ref_input=None,
         ref_weight=None,
         ref_bias=None,
@@ -180,15 +217,20 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
         return LigerFusedLinearGRPOFunction.apply(
             _input,
             lin_weight,
+            selected_token_ids,
             attention_mask,
-            rewards,
+            advantages,
             bias,
+            ref_per_token_logps,
+            old_per_token_logps,
             ref_input,
             ref_weight,
             ref_bias,
             self.beta,
+            self.epsilon_low,
+            self.epsilon_high,
+            self.temperature,
             self.compiled,
             self.use_ref_model,
-            self.num_generations,
             self.chunk_size,
         )

liger_kernel/ops/cross_entropy.py

@@ -9,6 +9,7 @@ import triton.language as tl
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import element_mul_kernel
 from liger_kernel.ops.utils import is_hip
+from liger_kernel.utils import infer_device
 
 if compare_version("triton", operator.ge, "3.0.0"):
     try:
@@ -59,7 +60,7 @@ def liger_cross_entropy_kernel(
     z_loss_ptr: Pointer to tensor to store the z loss. No operation if RETURN_Z_LOSS is 0.
     loss_stride (int): The stride of the loss tensor.
     n_cols (int): The number of columns in the input tensor.
-    n_non_ignore (flaot): The number of non-ignored elements in the batch.
+    n_non_ignore (float): The number of non-ignored elements in the batch.
     sum_non_ignore_weight (float): The sum of non-ignored target's weights in the batch.
     weight_sum (float): The sum of weight tensor.
     ignore_index (int): The index to ignore in the target.
@@ -258,7 +259,7 @@ def liger_cross_entropy_kernel(
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 65536 // 2  # the best size we found by manually tuning
+MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2  # the best size we found by manually tuning
 
 
 def cross_entropy_forward(

{liger_kernel_nightly-0.5.5.dev20250331170510.dist-info → liger_kernel_nightly-0.5.5.dev20250402185606.dist-info}/METADATA

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

{liger_kernel_nightly-0.5.5.dev20250331170510.dist-info → liger_kernel_nightly-0.5.5.dev20250402185606.dist-info}/RECORD

@@ -5,18 +5,18 @@ liger_kernel/chunked_loss/README.md,sha256=0FmkFC3hKBqyoDT5uTlIYmrvRkF-EOCR1y-EB
 liger_kernel/chunked_loss/__init__.py,sha256=ATu-xX5Fc49Cr6yBOGBRNTo593ZrU5ZCsIuvoIbJWw4,603
 liger_kernel/chunked_loss/cpo_loss.py,sha256=Gzz1eU4kgcbdubFVRy55e8A1Cr-r45UgNicXwZIjmBU,5454
 liger_kernel/chunked_loss/dpo_loss.py,sha256=xZwGqS04si9zXyob95SAdalC-hajZg8fWINqiqffN8k,5855
-liger_kernel/chunked_loss/functional.py,sha256=THWWpCnRVhTVfnPnyvQjdBvo1JDtxhwLmtZE_yiBBqM,817
+liger_kernel/chunked_loss/functional.py,sha256=9G3nKm-Bi7uoZRFkL8wwGMl6juDl4bSzDvTa5GHZPzg,955
 liger_kernel/chunked_loss/fused_linear_distillation.py,sha256=ooR-qnZCyWJN935oHCSWLaKKKyaYERyhNczRGi1VOiw,11935
+liger_kernel/chunked_loss/fused_linear_ppo.py,sha256=2_UvvIksUP45RBw3c-88-jOtjGATf04vaWopcqtX4Oo,12688
 liger_kernel/chunked_loss/fused_linear_preference.py,sha256=ojB42jYPu0c4ki96Ft-hy7Sf6fh_WikG-aWNrlZzSio,18362
-liger_kernel/chunked_loss/fused_linear_rlhf.py,sha256=wGujqwLz91mOE9MmdenhBIKvbmswhwtINMCpcP7D74c,9050
 liger_kernel/chunked_loss/fused_linear_unpaired_preference.py,sha256=RiuK3UtRwH9T6jZ36sA8Urj-TVuOLOO2syLg_JOQapY,13437
-liger_kernel/chunked_loss/grpo_loss.py,sha256=axED3628yKODu1v7PMAvSd08WZqwNQvJOTUYMgcihdQ,6665
+liger_kernel/chunked_loss/grpo_loss.py,sha256=6Mb4ZT6MfnOr4Xo681rMR0LKkhzJhInvQp8wp2YVMK0,8913
 liger_kernel/chunked_loss/jsd_loss.py,sha256=u2ahkuHsbhpNaKcpBCz5gCMDk9ou-P04DHji592dIBo,7067
 liger_kernel/chunked_loss/kto_loss.py,sha256=llVCe6DkcpCo57seGWoMikaQVFApx764jsmSbQyqwQY,7529
 liger_kernel/chunked_loss/orpo_loss.py,sha256=nu9UYG16dcMw93lvHi4_hYs3Q0FK1KnlmMRj7OpYU8s,4872
 liger_kernel/chunked_loss/simpo_loss.py,sha256=fy2w8KbhMrBv7b1jdIeH3bBFxY52bPQPZb3KwBvmurM,5385
 liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-liger_kernel/ops/cross_entropy.py,sha256=yKKhN63I7r9NxJye4wTLBvvKAyrXQt6jf4nBo3lJyVg,18860
+liger_kernel/ops/cross_entropy.py,sha256=T5oSsqOS1y-Iea5o9v_BSU-_mIEXqWAT1oX_m59NcA4,18941
 liger_kernel/ops/dyt.py,sha256=YD1-buHz9VmIX838VKzLc-lm5CeUQ4LAskGDWBUMQHA,6187
 liger_kernel/ops/fused_linear_cross_entropy.py,sha256=1Y3Uk_TCSjqKgoG2eot1ptnWXJXXQESqGvOmqAW1gsM,10912
 liger_kernel/ops/fused_linear_jsd.py,sha256=Seshez2qaM6HiTQ8_HEqSwhaeVruNT1SvIM4ZrAPBEU,9602
@@ -72,9 +72,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=pdekW7l6Qg_aqa5SYKYlSWUF8m3lkOFvFLcIMEHrz9s,8338
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/METADATA,sha256=KEjXLNI8PYfmvipid4KUVeM0XE5oKXd5Pl7ikrZbAqU,22959
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.5.dev20250331170510.dist-info/RECORD,,
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/METADATA,sha256=XQaGc9bnsEFdwtLh1Mv5_fX-TIejLbcHk1SP-FEY5ew,22959
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.5.dev20250402185606.dist-info/RECORD,,

liger_kernel/chunked_loss/fused_linear_rlhf.py

@@ -1,240 +0,0 @@
-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
-        )