# Tested successfully on the hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.4-flashinfer0.2.2-cxx11abi0 image.
# Pairwise comparison training for HelpSteer2 dataset

TIMESTAMP=$(date "+%m%dT%H%M")

TRAIN_FILE=./examples/data/exports/helpsteer2_pairwise_train.parquet
VAL_FILE=./examples/data/exports/helpsteer2_pairwise_test.parquet
MODEL_PATH=/mnt/data_cpfs/xielipeng.xlp/models/Qwen2.5-14B-Instruct

PROJECT_NAME=pairwise_train
EXPERIMENT_NAME=rm-gallery-pairwise-qwen2.5-14b-${TIMESTAMP}

CUSTOM_REWARD_FUNCTION_PATH=./examples/train/pairwise/reward_fn.py
CUSTOM_CHAT_RL_DATASET_PATH=./examples/train/pairwise/dataset.py
CUSTOM_CHAT_RL_DATASET_NAME=HelpfulnessPairwiseTrainDataset
REWARD_MANAGER=naive
REWARD_FUNCTION_NAME=compute_score

DEFAULT_LOCAL_DIR=./checkpoints/${TIMESTAMP}

N_GPUS_PER_NODE=8
N_NODES=1

set -x

ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env=./examples/train/pairwise/runtime_env.yaml \
    -- \
    python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files="$TRAIN_FILE" \
    data.val_files="$VAL_FILE" \
    data.train_batch_size=96 \
    data.val_batch_size=192 \
    data.max_prompt_length=4096 \
    data.max_response_length=2048 \
    data.filter_overlong_prompts=True \
    data.truncation='right' \
    data.prompt_key='input' \
    data.custom_cls.path="${CUSTOM_CHAT_RL_DATASET_PATH}" \
    data.custom_cls.name="${CUSTOM_CHAT_RL_DATASET_NAME}" \
    reward_model.reward_manager=${REWARD_MANAGER} \
    custom_reward_function.path=${CUSTOM_REWARD_FUNCTION_PATH} \
    custom_reward_function.name=${REWARD_FUNCTION_NAME} \
    actor_rollout_ref.model.path=${MODEL_PATH} \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=24 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
    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.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=1 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
    actor_rollout_ref.rollout.n=4 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
    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=${PROJECT_NAME} \
    trainer.experiment_name=${EXPERIMENT_NAME} \
    trainer.n_gpus_per_node=${N_GPUS_PER_NODE} \
    trainer.nnodes=${N_NODES} \
    trainer.save_freq=20 \
    trainer.test_freq=2 \
    trainer.total_epochs=10 \
    trainer.val_before_train=False \
    trainer.default_local_dir=${DEFAULT_LOCAL_DIR}