diff --git a/--bind b/--bind new file mode 100644 index 0000000000..e69de29bb2 diff --git a/--config b/--config new file mode 100644 index 0000000000..e69de29bb2 diff --git a/--env b/--env new file mode 100644 index 0000000000..e69de29bb2 diff --git a/.gitignore b/.gitignore index dc8f74212f..31e5e9b1cc 100644 --- a/.gitignore +++ b/.gitignore @@ -203,3 +203,7 @@ wandb/ outputs/ sympy/ !/docs/figures/* +*.whl + +# Apptainer +areal.sif diff --git a/Python-3.12.7.tgz b/Python-3.12.7.tgz new file mode 100644 index 0000000000..623ae12f8c Binary files /dev/null and b/Python-3.12.7.tgz differ diff --git a/actor.backend=fsdp:d1p1t1 b/actor.backend=fsdp:d1p1t1 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/actor.cispo_epsilon_high=1.2 b/actor.cispo_epsilon_high=1.2 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/actor.path= b/actor.path= new file mode 100644 index 0000000000..e69de29bb2 diff --git a/actor.use_cispo_loss=true b/actor.use_cispo_loss=true new file mode 100644 index 0000000000..e69de29bb2 diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 7e08f2d63a..a3609cc264 100644 --- a/areal/api/cli_args.py +++ b/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, diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 07944a31a5..b7693f6e04 100644 --- a/areal/trainer/ppo/actor.py +++ b/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`. " diff --git a/areal/utils/functional/__init__.py b/areal/utils/functional/__init__.py index 369eaf9149..6fd949f6b5 100644 --- a/areal/utils/functional/__init__.py +++ b/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", diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py index a79aab6024..2ce41270fe 100644 --- a/areal/utils/functional/functional.py +++ b/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, diff --git a/cluster.n_gpus_per_node=2 b/cluster.n_gpus_per_node=2 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/cluster.n_nodes=1 b/cluster.n_nodes=1 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/--config b/examples/--config new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/math/gsm8k_grpo_zaratan.yaml b/examples/math/gsm8k_grpo_zaratan.yaml new file mode 100644 index 0000000000..941ee0cb39 --- /dev/null +++ b/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 diff --git a/examples/scheduler.type=local b/examples/scheduler.type=local new file mode 100644 index 0000000000..e69de29bb2 diff --git a/experiment_name=cispo_eta2 b/experiment_name=cispo_eta2 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/final_script.sh b/final_script.sh new file mode 100644 index 0000000000..53b400acca --- /dev/null +++ b/final_script.sh @@ -0,0 +1,96 @@ +#!/bin/bash +# ============================================================ +# AReaL Experiment Matrix: GRPO vs CISPO vs SAPO x eta={0,2,4} +# 10 epochs with test set evaluation every 20 steps +# Usage: bash submit_all_v2.sh +# ============================================================ + +USER_NAME=${1:-sgnanaku} +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/$USER_NAME +LOGS=$SCRATCH/slurm_logs_v2 +mkdir -p $LOGS + +echo "Submitting jobs for user: $USER_NAME" +echo "Logs: $LOGS" +echo "" + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + total_train_epochs=10 \\\\ + train_dataset.batch_size=64 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + evaluator.freq_steps=20 \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# GRPO (baseline) +submit_job "grpo_eta0" "" "0" +submit_job "grpo_eta2" "" "2" +submit_job "grpo_eta4" "" "4" + +# CISPO +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# SAPO +submit_job "sapo_eta0" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "4" + +# M2PO +submit_job "m2po_eta0" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "0" +submit_job "m2po_eta2" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "2" +submit_job "m2po_eta4" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "4" + +echo "" +echo "All 9 jobs submitted! Each runs for 10 epochs (~10 hours)" +echo "Monitor with: squeue -u ${USER_NAME}" +echo "Logs at: ${LOGS}" diff --git a/gconfig.max_new_tokens=512 b/gconfig.max_new_tokens=512 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/get_results.py b/get_results.py new file mode 100644 index 0000000000..618e52c1d0 --- /dev/null +++ b/get_results.py @@ -0,0 +1,81 @@ +# !/usr/bin/env python3 +# ============================================================ +# Extract and display results from all AReaL experiment logs +# Usage: python3 get_results.py [username] +# Example: python3 get_results.py sgnanaku +# ============================================================ + +import re +import glob +import sys +import os + +username = sys.argv[1] if len(sys.argv) > 1 else "sgnanaku" +log_dir = f"/scratch/zt1/project/zaoxing-prj/user/{username}/slurm_logs" + +# Map experiment names to their log files (pick the latest one) +experiments = [ + "grpo_eta0", "grpo_eta2", "grpo_eta4", + "cispo_eta0", "cispo_eta2", "cispo_eta4", + "sapo_eta0", "sapo_eta2", "sapo_eta4", +] + +def find_latest_log(exp_name, log_dir): + pattern = os.path.join(log_dir, f"{exp_name}_*.out") + files = glob.glob(pattern) + if not files: + # also try with _final_ suffix + pattern = os.path.join(log_dir, f"{exp_name}_final_*.out") + files = glob.glob(pattern) + if not files: + return None + # Return the one with training data, preferring latest + files_with_data = [] + for f in files: + with open(f) as fh: + content = fh.read() + if "Train step" in content: + files_with_data.append(f) + if files_with_data: + return sorted(files_with_data)[-1] + return sorted(files)[-1] + +print(f"\n{'='*60}") +print(f"AReaL Experiment Results - User: {username}") +print(f"{'='*60}") +print(f"{'Experiment':<15} {'Final Reward':>12} {'Max Reward':>10} {'Steps':>8} {'Status':>10}") +print("-" * 60) + +results = {} +for exp in experiments: + log_file = find_latest_log(exp, log_dir) + if not log_file: + print(f"{exp:<15} {'N/A':>12} {'N/A':>10} {'N/A':>8} {'NO LOG':>10}") + continue + + with open(log_file) as f: + content = f.read() + + rewards = re.findall(r'task_reward/avg\s*│\s*([\d.e+\-]+)', content) + steps = re.findall(r'Train step (\d+)/\d+ done', content) + + if rewards: + final_reward = float(rewards[-1]) + max_reward = max(float(r) for r in rewards) + n_steps = steps[-1] if steps else "?" + status = "OK" + results[exp] = final_reward + print(f"{exp:<15} {final_reward:>12.4f} {max_reward:>10.4f} {n_steps:>8} {status:>10}") + else: + print(f"{exp:<15} {'N/A':>12} {'N/A':>10} {'N/A':>8} {'CRASHED':>10}") + +print(f"\n{'='*60}") +print("SUMMARY TABLE (Final Reward)") +print(f"{'='*60}") +print(f"{'Algorithm':<10} {'eta=0':>10} {'eta=2':>10} {'eta=4':>10}") +print("-" * 45) +for algo in ["grpo", "cispo", "sapo"]: + r0 = results.get(f"{algo}_eta0", float('nan')) + r2 = results.get(f"{algo}_eta2", float('nan')) + r4 = results.get(f"{algo}_eta4", float('nan')) + print(f"{algo.upper():<10} {r0:>10.4f} {r2:>10.4f} {r4:>10.4f}") diff --git a/no_trust_region.sh b/no_trust_region.sh new file mode 100755 index 0000000000..d9d4f0cd42 --- /dev/null +++ b/no_trust_region.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# ============================================================ +# AReaL Experiment Matrix: GRPO vs CISPO vs SAPO x eta={0,2,4} +# 10 epochs with test set evaluation every 20 steps +# Usage: bash submit_all_v2.sh +# ============================================================ + +USER_NAME=${1:-sgnanaku} +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/$USER_NAME +LOGS=$SCRATCH/slurm_logs_v2 +mkdir -p $LOGS + +echo "Submitting jobs for user: $USER_NAME" +echo "Logs: $LOGS" +echo "" + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + total_train_epochs=10 \\\\ + train_dataset.batch_size=32 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + evaluator.freq_steps=20 \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# GRPO (baseline) +submit_job "grpo_eta0" "+eps_clip=1e9" "0" +submit_job "grpo_eta2" "+eps_clip=1e9" "2" +submit_job "grpo_eta4" "+eps_clip=1e9" "4" + +# CISPO +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1e9" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1e9" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1e9" "4" + +# SAPO +submit_job "sapo_eta0" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false ++tau_pos=0" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false ++tau_pos=0" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false ++tau_pos=0" "4" + +# M2PO +submit_job "m2po_eta0" "++actor.m2_threshold=0.04" "0" +submit_job "m2po_eta2" "++actor.m2_threshold=0.04" "2" +submit_job "m2po_eta4" "++actor.m2_threshold=0.04" "4" + +echo "" +echo "All 9 jobs submitted! Each runs for 10 epochs (~10 hours)" +echo "Monitor with: squeue -u ${USER_NAME}" +echo "Logs at: ${LOGS}" diff --git a/plot_results.py b/plot_results.py new file mode 100644 index 0000000000..fb31573866 --- /dev/null +++ b/plot_results.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +# ============================================================ +# Plot results from AReaL experiment logs +# Usage: python3 plot_results.py [username] [log_subdir] +# Saves plots to: /plots/ (default log_subdir: slurm_logs) +# ============================================================ + +import re +import glob +import sys +import os +import json + +username = sys.argv[1] if len(sys.argv) > 1 else "sgnanaku" +log_subdir = sys.argv[2] if len(sys.argv) > 2 else "slurm_logs" +log_dir = f"/scratch/zt1/project/zaoxing-prj/user/{username}/{log_subdir}" +plot_dir = os.path.join(log_dir, "plots") +os.makedirs(plot_dir, exist_ok=True) + +try: + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + import numpy as np +except ImportError: + print("Installing matplotlib...") + os.system("pip install matplotlib numpy --quiet") + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + import numpy as np + +# ============================================================ +# Config +# ============================================================ +COLORS = { + 'grpo': '#2196F3', # blue + 'cispo': '#FF5722', # orange-red + 'sapo': '#4CAF50', # green + 'm2po': '#9C27B0', # purple +} +LINESTYLES = { + 0: '-', + 2: '--', + 4: ':', +} +ETA_LABELS = {0: 'η=0 (sync)', 2: 'η=2', 4: 'η=4'} +ALGO_LABELS = {'grpo': 'GRPO', 'cispo': 'CISPO', 'sapo': 'SAPO', 'm2po': 'M2PO'} + +EXPERIMENTS = { + 'grpo_eta0': ('grpo', 0), + 'grpo_eta2': ('grpo', 2), + 'grpo_eta4': ('grpo', 4), + 'cispo_eta0': ('cispo', 0), + 'cispo_eta2': ('cispo', 2), + 'cispo_eta4': ('cispo', 4), + 'sapo_eta0': ('sapo', 0), + 'sapo_eta2': ('sapo', 2), + 'sapo_eta4': ('sapo', 4), + 'm2po_eta0': ('m2po', 0), + 'm2po_eta2': ('m2po', 2), + 'm2po_eta4': ('m2po', 4), +} + +# ============================================================ +# Parse logs +# ============================================================ +def find_latest_log(exp_name, log_dir): + """Find the log file with the most training steps (handles multiple job submissions)""" + best_file = None + max_rewards = 0 + + for pattern in [f"{exp_name}_final_*.out", f"{exp_name}_*.out"]: + files = glob.glob(os.path.join(log_dir, pattern)) + for f in files: + try: + with open(f) as file: + content = file.read() + if "Train step" in content: + reward_count = len(re.findall(r'(?:ppo_actor/)?task_reward/avg', content)) + if reward_count > max_rewards: + max_rewards = reward_count + best_file = f + except: + pass + + return best_file + +def parse_log(log_file): + with open(log_file) as f: + content = f.read() + rewards = [float(x) for x in re.findall(r'(?:ppo_actor/)?task_reward/avg\s*│\s*([\d.e+\-]+)', content)] + grad_norms = [float(x) for x in re.findall(r'update/grad_norm\s*│\s*([\d.e+\-]+)', content)] + entropy = [float(x) for x in re.findall(r'update/entropy/avg\s*│\s*([\d.e+\-]+)', content)] + + # Find last non-zero reward for crashed jobs + final_reward = None + if rewards: + for r in reversed(rewards): + if r > 0: + final_reward = r + break + if final_reward is None: + final_reward = rewards[-1] # Use last reward if all are zero + + # Eval reward appears every ~20 train steps inside one of the per-step tables. + # Align each eval value to the training step it was logged in by counting + # how many `task_reward/avg` markers precede the eval's position in the file. + train_positions = [m.start() for m in re.finditer(r'(?:ppo_actor/)?task_reward/avg', content)] + eval_rewards, eval_steps = [], [] + for em in re.finditer(r'eval-rollout/reward\s*│\s*([\d.e+\-]+)', content): + step = sum(1 for tp in train_positions if tp <= em.start()) + eval_rewards.append(float(em.group(1))) + eval_steps.append(step) + + return { + 'rewards': rewards, + 'grad_norms': grad_norms, + 'entropy': entropy, + 'steps': list(range(1, len(rewards) + 1)), + 'eval_rewards': eval_rewards, + 'eval_steps': eval_steps, + 'final_reward': final_reward, + } + +data = {} +for exp_name, (algo, eta) in EXPERIMENTS.items(): + log_file = find_latest_log(exp_name, log_dir) + if log_file: + data[exp_name] = parse_log(log_file) + data[exp_name]['algo'] = algo + data[exp_name]['eta'] = eta + print(f"Loaded {exp_name}: {len(data[exp_name]['rewards'])} steps, " + f"{len(data[exp_name]['eval_rewards'])} evals") + else: + print(f"Missing: {exp_name}") + +# ============================================================ +# Plot 1: Reward curves per algorithm (4 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +fig.suptitle('Reward Curves by Algorithm (GRPO vs CISPO vs SAPO vs M2PO)', fontsize=14, fontweight='bold') + +for ax, algo in zip(axes, ['grpo', 'cispo', 'sapo', 'm2po']): + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + d = data[exp] + ax.plot(d['steps'], d['rewards'], + color=COLORS[algo], + linestyle=LINESTYLES[eta], + linewidth=2, + label=ETA_LABELS[eta]) + ax.set_title(ALGO_LABELS[algo], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Task Reward' if ax is axes[0] else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.5, 1.0) + +plt.tight_layout() +path1 = os.path.join(plot_dir, 'reward_curves_by_algo.png') +plt.savefig(path1, dpi=150, bbox_inches='tight') +print(f"Saved: {path1}") +plt.close() + +# ============================================================ +# Plot 2: Reward curves per staleness level (3 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) +fig.suptitle('Reward Curves by Staleness Level (η=0, 2, 4)', fontsize=14, fontweight='bold') + +for ax, eta in zip(axes, [0, 2, 4]): + for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + d = data[exp] + ax.plot(d['steps'], d['rewards'], + color=COLORS[algo], + linewidth=2, + label=ALGO_LABELS[algo]) + ax.set_title(ETA_LABELS[eta], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Task Reward' if eta == 0 else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.5, 1.0) + +plt.tight_layout() +path2 = os.path.join(plot_dir, 'reward_curves_by_eta.png') +plt.savefig(path2, dpi=150, bbox_inches='tight') +print(f"Saved: {path2}") +plt.close() + +# ============================================================ +# Plot 2b: Eval reward curves per algorithm (4 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +fig.suptitle('Eval Reward Curves by Algorithm (GRPO vs CISPO vs SAPO vs M2PO)', fontsize=14, fontweight='bold') + +for ax, algo in zip(axes, ['grpo', 'cispo', 'sapo', 'm2po']): + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp].get('eval_rewards'): + d = data[exp] + ax.plot(d['eval_steps'], d['eval_rewards'], + color=COLORS[algo], + linestyle=LINESTYLES[eta], + linewidth=2, + marker='o', + markersize=4, + label=ETA_LABELS[eta]) + ax.set_title(ALGO_LABELS[algo], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Eval Reward' if ax is axes[0] else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + +plt.tight_layout() +path_eval_algo = os.path.join(plot_dir, 'eval_curves_by_algo.png') +plt.savefig(path_eval_algo, dpi=150, bbox_inches='tight') +print(f"Saved: {path_eval_algo}") +plt.close() + +# ============================================================ +# Plot 2c: Eval reward curves per staleness level (3 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) +fig.suptitle('Eval Reward Curves by Staleness Level (η=0, 2, 4)', fontsize=14, fontweight='bold') + +for ax, eta in zip(axes, [0, 2, 4]): + for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp].get('eval_rewards'): + d = data[exp] + ax.plot(d['eval_steps'], d['eval_rewards'], + color=COLORS[algo], + linewidth=2, + marker='o', + markersize=4, + label=ALGO_LABELS[algo]) + ax.set_title(ETA_LABELS[eta], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Eval Reward' if eta == 0 else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + +plt.tight_layout() +path_eval_eta = os.path.join(plot_dir, 'eval_curves_by_eta.png') +plt.savefig(path_eval_eta, dpi=150, bbox_inches='tight') +print(f"Saved: {path_eval_eta}") +plt.close() + +# ============================================================ +# Plot 3: Final reward heatmap (4x3 grid) +# ============================================================ +fig, ax = plt.subplots(figsize=(8, 6)) + +algos = ['GRPO', 'CISPO', 'SAPO', 'M2PO'] +etas = ['η=0', 'η=2', 'η=4'] +matrix = np.zeros((4, 3)) + +for i, algo in enumerate(['grpo', 'cispo', 'sapo', 'm2po']): + for j, eta in enumerate([0, 2, 4]): + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + matrix[i, j] = data[exp].get('final_reward', data[exp]['rewards'][-1]) + +im = ax.imshow(matrix, cmap='RdYlGn', vmin=0.7, vmax=0.95) +ax.set_xticks(range(3)) +ax.set_yticks(range(4)) +ax.set_xticklabels(etas, fontsize=12) +ax.set_yticklabels(algos, fontsize=12) +ax.set_xlabel('Staleness (η)', fontsize=12) +ax.set_ylabel('Algorithm', fontsize=12) +ax.set_title('Final Reward Heatmap', fontsize=14, fontweight='bold') + +for i in range(4): + for j in range(3): + val = matrix[i, j] + ax.text(j, i, f'{val:.4f}', ha='center', va='center', + fontsize=11, fontweight='bold', + color='black' if 0.75 < val < 0.92 else 'white') + +plt.colorbar(im, ax=ax, label='Final Reward') +plt.tight_layout() +path3 = os.path.join(plot_dir, 'final_reward_heatmap.png') +plt.savefig(path3, dpi=150, bbox_inches='tight') +print(f"Saved: {path3}") +plt.close() + +# ============================================================ +# Plot 4: Staleness effect (line plot, final reward vs eta) +# ============================================================ +fig, ax = plt.subplots(figsize=(8, 6)) + +for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + rewards = [] + valid_etas = [] + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + rewards.append(data[exp].get('final_reward', data[exp]['rewards'][-1])) + valid_etas.append(eta) + if rewards: + ax.plot(valid_etas, rewards, + color=COLORS[algo], + marker='o', + markersize=10, + linewidth=2.5, + label=ALGO_LABELS[algo]) + for eta, r in zip(valid_etas, rewards): + ax.annotate(f'{r:.4f}', (eta, r), + textcoords="offset points", + xytext=(0, 10), + ha='center', fontsize=9) + +ax.set_xlabel('Staleness Level (η)', fontsize=12) +ax.set_ylabel('Final Task Reward', fontsize=12) +ax.set_title('Effect of Staleness on Final Reward', fontsize=14, fontweight='bold') +ax.set_xticks([0, 2, 4]) +ax.set_xticklabels(['η=0\n(sync)', 'η=2\n(default)', 'η=4\n(high)']) +ax.legend(fontsize=11) +ax.grid(True, alpha=0.3) +ax.set_ylim(0.7, 1.0) + +plt.tight_layout() +path4 = os.path.join(plot_dir, 'staleness_effect.png') +plt.savefig(path4, dpi=150, bbox_inches='tight') +print(f"Saved: {path4}") +plt.close() + +print(f"\nAll plots saved to: {plot_dir}") diff --git a/plot_results.py.backup b/plot_results.py.backup new file mode 100644 index 0000000000..37cdae716e --- /dev/null +++ b/plot_results.py.backup @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +# ============================================================ +# Plot results from AReaL experiment logs +# Usage: python3 plot_results.py [username] [log_subdir] +# Saves plots to: /plots/ (default log_subdir: slurm_logs) +# ============================================================ + +import re +import glob +import sys +import os +import json + +username = sys.argv[1] if len(sys.argv) > 1 else "sgnanaku" +log_subdir = sys.argv[2] if len(sys.argv) > 2 else "slurm_logs" +log_dir = f"/scratch/zt1/project/zaoxing-prj/user/{username}/{log_subdir}" +plot_dir = os.path.join(log_dir, "plots") +os.makedirs(plot_dir, exist_ok=True) + +try: + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + import numpy as np +except ImportError: + print("Installing matplotlib...") + os.system("pip install matplotlib numpy --quiet") + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + import numpy as np + +# ============================================================ +# Config +# ============================================================ +COLORS = { + 'grpo': '#2196F3', # blue + 'cispo': '#FF5722', # orange-red + 'sapo': '#4CAF50', # green + 'm2po': '#9C27B0', # purple +} +LINESTYLES = { + 0: '-', + 2: '--', + 4: ':', +} +ETA_LABELS = {0: 'η=0 (sync)', 2: 'η=2', 4: 'η=4'} +ALGO_LABELS = {'grpo': 'GRPO', 'cispo': 'CISPO', 'sapo': 'SAPO', 'm2po': 'M2PO'} + +EXPERIMENTS = { + 'grpo_eta0': ('grpo', 0), + 'grpo_eta2': ('grpo', 2), + 'grpo_eta4': ('grpo', 4), + 'cispo_eta0': ('cispo', 0), + 'cispo_eta2': ('cispo', 2), + 'cispo_eta4': ('cispo', 4), + 'sapo_eta0': ('sapo', 0), + 'sapo_eta2': ('sapo', 2), + 'sapo_eta4': ('sapo', 4), + 'm2po_eta0': ('m2po', 0), + 'm2po_eta2': ('m2po', 2), + 'm2po_eta4': ('m2po', 4), +} + +# ============================================================ +# Parse logs +# ============================================================ +def find_latest_log(exp_name, log_dir): + for pattern in [f"{exp_name}_final_*.out", f"{exp_name}_*.out"]: + files = glob.glob(os.path.join(log_dir, pattern)) + files_with_data = [f for f in files if "Train step" in open(f).read()] + if files_with_data: + return sorted(files_with_data)[-1] + return None + +def parse_log(log_file): + with open(log_file) as f: + content = f.read() + rewards = [float(x) for x in re.findall(r'task_reward/avg\s*│\s*([\d.e+\-]+)', content)] + grad_norms = [float(x) for x in re.findall(r'update/grad_norm\s*│\s*([\d.e+\-]+)', content)] + entropy = [float(x) for x in re.findall(r'update/entropy/avg\s*│\s*([\d.e+\-]+)', content)] + + # Eval reward appears every ~20 train steps inside one of the per-step tables. + # Align each eval value to the training step it was logged in by counting + # how many `task_reward/avg` markers precede the eval's position in the file. + train_positions = [m.start() for m in re.finditer(r'task_reward/avg', content)] + eval_rewards, eval_steps = [], [] + for em in re.finditer(r'eval-rollout/reward\s*│\s*([\d.e+\-]+)', content): + step = sum(1 for tp in train_positions if tp <= em.start()) + eval_rewards.append(float(em.group(1))) + eval_steps.append(step) + + return { + 'rewards': rewards, + 'grad_norms': grad_norms, + 'entropy': entropy, + 'steps': list(range(1, len(rewards) + 1)), + 'eval_rewards': eval_rewards, + 'eval_steps': eval_steps, + } + +data = {} +for exp_name, (algo, eta) in EXPERIMENTS.items(): + log_file = find_latest_log(exp_name, log_dir) + if log_file: + data[exp_name] = parse_log(log_file) + data[exp_name]['algo'] = algo + data[exp_name]['eta'] = eta + print(f"Loaded {exp_name}: {len(data[exp_name]['rewards'])} steps, " + f"{len(data[exp_name]['eval_rewards'])} evals") + else: + print(f"Missing: {exp_name}") + +# ============================================================ +# Plot 1: Reward curves per algorithm (3 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +fig.suptitle('Reward Curves by Algorithm (GRPO vs CISPO vs SAPO vs M2PO)', fontsize=14, fontweight='bold') + +for ax, algo in zip(axes, ['grpo', 'cispo', 'sapo', 'm2po']): + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + d = data[exp] + ax.plot(d['steps'], d['rewards'], + color=COLORS[algo], + linestyle=LINESTYLES[eta], + linewidth=2, + label=ETA_LABELS[eta]) + ax.set_title(ALGO_LABELS[algo], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Task Reward' if ax is axes[0] else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.5, 1.0) + +plt.tight_layout() +path1 = os.path.join(plot_dir, 'reward_curves_by_algo.png') +plt.savefig(path1, dpi=150, bbox_inches='tight') +print(f"Saved: {path1}") +plt.close() + +# ============================================================ +# Plot 2: Reward curves per staleness level (3 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) +fig.suptitle('Reward Curves by Staleness Level (η=0, 2, 4)', fontsize=14, fontweight='bold') + +for ax, eta in zip(axes, [0, 2, 4]): + for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + d = data[exp] + ax.plot(d['steps'], d['rewards'], + color=COLORS[algo], + linewidth=2, + label=ALGO_LABELS[algo]) + ax.set_title(ETA_LABELS[eta], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Task Reward' if eta == 0 else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.5, 1.0) + +plt.tight_layout() +path2 = os.path.join(plot_dir, 'reward_curves_by_eta.png') +plt.savefig(path2, dpi=150, bbox_inches='tight') +print(f"Saved: {path2}") +plt.close() + +# ============================================================ +# Plot 2b: Eval reward curves per algorithm (4 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +fig.suptitle('Eval Reward Curves by Algorithm (GRPO vs CISPO vs SAPO vs M2PO)', fontsize=14, fontweight='bold') + +for ax, algo in zip(axes, ['grpo', 'cispo', 'sapo', 'm2po']): + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp].get('eval_rewards'): + d = data[exp] + ax.plot(d['eval_steps'], d['eval_rewards'], + color=COLORS[algo], + linestyle=LINESTYLES[eta], + linewidth=2, + marker='o', + markersize=4, + label=ETA_LABELS[eta]) + ax.set_title(ALGO_LABELS[algo], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Eval Reward' if ax is axes[0] else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + +plt.tight_layout() +path_eval_algo = os.path.join(plot_dir, 'eval_curves_by_algo.png') +plt.savefig(path_eval_algo, dpi=150, bbox_inches='tight') +print(f"Saved: {path_eval_algo}") +plt.close() + +# ============================================================ +# Plot 2c: Eval reward curves per staleness level (3 subplots) +# ============================================================ +fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) +fig.suptitle('Eval Reward Curves by Staleness Level (η=0, 2, 4)', fontsize=14, fontweight='bold') + +for ax, eta in zip(axes, [0, 2, 4]): + for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp].get('eval_rewards'): + d = data[exp] + ax.plot(d['eval_steps'], d['eval_rewards'], + color=COLORS[algo], + linewidth=2, + marker='o', + markersize=4, + label=ALGO_LABELS[algo]) + ax.set_title(ETA_LABELS[eta], fontsize=12, fontweight='bold') + ax.set_xlabel('Training Step') + ax.set_ylabel('Eval Reward' if eta == 0 else '') + ax.legend() + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + +plt.tight_layout() +path_eval_eta = os.path.join(plot_dir, 'eval_curves_by_eta.png') +plt.savefig(path_eval_eta, dpi=150, bbox_inches='tight') +print(f"Saved: {path_eval_eta}") +plt.close() + +# ============================================================ +# Plot 3: Final reward heatmap (3x3 grid) +# ============================================================ +fig, ax = plt.subplots(figsize=(8, 6)) + +algos = ['GRPO', 'CISPO', 'SAPO', 'M2PO'] +etas = ['η=0', 'η=2', 'η=4'] +matrix = np.zeros((4, 3)) + +for i, algo in enumerate(['grpo', 'cispo', 'sapo', 'm2po']): + for j, eta in enumerate([0, 2, 4]): + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + matrix[i, j] = data[exp]['rewards'][-1] + +im = ax.imshow(matrix, cmap='RdYlGn', vmin=0.7, vmax=0.95) +ax.set_xticks(range(3)) +ax.set_yticks(range(4)) +ax.set_xticklabels(etas, fontsize=12) +ax.set_yticklabels(algos, fontsize=12) +ax.set_xlabel('Staleness (η)', fontsize=12) +ax.set_ylabel('Algorithm', fontsize=12) +ax.set_title('Final Reward Heatmap', fontsize=14, fontweight='bold') + +for i in range(4): + for j in range(3): + val = matrix[i, j] + ax.text(j, i, f'{val:.4f}', ha='center', va='center', + fontsize=11, fontweight='bold', + color='black' if 0.75 < val < 0.92 else 'white') + +plt.colorbar(im, ax=ax, label='Final Reward') +plt.tight_layout() +path3 = os.path.join(plot_dir, 'final_reward_heatmap.png') +plt.savefig(path3, dpi=150, bbox_inches='tight') +print(f"Saved: {path3}") +plt.close() + +# ============================================================ +# Plot 4: Staleness effect (line plot, final reward vs eta) +# ============================================================ +fig, ax = plt.subplots(figsize=(8, 6)) + +for algo in ['grpo', 'cispo', 'sapo', 'm2po']: + rewards = [] + valid_etas = [] + for eta in [0, 2, 4]: + exp = f"{algo}_eta{eta}" + if exp in data and data[exp]['rewards']: + rewards.append(data[exp]['rewards'][-1]) + valid_etas.append(eta) + if rewards: + ax.plot(valid_etas, rewards, + color=COLORS[algo], + marker='o', + markersize=10, + linewidth=2.5, + label=ALGO_LABELS[algo]) + for eta, r in zip(valid_etas, rewards): + ax.annotate(f'{r:.4f}', (eta, r), + textcoords="offset points", + xytext=(0, 10), + ha='center', fontsize=9) + +ax.set_xlabel('Staleness Level (η)', fontsize=12) +ax.set_ylabel('Final Task Reward', fontsize=12) +ax.set_title('Effect of Staleness on Final Reward', fontsize=14, fontweight='bold') +ax.set_xticks([0, 2, 4]) +ax.set_xticklabels(['η=0\n(sync)', 'η=2\n(default)', 'η=4\n(high)']) +ax.legend(fontsize=11) +ax.grid(True, alpha=0.3) +ax.set_ylim(0.7, 1.0) + +plt.tight_layout() +path4 = os.path.join(plot_dir, 'staleness_effect.png') +plt.savefig(path4, dpi=150, bbox_inches='tight') +print(f"Saved: {path4}") +plt.close() + +print(f"\nAll plots saved to: {plot_dir}") diff --git a/resubmit_failed.sh b/resubmit_failed.sh new file mode 100644 index 0000000000..07f8004bc6 --- /dev/null +++ b/resubmit_failed.sh @@ -0,0 +1,70 @@ +#!/bin/bash +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/sgnanaku +LOGS=$SCRATCH/slurm_logs + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + train_dataset.batch_size=64 \\\\ + total_train_epochs=1 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + ${ALGO_FLAGS} + " +SBATCH + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# CISPO +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# SAPO +submit_job "sapo_eta0" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "4" + +echo "All 6 jobs resubmitted!" diff --git a/resubmit_failed_v2.sh b/resubmit_failed_v2.sh new file mode 100644 index 0000000000..2182e2e686 --- /dev/null +++ b/resubmit_failed_v2.sh @@ -0,0 +1,91 @@ +#!/bin/bash +# ============================================================ +# Resubmit failed jobs with batch_size=32 to fix OOM +# Usage: bash resubmit_failed_v2.sh +# ============================================================ + +USER_NAME=${1:-sgnanaku} +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/$USER_NAME +LOGS=$SCRATCH/slurm_logs_v2 +mkdir -p $LOGS + +echo "Resubmitting failed jobs for user: $USER_NAME" +echo "" + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + total_train_epochs=10 \\\\ + train_dataset.batch_size=32 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + evaluator.freq_steps=20 \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# OOM jobs - resubmit with batch_size=32 +submit_job "grpo_eta0" "" "0" +submit_job "grpo_eta2" "" "2" +submit_job "grpo_eta4" "" "4" +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# Missing SAPO jobs +submit_job "sapo_eta0" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "4" + +# M2PO - fixed (no use_m2po_loss flag, just set threshold) +submit_job "m2po_eta0" "++actor.m2_threshold=0.04" "0" +submit_job "m2po_eta2" "++actor.m2_threshold=0.04" "2" +submit_job "m2po_eta4" "++actor.m2_threshold=0.04" "4" + +echo "" +echo "All failed jobs resubmitted with batch_size=32!" +echo "Monitor with: squeue -u ${USER_NAME}" +echo "Logs at: ${LOGS}" diff --git a/resubmit_failed_v3.sh b/resubmit_failed_v3.sh new file mode 100644 index 0000000000..d6be68212d --- /dev/null +++ b/resubmit_failed_v3.sh @@ -0,0 +1,89 @@ +#!/bin/bash +# ============================================================ +# Resubmit only truly failed/incomplete jobs with batch_size=32 +# Usage: bash resubmit_failed_v3.sh +# ============================================================ + +USER_NAME=${1:-sgnanaku} +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/$USER_NAME +LOGS=$SCRATCH/slurm_logs_v2 +mkdir -p $LOGS + +echo "Resubmitting failed jobs for user: $USER_NAME" +echo "" + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + total_train_epochs=10 \\\\ + train_dataset.batch_size=32 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + evaluator.freq_steps=20 \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# CISPO - incomplete/OOM +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# SAPO - missing/incomplete +submit_job "sapo_eta0" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "4" + +# M2PO - wrong config flag before, fixed now +submit_job "m2po_eta0" "++actor.m2_threshold=0.04" "0" +submit_job "m2po_eta2" "++actor.m2_threshold=0.04" "2" +submit_job "m2po_eta4" "++actor.m2_threshold=0.04" "4" + +echo "" +echo "9 jobs resubmitted with batch_size=32!" +echo "Monitor with: squeue -u ${USER_NAME}" +echo "Logs at: ${LOGS}" diff --git a/resubmit_sapo.sh b/resubmit_sapo.sh new file mode 100644 index 0000000000..1b7e128426 --- /dev/null +++ b/resubmit_sapo.sh @@ -0,0 +1,62 @@ +#!/bin/bash +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/sgnanaku +LOGS=$SCRATCH/slurm_logs + +submit_job() { + local EXP_NAME=$1 + local ETA=$2 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + train_dataset.batch_size=64 \\\\ + total_train_epochs=1 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + +actor.use_sapo_loss=true \\\\ + ++actor.use_decoupled_loss=false + " +SBATCH + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +submit_job "sapo_eta0" "0" +submit_job "sapo_eta2" "2" +submit_job "sapo_eta4" "4" diff --git a/rollout.backend=sglang:d1p1t1 b/rollout.backend=sglang:d1p1t1 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/rollout.max_head_offpolicyness=2 b/rollout.max_head_offpolicyness=2 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scheduler.type=local b/scheduler.type=local new file mode 100644 index 0000000000..e69de29bb2 diff --git a/submit_all.sh b/submit_all.sh new file mode 100755 index 0000000000..fbc5bda348 --- /dev/null +++ b/submit_all.sh @@ -0,0 +1,91 @@ +#!/bin/bash +# ============================================================ +# Submit all 9 experiments to Slurm +# Run this from the login node: bash submit_all.sh +# ============================================================ + +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/sgnanaku +AREAL=$SCRATCH/AReaL +LOGS=$SCRATCH/slurm_logs +mkdir -p $LOGS + +# Function to submit one job +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + train_dataset.batch_size=64 \\\\ + total_train_epochs=10 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# GRPO (baseline) +submit_job "grpo_eta0" "" "0" +submit_job "grpo_eta2" "" "2" +submit_job "grpo_eta4" "" "4" + +# CISPO +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# SAPO +submit_job "sapo_eta0" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true +actor.use_decoupled_loss=false" "4" + +# M2PO +submit_job "m2po_eta0" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "0" +submit_job "m2po_eta2" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "2" +submit_job "m2po_eta4" "+actor.use_m2po_loss=true ++actor.m2_threshold=0.04" "4" + +echo "" +echo "All 9 jobs submitted! Monitor with: squeue -u sgnanaku" +echo "Logs at: $LOGS" diff --git a/submit_all_final.sh b/submit_all_final.sh new file mode 100644 index 0000000000..1c411a9b64 --- /dev/null +++ b/submit_all_final.sh @@ -0,0 +1,99 @@ +#!/bin/bash +# ============================================================ +# AReaL Experiment Matrix: GRPO vs CISPO vs SAPO x eta={0,2,4} +# Usage: bash submit_all.sh +# Example: bash submit_all.sh sgnanaku +# ============================================================ + +USER_NAME=${1:-sgnanaku} +SCRATCH=/scratch/zt1/project/zaoxing-prj/user/$USER_NAME +LOGS=$SCRATCH/slurm_logs +mkdir -p $LOGS + +MODEL_PATH=$SCRATCH/.cache/huggingface/hub/models--Qwen--Qwen2.5-1.5B-Instruct/snapshots/989aa7980e4cf806f80c7fef2b1adb7bc71aa306 + +echo "Submitting jobs for user: $USER_NAME" +echo "Scratch: $SCRATCH" +echo "Logs: $LOGS" +echo "" + +submit_job() { + local EXP_NAME=$1 + local ALGO_FLAGS=$2 + local ETA=$3 + + sbatch < /tmp/te_stub.py + export LD_LIBRARY_PATH=/.singularity.d/libs:\\\$LD_LIBRARY_PATH + source /AReaL/.venv/bin/activate + cd \$SCRATCH/AReaL + uv pip install -e . --no-deps -q + + python3 examples/math/gsm8k_rl.py \\\\ + --config examples/math/gsm8k_grpo.yaml \\\\ + scheduler.type=local \\\\ + experiment_name=${EXP_NAME} \\\\ + trial_name=run1 \\\\ + rollout.backend=sglang:d1p1t1 \\\\ + actor.backend=fsdp:d1p1t1 \\\\ + cluster.n_nodes=1 \\\\ + cluster.n_gpus_per_node=2 \\\\ + actor.path=\$MODEL_PATH \\\\ + gconfig.max_new_tokens=512 \\\\ + train_dataset.batch_size=64 \\\\ + total_train_epochs=1 \\\\ + rollout.max_head_offpolicyness=${ETA} \\\\ + ${ALGO_FLAGS} + " +EOF + echo "Submitted: ${EXP_NAME} (eta=${ETA})" +} + +# ============================================================ +# GRPO (baseline) +# ============================================================ +submit_job "grpo_eta0" "" "0" +submit_job "grpo_eta2" "" "2" +submit_job "grpo_eta4" "" "4" + +# ============================================================ +# CISPO +# ============================================================ +submit_job "cispo_eta0" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "0" +submit_job "cispo_eta2" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "2" +submit_job "cispo_eta4" "+actor.use_cispo_loss=true +actor.cispo_epsilon_high=1.2" "4" + +# ============================================================ +# SAPO +# ============================================================ +submit_job "sapo_eta0" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "0" +submit_job "sapo_eta2" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "2" +submit_job "sapo_eta4" "+actor.use_sapo_loss=true ++actor.use_decoupled_loss=false" "4" + +echo "" +echo "All 9 jobs submitted!" +echo "Monitor with: squeue -u ${USER_NAME}" +echo "Logs at: ${LOGS}" diff --git a/total_train_epochs=1 b/total_train_epochs=1 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/train_dataset.batch_size=64 b/train_dataset.batch_size=64 new file mode 100644 index 0000000000..e69de29bb2 diff --git a/trial_name=run1 b/trial_name=run1 new file mode 100644 index 0000000000..e69de29bb2