[algo] feat: add GSPO-token policy loss computation function (volcengine#2775)

### What does this PR do?

This PR implements the GSPO-token policy loss calculation proposed by
paper https://arxiv.org/pdf/2507.18071

### Test

<img width="1341" height="637" alt="image"
src="https://github.com/user-attachments/assets/bc5e2245-b0f5-4a1f-aa7c-4c2b28d95142"
/>

Compared GRPO and GSPO under the same settings. GRPO uses the following
script:
``` sh
python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.train_batch_size=512 \
    data.max_prompt_length=512 \
    data.max_response_length=1024 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=128 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.actor.policy_loss.loss_mode="vanilla" \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
    actor_rollout_ref.rollout.n=10 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    algorithm.use_kl_in_reward=False \
    trainer.critic_warmup=0 \
    trainer.logger='["console","wandb"]' \
    trainer.project_name='verl_gspo_cmp' \
    trainer.experiment_name='qwen2.5-3B-GRPO' \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=1 \
    trainer.save_freq=20 \
    trainer.test_freq=5 \
    trainer.total_epochs=15 $@
```

GSPO uses the following script:
```sh
python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.train_batch_size=512 \
    data.max_prompt_length=512 \
    data.max_response_length=1024 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=128 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.actor.policy_loss.loss_mode="gspo" \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
    actor_rollout_ref.rollout.n=10 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    algorithm.use_kl_in_reward=False \
    trainer.critic_warmup=0 \
    trainer.logger='["console","wandb"]' \
    trainer.project_name='verl_gspo_cmp' \
    trainer.experiment_name='qwen2.5-3B-GRPO' \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=1 \
    trainer.save_freq=20 \
    trainer.test_freq=5 \
    trainer.total_epochs=15 $@

```

### API and Usage Example

To use GSPO, users only need to set
`actor_rollout_ref.actor.policy_loss.loss_mode` to `gspo`.

```shell
python3 -m verl.trainer.main_ppo \
  ... \
  actor_rollout_ref.actor.policy_loss.loss_mode="gspo" \
  ...
```


### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [x] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [x] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [x] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

---------

Co-authored-by: BounharAbdelaziz <[email protected]>

Author: 2 people

recipe/gspo/test_gspo_3b_math.sh

@@ -0,0 +1,191 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=rl-gspo-3B
+#SBATCH --partition=main
+#SBATCH --nodes=1                # Number of nodes
+#SBATCH --ntasks-per-node=1      # One task per node
+#SBATCH --cpus-per-task=128      # cpu-cores per task
+#SBATCH --gres=gpu:8
+#SBATCH --mem=0
+#SBATCH --exclusive
+#SBATCH --time=500:00:00
+#SBATCH --output=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.out
+#SBATCH --error=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.err
+
+set -xeuo pipefail
+
+# activate the venv
+echo "Activating verl environment..."
+eval "$(conda shell.bash hook)"
+conda deactivate
+conda activate verl
+
+# can make training faster, depends on your infrastructure
+export NCCL_IBEXT_DISABLE=1
+export NCCL_NVLS_ENABLE=1
+export NCCL_IB_HCA=mlx5
+export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1
+
+# Set how many GPUs we actually have on this node.
+export GPUS_PER_NODE=8
+
+NNODES=${SLURM_JOB_NUM_NODES}
+export NNODES
+
+export VLLM_ATTENTION_BACKEND=FLASH_ATTN
+export RAY_LOGGING_LEVEL=DEBUG
+export HYDRA_FULL_ERROR=1
+export WANDB_API_KEY=... # your wandb API key
+
+echo "Using $NNODES nodes for training..."
+
+# ------------------------------------- Setup xp params ---------------------------------------
+project_name='RL-GSPO'
+
+adv_estimator=grpo
+loss_mode=gspo
+loss_agg_mode="seq-mean-token-mean"
+MODEL_PATH=Qwen/Qwen2.5-3B-Instruct
+offload=false # it's a small model, offloading will just slow-down training
+rollout_engine=vllm
+rollout_mode=sync # can be async to speedup large scale xps
+gpu_memory_utilization=0.8
+reward_manager=dapo
+adv_estimator=grpo
+shuffle_dataset=true
+first_time_dataset_prep=true # prepare dataset
+
+test_freq=10
+save_freq=10
+total_epochs=10
+total_training_steps=500
+val_before_train=false
+
+use_kl_in_reward=false
+kl_coef=0.0
+use_kl_loss=false
+kl_loss_coef=0.0
+
+clip_ratio_low=0.0003 # as recommended by the paper, see Sec. 5.1
+clip_ratio_high=0.0004 # as recommended by the paper, see Sec. 5.1
+train_batch_size=512
+ppo_mini_batch_size=128 # maintain 4 mini-batches as recommended by the paper, see Sec. 5.1
+ppo_micro_batch_size_per_gpu=8 # setup depending on your GPU memory
+n_resp_per_prompt=16
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+# dapo reward manager params
+enable_overlong_buffer=false # true
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Paths and namings
+SFT_MODEL=$(basename $MODEL_PATH)
+exp_name="${loss_mode}-epslow-${clip_ratio_low}-epshigh-${clip_ratio_high}-${SFT_MODEL}-RL"
+CKPTS_DIR=/rl/checkpoints/experimental/4b/${loss_mode}/${exp_name}
+
+# Sampling params at rollouts
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+sp_size=1
+use_dynamic_bsz=true
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+offload=true
+gen_tp=1
+entropy_checkpointing=true # This enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training.
+
+# ------------------------------------- train/val data preparation ---------------------------------------
+if [ "$first_time_dataset_prep" = true ]; then
+    echo "Preprocessing GSM8K dataset..."
+    python examples/data_preprocess/gsm8k.py --local_dir /data/gsm8k/
+fi
+
+gsm8k_train_path=/data/gsm8k/train.parquet
+gsm8k_test_path=/data/gsm8k/test.parquet
+
+# set the paths
+train_files="['$gsm8k_train_path']"
+test_files="['$gsm8k_test_path']"
+
+python3 -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=${adv_estimator} \
+    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \
+    data.train_files="${train_files}" \
+    data.val_files="${test_files}" \
+    data.shuffle=$shuffle_dataset \
+    data.prompt_key=prompt \
+    data.truncation='error' \
+    data.filter_overlong_prompts=true \
+    data.train_batch_size=${train_batch_size} \
+    data.max_prompt_length=${max_prompt_length} \
+    data.max_response_length=${max_response_length} \
+    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+    algorithm.use_kl_in_reward=${use_kl_in_reward} \
+    algorithm.kl_ctrl.kl_coef=${kl_coef} \
+    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+    actor_rollout_ref.model.use_remove_padding=true \
+    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.name=${rollout_engine} \
+    actor_rollout_ref.rollout.mode=${rollout_mode} \
+    actor_rollout_ref.model.path="${MODEL_PATH}" \
+    actor_rollout_ref.model.enable_gradient_checkpointing=true \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.05 \
+    actor_rollout_ref.actor.optim.weight_decay=0.1 \
+    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ppo_micro_batch_size_per_gpu} \
+    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.actor.grad_clip=1.0 \
+    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+    actor_rollout_ref.rollout.enable_chunked_prefill=true \
+    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+    actor_rollout_ref.rollout.temperature=${temperature} \
+    actor_rollout_ref.rollout.top_p=${top_p} \
+    actor_rollout_ref.rollout.top_k=${top_k} \
+    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+    actor_rollout_ref.rollout.val_kwargs.do_sample=true \
+    actor_rollout_ref.rollout.val_kwargs.n=1 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \
+    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+    actor_rollout_ref.actor.entropy_checkpointing=${entropy_checkpointing} \
+    reward_model.reward_manager=${reward_manager} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.log=false \
+    +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+    trainer.logger='["console","wandb"]' \
+    trainer.project_name="${project_name}" \
+    trainer.experiment_name="${exp_name}" \
+    trainer.n_gpus_per_node="${GPUS_PER_NODE}" \
+    trainer.nnodes="${NNODES}" \
+    trainer.val_before_train=${val_before_train} \
+    trainer.test_freq=${test_freq} \
+    trainer.save_freq=${save_freq} \
+    trainer.total_epochs=${total_epochs} \
+    trainer.total_training_steps=${total_training_steps} \
+    trainer.default_local_dir="${CKPTS_DIR}" \
+    trainer.resume_mode=auto \
+    trainer.log_val_generations=2 \
+    $@

