-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathtrain_7b_grpo_megatron.sh
More file actions
145 lines (139 loc) · 6.58 KB
/
Copy pathtrain_7b_grpo_megatron.sh
File metadata and controls
145 lines (139 loc) · 6.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
set -x
dataset_name=deepmath_torl # or math_torl_offical to use torl training data
train_data=$(pwd)/data/${dataset_name}/train.parquet
val_data=[$(pwd)/data/${dataset_name}/test.parquet,\
$(pwd)/data/${dataset_name}/math500_test.parquet,\
$(pwd)/data/${dataset_name}/aime24_test.parquet,\
$(pwd)/data/${dataset_name}/aime25_test.parquet]
model_name=Qwen/Qwen2.5-Math-7B
rl_alg=grpo # gae(ppo) or grpo, if grpo, then better set n>1 otherwise the group norm can not be effective
n_gpus_per_node=8
n_nodes=1
n=16
batch_size=128
ppo_mini_batch_size=128
max_prompt_length=1024
max_response_length=3072
max_obs_length=512
temperature=1.0
top_p=1.0
enable_agent=True # enable agent for tool use
strategy="megatron"
action_stop_tokens='```output'
max_turns=1
kl_loss_coef=0.0
kl_coef=0
entropy_coeff=0
kl_loss_type=low_var_kl
lr=1e-6
reward_manager=torl
ppo_micro_batch_size_per_gpu=1
log_prob_micro_batch_size_per_gpu=8
tensor_model_parallel_size=1
megatron_TP=2
megatron_PP=2
gpu_memory_utilization=0.7 # higher gpu_memory_utilization will likely cause the vllm to OOM and get stuck, so set it to a lower value like 0.4 or 0.5
do_offload=False
use_dynamic_bsz=False # faster
mask_observations=True # mask observations for kl loss and gradient descent
enable_mtrl=False # enable multi-turn training
max_action_length=2048
model_pretty_name=$(echo $model_name | tr '/' '_' | tr '[:upper:]' '[:lower:]')
run_name_postfix="acc-only-max-200-steps"
if [ "$enable_agent" = "True" ]; then
run_name="${reward_manager}-${strategy}-agent-${model_pretty_name}-${rl_alg}-n${n}-b${batch_size}-t${temperature}-lr${lr}${run_name_postfix}"
else
run_name="${reward_manager}-${strategy}-${model_pretty_name}-${rl_alg}-n${n}-b${batch_size}-t${temperature}-lr${lr}${run_name_postfix}"
fi
export VERL_RUN_ID=$run_name
export NCCL_DEBUG=INFO
export VLLM_USE_V1=1
rollout_mode='async' # for megatron, only async mode is supported currently
# temp file for action tokens as verl cannot pass special strs as params
action_stop_tokens_file="$(pwd)$(mktemp)"
mkdir -p $(dirname $action_stop_tokens_file)
echo -e -n "$action_stop_tokens" | tee $action_stop_tokens_file
echo "action_stop_tokens_file=$action_stop_tokens_file"
host=$(hostname -i | awk '{print $1}')
port=$(shuf -i 30000-31000 -n 1)
tool_server_url=http://$host:$port/get_observation
python -m verl_tool.servers.serve --host $host --port $port --tool_type "python_code" --workers_per_tool 8 &
server_pid=$!
echo "Server (pid=$server_pid) started at $tool_server_url"
PYTHONUNBUFFERED=1 python3 -m verl_tool.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer' \
algorithm.adv_estimator=$rl_alg \
data.train_files=$train_data \
data.val_files=$val_data \
data.train_batch_size=$batch_size \
data.val_batch_size=1024 \
data.max_prompt_length=$max_prompt_length \
data.max_response_length=$max_response_length \
data.truncation='right' \
reward_model.reward_manager=$reward_manager \
reward_model.launch_reward_fn_async=True \
actor_rollout_ref.model.path=$model_name \
actor_rollout_ref.model.enable_gradient_checkpointing=True \
actor_rollout_ref.actor.optim.lr=$lr \
actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
actor_rollout_ref.model.trust_remote_code=True \
actor_rollout_ref.actor.checkpoint.save_contents=['model','optimizer','extra','hf_model'] \
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.use_dynamic_bsz=$use_dynamic_bsz \
actor_rollout_ref.actor.use_kl_loss=True \
actor_rollout_ref.actor.strategy=$strategy \
actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \
actor_rollout_ref.actor.kl_loss_type=$kl_loss_type \
actor_rollout_ref.actor.entropy_coeff=$entropy_coeff \
actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$megatron_PP \
actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$megatron_TP \
actor_rollout_ref.actor.megatron.param_offload=$do_offload \
actor_rollout_ref.actor.megatron.grad_offload=$do_offload \
actor_rollout_ref.actor.megatron.optimizer_offload=$do_offload \
actor_rollout_ref.agent.enable_agent=$enable_agent \
actor_rollout_ref.agent.tool_server_url=$tool_server_url \
actor_rollout_ref.agent.max_prompt_length=$max_prompt_length \
actor_rollout_ref.agent.max_response_length=$max_response_length \
actor_rollout_ref.agent.max_start_length=$max_prompt_length \
actor_rollout_ref.agent.max_obs_length=$max_obs_length \
actor_rollout_ref.agent.max_turns=$max_turns \
actor_rollout_ref.agent.mask_observations=$mask_observations \
actor_rollout_ref.agent.action_stop_tokens=$action_stop_tokens_file \
actor_rollout_ref.agent.enable_mtrl=$enable_mtrl \
actor_rollout_ref.agent.max_action_length=$max_action_length \
actor_rollout_ref.rollout.tensor_model_parallel_size=$tensor_model_parallel_size \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$log_prob_micro_batch_size_per_gpu \
actor_rollout_ref.rollout.enforce_eager=False \
actor_rollout_ref.rollout.free_cache_engine=True \
actor_rollout_ref.rollout.name=vllm \
actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \
actor_rollout_ref.rollout.temperature=$temperature \
actor_rollout_ref.rollout.top_p=$top_p \
actor_rollout_ref.rollout.top_k=-1 \
actor_rollout_ref.rollout.n=$n \
actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=$use_dynamic_bsz \
actor_rollout_ref.rollout.max_num_seqs=512 \
actor_rollout_ref.rollout.mode=$rollout_mode \
actor_rollout_ref.ref.log_prob_use_dynamic_bsz=$use_dynamic_bsz \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=$log_prob_micro_batch_size_per_gpu \
actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$megatron_PP \
actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$megatron_TP \
critic.optim.lr=1e-5 \
critic.strategy=$strategy \
critic.model.path=$model_name \
critic.ppo_micro_batch_size_per_gpu=$ppo_micro_batch_size_per_gpu \
algorithm.kl_ctrl.kl_coef=$kl_coef \
trainer.logger=['console','wandb'] \
trainer.project_name=$reward_manager \
trainer.experiment_name=$run_name \
trainer.val_before_train=True \
trainer.default_hdfs_dir=null \
trainer.n_gpus_per_node=$n_gpus_per_node \
trainer.nnodes=$n_nodes \
+trainer.remove_previous_ckpt_in_save=True \
trainer.save_freq=10 \
trainer.test_freq=10 \
trainer.total_epochs=10 \
trainer.total_training_steps=200
pkill -P -9 $server_pid
kill -9 $kill $server_pid