Adapt to own scripts and config for A100.

.gitignore

@@ -111,3 +111,9 @@ tests/e2e/toy_examples/deepspeed/synchronous/output.txt
 
 # vim
 *.swp
+
+# large data
+*.parquet
+
+# model
+model/

scripts/train_tiny_zero_a100.sh

@@ -0,0 +1,29 @@
+/home/weiji/anaconda3/envs/zero/bin/python3 -m verl.trainer.main_ppo \
+data.train_files=$DATA_DIR/train.parquet \
+data.val_files=$DATA_DIR/test.parquet \
+data.train_batch_size=128 \
+data.val_batch_size=640 \
+data.max_prompt_length=256 \
+data.max_response_length=1024 \
+actor_rollout_ref.model.path=$BASE_MODEL \
+actor_rollout_ref.actor.optim.lr=1e-6 \
+actor_rollout_ref.actor.ppo_mini_batch_size=64 \
+actor_rollout_ref.actor.ppo_micro_batch_size=4 \
+actor_rollout_ref.rollout.log_prob_micro_batch_size=4 \
+actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP_SIZE \
+actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
+actor_rollout_ref.ref.log_prob_micro_batch_size=2 \
+critic.optim.lr=1e-5 \
+critic.model.path=$BASE_MODEL \
+critic.ppo_micro_batch_size=4 \
+algorithm.kl_ctrl.kl_coef=0.001 \
+trainer.logger=['wandb'] \
++trainer.val_before_train=False \
+trainer.default_hdfs_dir=null \
+trainer.n_gpus_per_node=$N_GPUS \
+trainer.nnodes=1 \
+trainer.save_freq=100 \
+trainer.test_freq=100 \
+trainer.project_name=TinyZero \
+trainer.experiment_name=$EXPERIMENT_NAME \
+trainer.total_epochs=15 2>&1 | tee verl_demo.log

scripts/train_tiny_zero.sh → scripts/train_tiny_zero_h200.sh

@@ -1,4 +1,4 @@
-python3 -m verl.trainer.main_ppo \
+/home/weiji/anaconda3/envs/zero/bin/python3 -m verl.trainer.main_ppo \
 data.train_files=$DATA_DIR/train.parquet \
 data.val_files=$DATA_DIR/test.parquet \
 data.train_batch_size=256 \

train_0.5b_ppo.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# alias python='/home/weiji/anaconda3/envs/zero/bin/python'
+# alias python3='/home/weiji/anaconda3/envs/zero/bin/python3'
+# alias pip='/home/weiji/anaconda3/envs/zero/bin/pip'
+
+export N_GPUS=1
+export CUDA_VISIBLE_DEVICES=5
+ray stop --force && ray start --head --include-dashboard=True
+export BASE_MODEL="model/Qwen2.5-0.5B"
+export DATA_DIR="data/countdown"
+export ROLLOUT_TP_SIZE=1
+export EXPERIMENT_NAME=countdown-qwen2.5-0.5b
+export VLLM_ATTENTION_BACKEND=XFORMERS
+
+bash ./scripts/train_tiny_zero.sh

train_3b_ppo.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# alias python='/home/weiji/anaconda3/envs/zero/bin/python'
+# alias python3='/home/weiji/anaconda3/envs/zero/bin/python3'
+# alias pip='/home/weiji/anaconda3/envs/zero/bin/pip'
+
+export N_GPUS=2
+export CUDA_VISIBLE_DEVICES=2,3
+ray stop --force && ray start --head --include-dashboard=True
+export BASE_MODEL="model/Qwen2.5-3B"
+export DATA_DIR="data/countdown"
+export ROLLOUT_TP_SIZE=2
+export EXPERIMENT_NAME=countdown-qwen2.5-3b
+export VLLM_ATTENTION_BACKEND=XFORMERS
+
+bash ./scripts/train_tiny_zero_a100.sh