feat: Add CISPO (Clipped IS-weight Policy Optimization)

Add support for CISPO algorithm from MiniMax-M1 paper, which addresses
PPO/GRPO's limitation of clipping out low-probability reasoning tokens.

Changes:
- Add compute_cispo_loss() in slime/utils/ppo_utils.py
- Add 'cispo' to advantage_estimator choices
- Update reward normalization to include CISPO
- Use CISPO loss when advantage_estimator='cispo'

Key implementation details:
- Token-level IS with stop-gradient on clipped ratios
- Explicit log probability: ratio_sg * advantages * log_probs
- Upper-only clipping with default eps_clip_high=5.0
- Direct clipfrac calculation: (ratio > eps_clip_high)

Reference: MiniMax-M1 paper (arxiv:2506.13585)

Author: kekmodel

slime/backends/megatron_utils/loss.py

@@ -10,6 +10,7 @@
 from slime.utils.ppo_utils import (
     calculate_log_probs_and_entropy,
     compute_approx_kl,
+    compute_cispo_loss,
     compute_gspo_kl,
     compute_opsm_mask,
     compute_policy_loss,
@@ -239,7 +240,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
             for i in range(len(log_probs))
         ]
 
-    if args.advantage_estimator in ["grpo", "gspo"]:
+    if args.advantage_estimator in ["grpo", "gspo", "cispo"]:
         rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
         returns = get_grpo_returns(rewards, kl)
         # TODO: is the copy necessary?
@@ -449,7 +450,11 @@ def policy_loss_function(
         log_probs = torch.cat(log_probs, dim=0)
         ppo_kl = old_log_probs - log_probs
 
-    pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high)
+    # Compute policy loss: CISPO uses upper truncation with stop-gradient
+    if args.advantage_estimator == "cispo":
+        pg_loss, pg_clipfrac = compute_cispo_loss(ppo_kl, log_probs, advantages, args.eps_clip_high)
+    else:
+        pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high)
 
     if args.use_opsm:
         pg_loss = pg_loss * opsm_mask

slime/ray/rollout.py

@@ -190,7 +190,7 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
 
         raw_rewards = [sample.get_reward_value(self.args) for sample in samples]
         if (
-            self.args.advantage_estimator in ["grpo", "gspo", "reinforce_plus_plus_baseline"]
+            self.args.advantage_estimator in ["grpo", "gspo", "cispo", "reinforce_plus_plus_baseline"]
             and self.args.rewards_normalization
         ):
             # group norm
@@ -203,7 +203,7 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
             mean = rewards.mean(dim=-1, keepdim=True)
             rewards = rewards - mean
 
-            if self.args.advantage_estimator in ["grpo", "gspo"] and self.args.grpo_std_normalization:
+            if self.args.advantage_estimator in ["grpo", "gspo", "cispo"] and self.args.grpo_std_normalization:
                 std = rewards.std(dim=-1, keepdim=True)
                 rewards = rewards / (std + 1e-6)
 

slime/utils/arguments.py

@@ -701,6 +701,7 @@ def add_algo_arguments(parser):
                 choices=[
                     "grpo",
                     "gspo",
+                    "cispo",
                     "reinforce_plus_plus",
                     "reinforce_plus_plus_baseline",
                     "ppo",

slime/utils/ppo_utils.py

@@ -144,6 +144,57 @@ def compute_policy_loss(
     return pg_losses, clipfrac
 
 
+@torch.compile(dynamic=True)
+def compute_cispo_loss(
+    ppo_kl: torch.Tensor,
+    log_probs: torch.Tensor,
+    advantages: torch.Tensor,
+    eps_clip_high: float,
+):
+    """Compute CISPO (Clipped IS-weight Policy Optimization) loss.
+
+    CISPO applies upper truncation on the importance sampling ratio with
+    stop-gradient, preventing the ratio itself from being learned. This differs
+    from PPO which uses both upper and lower clipping without stop-gradient.
+
+    The key formula from the paper:
+        ratio = exp(log π_current - log π_old)
+        ratio_truncated = min(ratio, ε_max)
+        loss = -sg(ratio_truncated) * advantages * log(π_current)
+
+    Note: log_probs is explicitly multiplied so gradient flows through it,
+    while ratio_sg is detached to prevent learning the ratio itself.
+
+    Args:
+        ppo_kl: Log-ratio (log π_old - log π_current) for each token
+        log_probs: Current policy log probabilities (requires gradient)
+        advantages: Advantage estimates for each token
+        eps_clip_high: Upper bound for clipping (ε_max), typically 5.0 (absolute value)
+
+    Returns:
+        Tuple of (pg_losses, clipfrac) where:
+            - pg_losses: Per-token CISPO policy gradient losses
+            - clipfrac: Fraction of ratios that were clipped
+    """
+    # Compute importance sampling ratio: π_current / π_old
+    ratio = (-ppo_kl).exp()
+
+    # Upper truncation: min(ratio, ε_max) where ε_max is absolute value
+    ratio_truncated = torch.clamp(ratio, max=eps_clip_high)
+
+    # Stop-gradient: prevent the ratio from being learned (CISPO's key feature)
+    ratio_sg = ratio_truncated.detach()
+
+    # CISPO formula: sg(ratio) * advantages * log_probs
+    # This ensures gradient flows through log_probs but not through ratio
+    pg_losses = -ratio_sg * advantages * log_probs
+
+    # Track clipping fraction for monitoring
+    clipfrac = (ratio > eps_clip_high).float()
+
+    return pg_losses, clipfrac
+
+
 def compute_log_probs(logits: torch.Tensor, tokens: torch.Tensor, process_group: dist.ProcessGroup | None):
     from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy