set -x
nnodes=
current_rank=
master_addr=
master_port=


policy_path=OctoThinker/OctoThinker-3B-Hybrid-Base
rollout_batch_size=256
n_samples_per_prompts=16
episode=100
temperature=1.0
batch_size=16
lr=1e-6
kl_loss_coef=0.0
kl_coef=0
entropy_coeff=0
max_gen_length=8192
dataset_name=math.8k
run_name=grpo_octothinker_3b_hybrid_${dataset_name}_$(TZ='UTC-8' date "+%m%d_%H%M")
echo $run_name


data_path=./data/${dataset_name}
samples_save_path=path/to/save/samples/$run_name

if [ "$current_rank" = "0" ]; then
    ray start --head --port=6379 --num-gpus 8
elif [ "$current_rank" != "0" ]; then
    ray start --address=${master_addr}:6379 --num-gpus 8 --block
fi

if [ "$current_rank" = "0" ]; then
python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$data_path/train.parquet \
    data.val_files=$data_path/test.parquet \
    data.train_batch_size=$rollout_batch_size \
    data.val_batch_size=2048 \
    data.shuffle=True \
    data.max_prompt_length=400 \
    data.template_type=orz_base \
    data.filter_overlong_prompts=True \
    data.max_response_length=$max_gen_length \
    actor_rollout_ref.model.path=$policy_path \
    actor_rollout_ref.actor.optim.lr=$lr \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=$batch_size \
    actor_rollout_ref.actor.use_dynamic_bsz=True \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \
    actor_rollout_ref.actor.entropy_coeff=$entropy_coeff \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.scheduler=stair \
    actor_rollout_ref.rollout.length_scheduler_dict="{0: 4096}" \
    actor_rollout_ref.rollout.temperature=$temperature \
    actor_rollout_ref.rollout.top_p=1 \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
    actor_rollout_ref.rollout.n=$n_samples_per_prompts \
    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \
    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    algorithm.kl_ctrl.kl_coef=$kl_coef \
    trainer.logger=['console','wandb'] \
    trainer.project_name='ct_rl' \
    trainer.experiment_name=${run_name} \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=$nnodes \
    trainer.save_freq=40 \
    trainer.test_freq=4 \
    trainer.default_local_dir=path/to/save/ckpts/$run_name \
    trainer.resume_mode=auto \
    trainer.samples_save_path=$samples_save_path \
    trainer.total_epochs=$episode $@