recipe/gspo/test_gspo_3b_math_slurm.sh

@@ -0,0 +1,195 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=rl-gspo-3B
+#SBATCH --partition=main
+#SBATCH --nodes=1                # Number of nodes
+#SBATCH --ntasks-per-node=1      # One task per node
+#SBATCH --cpus-per-task=128      # cpu-cores per task
+#SBATCH --gres=gpu:8
+#SBATCH --mem=0
+#SBATCH --exclusive
+#SBATCH --time=500:00:00
+#SBATCH --output=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.out
+#SBATCH --error=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.err
+
+set -xeuo pipefail
+
+# activate the venv
+echo "Activating verl environment..."
+eval "$(conda shell.bash hook)"
+conda deactivate
+conda activate verl
+
+# can make training faster, depends on your infrastructure
+export NCCL_IBEXT_DISABLE=1
+export NCCL_NVLS_ENABLE=1
+export NCCL_IB_HCA=mlx5
+export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1
+
+# Set how many GPUs we actually have on this node.
+export GPUS_PER_NODE=8
+
+NNODES=${SLURM_JOB_NUM_NODES}
+export NNODES
+
+export VLLM_ATTENTION_BACKEND=FLASH_ATTN
+export RAY_memory_monitor_refresh_ms=0
+export RAY_LOGGING_LEVEL=DEBUG
+export HYDRA_FULL_ERROR=1
+export WANDB_API_KEY=... # your wandb API key
+
+# Let Ray know how many nodes to expect
+export RAY_NUM_NODES=$NNODES
+
+echo "Using $NNODES nodes for training..."
+
+# ------------------------------------- Setup xp params ---------------------------------------
+project_name='RL-GSPO'
+
+adv_estimator=grpo
+loss_mode=gspo
+loss_agg_mode="seq-mean-token-mean"
+MODEL_PATH=Qwen/Qwen2.5-3B-Instruct
+offload=false # it's a small model, offloading will just slow-down training
+rollout_engine=vllm
+rollout_mode=sync # can be async to speedup large scale xps
+gpu_memory_utilization=0.8
+reward_manager=dapo
+adv_estimator=grpo
+shuffle_dataset=true
+first_time_dataset_prep=true # prepare dataset
+
+test_freq=10
+save_freq=10
+total_epochs=10
+total_training_steps=500
+val_before_train=false
+
+use_kl_in_reward=false
+kl_coef=0.0
+use_kl_loss=false
+kl_loss_coef=0.0
+
+clip_ratio_low=0.0003 # as recommended by the paper, see Sec. 5.1
+clip_ratio_high=0.0004 # as recommended by the paper, see Sec. 5.1
+train_batch_size=512
+ppo_mini_batch_size=128 # maintain 4 mini-batches as recommended by the paper, see Sec. 5.1
+ppo_micro_batch_size_per_gpu=8 # setup depending on your GPU memory
+n_resp_per_prompt=16
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+# dapo reward manager params
+enable_overlong_buffer=false # true
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Paths and namings
+SFT_MODEL=$(basename $MODEL_PATH)
+exp_name="${loss_mode}-epslow-${clip_ratio_low}-epshigh-${clip_ratio_high}-${SFT_MODEL}-RL"
+CKPTS_DIR=/rl/checkpoints/experimental/4b/${loss_mode}/${exp_name}
+
+# Sampling params at rollouts
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+sp_size=1
+use_dynamic_bsz=true
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+offload=true
+gen_tp=1
+entropy_checkpointing=true # This enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training.
+
+# ------------------------------------- train/val data preparation ---------------------------------------
+if [ "$first_time_dataset_prep" = true ]; then
+    echo "Preprocessing GSM8K dataset..."
+    python examples/data_preprocess/gsm8k.py --local_dir /data/gsm8k/
+fi
+
+gsm8k_train_path=/data/gsm8k/train.parquet
+gsm8k_test_path=/data/gsm8k/test.parquet
+
+# set the paths
+train_files="['$gsm8k_train_path']"
+test_files="['$gsm8k_test_path']"
+
+python3 -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=${adv_estimator} \
+    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \
+    data.train_files="${train_files}" \
+    data.val_files="${test_files}" \
+    data.shuffle=$shuffle_dataset \
+    data.prompt_key=prompt \
+    data.truncation='error' \
+    data.filter_overlong_prompts=true \
+    data.train_batch_size=${train_batch_size} \
+    data.max_prompt_length=${max_prompt_length} \
+    data.max_response_length=${max_response_length} \
+    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+    algorithm.use_kl_in_reward=${use_kl_in_reward} \
+    algorithm.kl_ctrl.kl_coef=${kl_coef} \
+    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+    actor_rollout_ref.model.use_remove_padding=true \
+    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.name=${rollout_engine} \
+    actor_rollout_ref.rollout.mode=${rollout_mode} \
+    actor_rollout_ref.model.path="${MODEL_PATH}" \
+    actor_rollout_ref.model.enable_gradient_checkpointing=true \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.05 \
+    actor_rollout_ref.actor.optim.weight_decay=0.1 \
+    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ppo_micro_batch_size_per_gpu} \
+    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.actor.grad_clip=1.0 \
+    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+    actor_rollout_ref.rollout.enable_chunked_prefill=true \
+    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+    actor_rollout_ref.rollout.temperature=${temperature} \
+    actor_rollout_ref.rollout.top_p=${top_p} \
+    actor_rollout_ref.rollout.top_k=${top_k} \
+    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+    actor_rollout_ref.rollout.val_kwargs.do_sample=true \
+    actor_rollout_ref.rollout.val_kwargs.n=1 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \
+    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+    actor_rollout_ref.actor.entropy_checkpointing=${entropy_checkpointing} \
+    reward_model.reward_manager=${reward_manager} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+    +reward_model.reward_kwargs.overlong_buffer_cfg.log=false \
+    +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+    trainer.logger='["console","wandb"]' \
+    trainer.project_name="${project_name}" \
+    trainer.experiment_name="${exp_name}" \
+    trainer.n_gpus_per_node="${GPUS_PER_NODE}" \
+    trainer.nnodes="${NNODES}" \
+    trainer.val_before_train=${val_before_train} \
+    trainer.test_freq=${test_freq} \
+    trainer.save_freq=${save_freq} \
+    trainer.total_epochs=${total_epochs} \
+    trainer.total_training_steps=${total_training_steps} \
+    trainer.default_local_dir="${CKPTS_DIR}" \
+    trainer.resume_mode=auto \
+    trainer.log_val_generations=2 \
+    $@