Feature/cispo algorithm

.gitignore

@@ -203,3 +203,7 @@ wandb/
 outputs/
 sympy/
 !/docs/figures/*
+*.whl
+
+# Apptainer
+areal.sif

areal/api/cli_args.py

@@ -1469,6 +1469,17 @@ class PPOActorConfig(TrainEngineConfig):
         metadata={"help": "SAPO temperature for negative advantages"},
     )
 
+
+    # CISPO (Clipped IS-weight Policy Optimization) - MiniMax-M1
+    use_cispo_loss: bool = field(
+        default=False,
+        metadata={"help": "Use CISPO loss (clipped importance sampling weight, detached)"},
+    )
+    cispo_epsilon_high: float = field(
+        default=1.2,
+        metadata={"help": "CISPO upper clamp for importance sampling weight"},
+    )
+
     # Asynchronous RL
     recompute_logprob: bool = field(
         default=False,

areal/trainer/ppo/actor.py

@@ -34,6 +34,7 @@
     ppo_actor_loss_fn,
     reward_overlong_penalty,
     sapo_loss_fn,
+    cispo_loss_fn,
 )
 from areal.utils.perf_tracer import trace_perf
 
@@ -355,6 +356,8 @@ def _ppo_update(self, data: dict[str, Any]) -> None:
                         importance_sampling_level=self.config.importance_sampling_level,
                         current_version=current_version,
                         prox_logp_method=self.config.prox_logp_method,
+                        use_cispo_loss=self.config.use_cispo_loss,
+                        cispo_epsilon_high=self.config.cispo_epsilon_high,
                         use_sapo_loss=self.config.use_sapo_loss,
                         sapo_tau_pos=self.config.sapo_tau_pos,
                         sapo_tau_neg=self.config.sapo_tau_neg,
@@ -420,6 +423,8 @@ def grpo_loss_fn(
     use_sapo_loss: bool = False,
     sapo_tau_pos: float = 1.0,
     sapo_tau_neg: float = 1.05,
+    use_cispo_loss: bool = False,
+    cispo_epsilon_high: float = 1.2,
     use_decoupled_loss: bool = False,
     vocab_min_logits: torch.Tensor | None = None,
     vocab_max_logits: torch.Tensor | None = None,
@@ -447,8 +452,16 @@ def grpo_loss_fn(
     if m2_threshold is not None:
         loss_mask = _apply_m2po_masking(old_logp, prox_logp, loss_mask, m2_threshold)
 
-    # Use SAPO or PPO loss
-    if use_sapo_loss:
+    # Use CISPO, SAPO, or PPO loss
+    if use_cispo_loss:
+        loss, stat = cispo_loss_fn(
+            logprobs=logprobs,
+            proximal_logprobs=prox_logp,
+            advantages=advantages,
+            epsilon_high=cispo_epsilon_high,
+            loss_mask=loss_mask,
+        )
+    elif use_sapo_loss:
         if use_decoupled_loss:
             raise ValueError(
                 "SAPO is not compatible with `use_decoupled_loss=True`. "

areal/utils/functional/__init__.py

@@ -8,6 +8,7 @@
     ppo_critic_loss_fn,
     reward_overlong_penalty,
     sapo_loss_fn,
+    cispo_loss_fn,
 )
 from areal.utils.functional.vocab_parallel import (
     gather_logprobs,
@@ -23,6 +24,7 @@
     "ppo_critic_loss_fn",
     "reward_overlong_penalty",
     "sapo_loss_fn",
+    "cispo_loss_fn",
     # logprobs.py
     "gather_logprobs",
     "gather_logprobs_entropy",

areal/utils/functional/functional.py

@@ -557,6 +557,40 @@ def ppo_actor_loss_fn(
     return pg_loss, stat
 
 
+def cispo_loss_fn(
+    logprobs: torch.Tensor,
+    proximal_logprobs: torch.Tensor,
+    advantages: torch.Tensor,
+    epsilon_high: float,
+    loss_mask: torch.Tensor,
+) -> tuple[torch.Tensor, dict]:
+    """CISPO (Clipped IS-weight Policy Optimization) loss from MiniMax-M1.
+
+    Instead of clipping the policy ratio like PPO/GRPO, CISPO clips the
+    importance sampling weight and detaches it, ensuring all tokens
+    (including rare reflective tokens) contribute gradients.
+
+    Loss = -clamp(exp(logp_new - logp_prox), max=epsilon_high).detach() * adv * logp_new
+    """
+    loss_mask_count = loss_mask.count_nonzero() or 1
+
+    log_ratio = torch.where(loss_mask, logprobs - proximal_logprobs, 0.0)
+    is_weight = torch.exp(log_ratio)
+    clamped_weight = torch.clamp(is_weight, max=epsilon_high).detach()
+
+    pg_loss = -clamped_weight * advantages * logprobs
+    pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count
+
+    stat = dict(
+        loss=(-clamped_weight * advantages * logprobs).detach(),
+        importance_weight=is_weight.detach(),
+        approx_kl=(logprobs - proximal_logprobs).detach(),
+        clip_mask=torch.zeros_like(loss_mask),
+        dual_clip_mask=torch.zeros_like(loss_mask),
+    )
+    return pg_loss, stat
+
+
 def sapo_loss_fn(
     logprobs: torch.Tensor,
     old_logprobs: torch.Tensor,

examples/math/gsm8k_grpo_zaratan.yaml

@@ -0,0 +1,189 @@
+experiment_name: gsm8k-grpo
+trial_name: trial0
+
+seed: 1
+enable_offload: false
+total_train_epochs: 10
+tokenizer_path: ${actor.path}
+
+cluster:
+  n_nodes: 1
+  n_gpus_per_node: 8
+  fileroot: /tmp/areal/experiments
+  name_resolve:
+    type: nfs
+    nfs_record_root: /tmp/areal/name_resolve
+
+
+scheduler:
+  type: null
+
+rollout:
+  backend: "sglang:d4p1t1"
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  max_concurrent_rollouts: 256
+  queue_size: null
+  consumer_batch_size: ${train_dataset.batch_size}
+  max_head_offpolicyness: 2
+  enable_rollout_tracing: false
+  scheduling_spec: ${actor.scheduling_spec}
+  fileroot: ${cluster.fileroot}
+  tokenizer_path: ${tokenizer_path}
+  dump_to_file: true
+
+gconfig:
+  n_samples: 4
+  min_new_tokens: 0
+  max_new_tokens: 1024
+  greedy: false
+  temperature: 1.0
+
+actor:
+  backend: "fsdp:d4p1t1"
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  path: Qwen/Qwen2.5-1.5B-Instruct
+  init_from_scratch: false
+  disable_dropout: true
+  gradient_checkpointing: true
+  dtype: bfloat16
+  mb_spec:
+    max_tokens_per_mb: 10240
+    packing_algorithm: ffd
+  optimizer:
+    type: adam
+    lr: 1.70e-5
+    weight_decay: 0.017
+    beta1: 0.9
+    beta2: 0.999
+    eps: 1e-8
+    lr_scheduler_type: constant
+    gradient_clipping: 1.0
+    warmup_steps_proportion: 0.001
+  eps_clip: 0.4
+  temperature: ${gconfig.temperature}
+  reward_scaling: 10.0
+  reward_bias: -0.5
+  kl_ctl: 0.0
+  ppo_n_minibatches: 1
+  recompute_logprob: true
+  use_decoupled_loss: true
+  rejection_sampling:
+    metric: ratio
+    upper: 5.0
+  reward_norm:
+    mean_level: group
+    std_level: group
+    group_size: ${gconfig.n_samples}
+  adv_norm:
+    mean_level: batch
+    std_level: batch
+  weight_update_mode: xccl
+  max_new_tokens: ${gconfig.max_new_tokens}
+  scheduling_spec:
+    - task_type: worker
+      port_count: 2
+      gpu: 1
+      mem: 32
+      cmd: python3 -m areal.infra.rpc.rpc_server
+      env_vars: {}
+
+ref:
+  backend: ${actor.backend}
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  path: ${actor.path}
+  init_from_scratch: false
+  disable_dropout: true
+  dtype: ${actor.dtype}
+  mb_spec:
+    max_tokens_per_mb: 10240
+    packing_algorithm: ffd
+  optimizer: null
+  scheduling_strategy:
+    type: colocation
+    target: actor
+  scheduling_spec: ${actor.scheduling_spec}
+
+# SGLang
+sglang:
+  model_path: ${actor.path}
+  random_seed: ${seed}
+  skip_tokenizer_init: true
+  dtype: ${actor.dtype}
+  max_running_requests: null
+  context_length: 32768
+  mem_fraction_static: 0.8
+
+vllm:
+  model: ${actor.path}
+  seed: ${seed}
+  skip_tokenizer_init: false
+  dtype: ${actor.dtype}
+  max_model_len: 32768
+  gpu_memory_utilization: 0.8
+
+# datasets
+train_dataset:
+  batch_size: 256
+  shuffle: true
+  pin_memory: true
+  num_workers: 4
+  path: openai/gsm8k
+  split: train
+  dataset_kwargs:
+    name: main
+  type: rl
+  max_length: 1024
+
+valid_dataset:
+  batch_size: 256
+  pin_memory: true
+  num_workers: 4
+  path: openai/gsm8k
+  split: test
+  dataset_kwargs:
+    name: main
+  type: rl
+
+# Utilities
+saver:
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  fileroot: ${cluster.fileroot}
+  freq_epochs: 1
+  freq_steps: null
+  freq_secs: null
+
+recover:
+  mode: disabled
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  fileroot: ${cluster.fileroot}
+  freq_epochs: 1
+  freq_steps: null
+  freq_secs: 3600
+
+evaluator:
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  fileroot: ${cluster.fileroot}
+  freq_epochs: 1
+  freq_steps: null
+  freq_secs: null
+
+stats_logger:
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  fileroot: ${cluster.fileroot}
+  wandb:
+    mode: disabled
+
+perf_tracer:
+  experiment_name: ${experiment_name}
+  trial_name: ${trial_name}
+  fileroot: ${cluster.fileroot}
+  enabled: false
+  session_tracer:
+    enabled: false