set -x

conda activate openrlhf


MODEL_SIZE=7B
MODEL_PATH=Qwen2.5-Math-${MODEL_SIZE}
KL_TYPE=k1

timestamp=$(date +"%Y%m%d-%H%M")
logfile="/${KL_COEF}-${MODEL_SIZE}-${KL_TYPE}-${timestamp}.log"

python3 -m openrlhf.cli.train_ppo_ray \
   --ref_num_nodes 1 \
   --ref_num_gpus_per_node 8 \
   --actor_num_nodes 1 \
   --actor_num_gpus_per_node 8 \
   --vllm_num_engines 8 \
   --vllm_tensor_parallel_size 1 \
   --colocate_all_models \
   --vllm_gpu_memory_utilization 0.6 \
   --init_kl_coef ${KL_COEF} \
   --entropy_loss_coef 0.0 \
   --gamma 1.0 \
   --use_kl_loss \
   --kl_estimator ${KL_TYPE} \
   --advantage_estimator group_norm \
   --pretrain ${MODEL_PATH} \
   --remote_rm_url /script/mathverify.py \
   --save_path /${KL_COEF}-${MODEL_SIZE}-${KL_TYPE}-${timestamp} \
   --ckpt_path /${KL_COEF}-${MODEL_SIZE}-${KL_TYPE}-${timestamp} \
   --save_steps 9999 \
   --save_hf_ckpt \
   --micro_train_batch_size 32 \
   --train_batch_size 256 \
   --micro_rollout_batch_size 32 \
   --rollout_batch_size 32 \
   --n_samples_per_prompt 8 \
   --max_epochs 1 \
   --prompt_max_len 1024 \
   --max_samples 999999999 \
   --generate_max_len 3072 \
   --zero_stage 3 \
   --bf16 \
   --actor_learning_rate 3e-6 \
   --prompt_data Openr1-Math-7k-4096.jsonl \
   --input_key question \
   --label_key label \
   --apply_chat_template \
   --gradient_checkpointing \
   --packing_samples \
   --vllm_sync_backend nccl \
   --enforce_eager \
   --vllm_enable_sleep \
   --deepspeed_enable_sleep \
   2>&1 | tee "$logfile" 2>&1