feat(grpo): optional entropy regularization for chunked GRPO loss

src/liger_kernel/chunked_loss/fused_linear_ppo.py

@@ -44,6 +44,7 @@ def forward(
         vllm_is_ratio=None,
         delta=None,
         use_bias_correction_kl=False,
+        entropy_coef=0.0,
     ):
         # TODO: check torch compile matmul
         """Chunked forward pass for PPO loss computation.
@@ -125,6 +126,7 @@ def forward(
             sapo_temperature_neg=sapo_temperature_neg,
             delta=delta,
             use_bias_correction_kl=use_bias_correction_kl,
+            entropy_coef=entropy_coef,
         )
 
         def fused_fwd_bwd(
@@ -327,6 +329,7 @@ def _compute_chunk_loss(
         sapo_temperature_neg=1.05,
         delta=None,
         use_bias_correction_kl=False,
+        entropy_coef=0.0,
     ):
         """Compute loss for a single chunk."""
         # Get policy log probabilities using chunk_forward
@@ -361,6 +364,7 @@ def _compute_chunk_loss(
             vllm_is_ratio=vllm_is_ratio_chunk,
             delta=delta,
             use_bias_correction_kl=use_bias_correction_kl,
+            entropy_coef=entropy_coef,
         )
 
         return chunk_loss, chunk_metrics

src/liger_kernel/chunked_loss/grpo_loss.py

@@ -78,6 +78,7 @@ def ppo_loss_fn(
         vllm_is_ratio=None,  # vLLM importance sampling ratio (chunk_size, seq_len) or (chunk_size, 1) or None
         delta=None,  # Upper clamp for two-sided clipping (INTELLECT-2)
         use_bias_correction_kl=False,  # Importance-sampling-corrected KL (DeepSeek-V3.2)
+        entropy_coef=0.0,  # Entropy regularization coefficient: loss -= entropy_coef * H
         **kwargs,
     ):
         """GRPO Loss Function matching GRPOTrainer implementation."""
@@ -165,6 +166,15 @@ def ppo_loss_fn(
             # Combine losses
             per_token_loss = per_token_loss + beta * kl_div
 
+        # Entropy regularization (a.k.a. PPO entropy bonus): higher per-token entropy H lowers the loss,
+        # counteracting entropy collapse. H = -sum_v p_v * log p_v, computed from the
+        # chunk's full log_probs (already materialized for the policy term -> no extra
+        # logit pass). per_token_loss is (B, T) at token level; padding is masked out in
+        # the reduction below. NOTE: assumes importance_sampling_level == "token".
+        if entropy_coef != 0.0:
+            per_token_entropy = -(log_probs.exp() * log_probs).sum(dim=-1)  # (B, T)
+            per_token_loss = per_token_loss - entropy_coef * per_token_entropy
+
         # 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
@@ -251,6 +261,7 @@ def forward(
         vllm_is_ratio=None,
         delta=None,
         use_bias_correction_kl=False,
+        entropy_coef=0.0,
     ):
         """
         Fused linear layer with GRPO loss.
@@ -317,6 +328,7 @@ def forward(
             vllm_is_ratio=vllm_is_ratio,
             delta=delta,
             use_bias_correction_kl=use_bias_correction_kl,
+            entropy_coef=entropy_coef,
         )
 
     @staticmethod
@@ -352,6 +364,7 @@ def backward(ctx, grad_output, *grad_metrics):
             None,  # grad_vllm_is_ratio
             None,  # grad_delta
             None,  # grad_use_bias_correction_kl
+            None,  # grad_entropy_coef
         )
 
 
@@ -374,6 +387,7 @@ def __init__(
         temperature: float = 1.0,
         delta: Optional[float] = None,
         use_bias_correction_kl: bool = False,
+        entropy_coef: float = 0.0,
     ):
         """
         Args:
@@ -416,6 +430,7 @@ def __init__(
         self.temperature = temperature
         self.delta = delta
         self.use_bias_correction_kl = use_bias_correction_kl
+        self.entropy_coef = entropy_coef
 
     def forward(
         self,
@@ -459,4 +474,5 @@ def forward(
             vllm_is_ratio,
             self.delta,
             self.use_bias_correction_kl,
+            self.entropy_coef,
         )