快速开始:在 GSM8K 数据集上进行 PPO 训练
使用 GSM8K 数据集对 LLM 进行后训练。
简介
在这个示例中,我们将 LLM 训练为处理 GSM8k 任务,并使用基于函数的奖励。 [1]
先决条件:
已按照安装指南安装最新版本的
verl及其依赖项。建议使用 docker 镜像。拥有至少 24 GB HBM 的 GPU
数据集简介
GSM8k 是一个数学问题数据集。提示是小学水平的应用题。要求 LLM 模型解决这些数学问题。示例如下:
提示
Katy 制作咖啡使用的糖的茶匙数和水的杯数之比为 7:13。如果她使用的糖的茶匙数和水的杯数总计为 120,请计算她使用的糖的茶匙数。
答案
她使用的配料的总比例为 7+13 = <<7+13=20>>20。由于她使用的茶匙数占总比例的 7/20,因此她使用了 7/20 * 120 = <<7/20*120=42>>42 #### 42
第一步:准备数据集
我们将数据集预处理为 parquet 格式,以便 (1) 包含计算 RL 奖励所需的字段,以及 (2) 加快读取速度。
python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k
第二步:下载模型用于后训练
在本例中,我们将从 Qwen2.5-0.5B-Instruct 模型开始。
如果您想在 RL 之前进行 SFT,请参阅 完整的 GSM8K 示例、sft 目录 和 SFT Trainer 以获取更多详细信息。
python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-0.5B-Instruct')"
第三步:使用指令模型执行 PPO 训练
奖励模型/函数
我们使用预定义的基于规则的奖励模型。我们强制模型按照答案中显示的 4 个“#”来生成最终答案。我们使用正则表达式匹配从解决方案和模型的输出来提取最终答案。我们将正确答案的奖励设置为 1,将错误答案的奖励设置为 0.0,将无答案的奖励设置为 0。
有关更多详细信息,请参阅 verl/utils/reward_score/gsm8k.py。
训练脚本
现在让我们使用上述数据集和模型运行 PPO 训练。[2]
根据您的数据集和模型名称或路径设置 data.train_files、data.val_files、actor_rollout_ref.model.path 和 critic.model.path。
您可以设置 VERL_USE_MODELSCOPE=True 以从 modelscope 而不是 huggingface 下载模型。
PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \
data.train_files=$HOME/data/gsm8k/train.parquet \
data.val_files=$HOME/data/gsm8k/test.parquet \
data.train_batch_size=256 \
data.max_prompt_length=512 \
data.max_response_length=256 \
actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
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_per_gpu=4 \
actor_rollout_ref.rollout.name=vllm \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
critic.optim.lr=1e-5 \
critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \
critic.ppo_micro_batch_size_per_gpu=4 \
algorithm.kl_ctrl.kl_coef=0.001 \
trainer.logger=console \
trainer.val_before_train=False \
trainer.n_gpus_per_node=1 \
trainer.nnodes=1 \
trainer.save_freq=10 \
trainer.test_freq=10 \
trainer.total_epochs=15 2>&1 | tee verl_demo.log
您应该会看到以下日志,表明训练正在进行中。关键指标 val/test_score/openai/gsm8k 每 trainer.test_freq 步计算一次:
step:0 - timing/gen:21.470 - timing/ref:4.360 - timing/values:5.800 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty_coeff:0.001 - timing/adv:0.109 - timing/update_critic:15.664 - critic/vf_loss:14.947 - critic/vf_clipfrac:0.000 - critic/vpred_mean:-2.056 - critic/grad_norm:1023.278 - critic/lr(1e-4):0.100 - timing/update_actor:20.314 - actor/entropy_loss:0.433 - actor/pg_loss:-0.005 - actor/pg_clipfrac:0.000 - actor/ppo_kl:0.000 - actor/grad_norm:1.992 - actor/lr(1e-4):0.010 - critic/score/mean:0.004 - critic/score/max:1.000 - critic/score/min:0.000 - critic/rewards/mean:0.004 - critic/rewards/max:1.000 - critic/rewards/min:0.000 - critic/advantages/mean:-0.000 - critic/advantages/max:2.360 - critic/advantages/min:-2.280 - critic/returns/mean:0.003 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.045 - critic/values/max:9.500 - critic/values/min:-14.000 - response_length/mean:239.133 - response_length/max:256.000 - response_length/min:77.000 - prompt_length/mean:104.883 - prompt_length/max:175.000 - prompt_length/min:68.000
step:1 - timing/gen:23.020 - timing/ref:4.322 - timing/values:5.953 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty:0.001 - timing/adv:0.118 - timing/update_critic:15.646 - critic/vf_loss:18.472 - critic/vf_clipfrac:0.384 - critic/vpred_mean:1.038 - critic/grad_norm:942.924 - critic/lr(1e-4):0.100 - timing/update_actor:20.526 - actor/entropy_loss:0.440 - actor/pg_loss:0.000 - actor/pg_clipfrac:0.002 - actor/ppo_kl:0.000 - actor/grad_norm:2.060 - actor/lr(1e-4):0.010 - critic/score/mean:0.000 - critic/score/max:0.000 - critic/score/min:0.000 - critic/rewards/mean:0.000 - critic/rewards/max:0.000 - critic/rewards/min:0.000 - critic/advantages/mean:0.000 - critic/advantages/max:2.702 - critic/advantages/min:-2.616 - critic/returns/mean:0.000 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.280 - critic/values/max:11.000 - critic/values/min:-16.000 - response_length/mean:232.242 - response_length/max:256.000 - response_length/min:91.000 - prompt_length/mean:102.398 - prompt_length/max:185.000 - prompt_length/min:70.000
查看 Algorithm Baselines 页面可获取完整的训练和验证日志作为参考。
检查点默认保存在以下目录:checkpoints/${trainer.project_name}/${trainer.experiment_name}。您可以使用 verl.model_merger 模块将保存的检查点合并到 huggingface 模型中,例如:
python3 -m verl.model_merger merge \
--backend fsdp \
--local_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor \
--target_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor/huggingface
有关检查点和模型合并的更多详细信息,请参阅 checkpoint-page。
要启用 wandb 进行实验跟踪,请设置以下配置:
trainer.logger='["console","wandb"]' \
trainer.project_name=$YOUR_PROJECT_NAME \
trainer.experiment_name=$YOUR_RUN_NAME \
如果您遇到 HBM 低于 32GB 的内存不足问题,启用以下配置将有所帮助:
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
critic.ppo_micro_batch_size_per_gpu=1 \
有关完整的配置集,请参阅 配置说明 以获取详细解释和性能调整。