GSM8K 示例

最后更新：03/25/2025。

介绍

在本示例中，我们训练一个 LLM 来处理 GSM8k 任务。

论文：https://arxiv.org/pdf/2110.14168

数据集：https://huggingface.co/datasets/gsm8k

请注意，原论文主要关注通过最佳 N 抽样 (Best-of-N sampling) 来训练一个验证器（奖励模型）以解决数学问题。在此示例中，我们使用基于规则的奖励模型来训练一个 RLHF 代理。

数据集介绍

GSM8k 是一个数学问题数据集。提示是一个小学水平的数学问题。LLM 模型需要回答这个数学问题。

训练集包含 7473 个样本，测试集包含 1319 个样本。

一个示例

Prompt

Katy 使用糖勺和水的杯子制作咖啡，其比例为 7:13。如果她总共使用了 120 勺糖和杯子水，请计算她使用了多少勺糖。

Solution

她用来制作咖啡的配料的总比例为 7+13 = <<7+13=20>>20。由于表示勺数的分数是 7/20，她使用了 7/20*120 = <<7/20*120=42>>42 #### 42

步骤 1：准备数据集

cd examples/data_preprocess
python3 gsm8k.py --local_save_dir ~/data/gsm8k

步骤 2：下载模型

准备用于后训练 (post-training) 的模型检查点有三种方式：

从 huggingface 或 modelscope 下载所需模型

huggingface-cli download deepseek-ai/deepseek-math-7b-instruct --local-dir ~/models/deepseek-math-7b-instruct --local-dir-use-symlinks False
# or
modelscope download --model deepseek-ai/deepseek-math-7b-instruct --local_dir ~/models/deepseek-math-7b-instruct

将模型存储在本地目录或 HDFS 路径中。
此外，您可以在运行脚本中的 actor_rollout_ref.model.path 和 critic.model.path 字段中直接使用 huggingface 中的模型名称（例如，deepseek-ai/deepseek-math-7b-instruct）。您也可以通过设置环境变量 VERL_USE_MODELSCOPE=True 从 modelscope 下载模型。请参阅 examples/ppo_trainer/run_deepseek7b_llm_modelscope.sh 以获取示例。

请注意，用户应准备演员 (actor)、批评者 (critic) 和奖励模型的检查点。

[可选] 步骤 3：对您的模型进行 SFT

我们在 fsdp_sft_trainer.py 中提供了一个使用 PyTorch FSDP 的 SFT 训练器。用户可以使用我们的 FSDP SFT 训练器自定义自己的 SFT 脚本。

我们还在 gsm8k sft directory 中为 GSM8k 数据集提供了各种 SFT 训练脚本。

set -x

torchrun -m verl.trainer.fsdp_sft_trainer \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.prompt_key=question \
    data.response_key=answer \
    data.micro_batch_size_per_gpu=8 \
    model.partial_pretrain=deepseek-ai/deepseek-coder-6.7b-instruct \
    trainer.project_name=gsm8k-sft \
    trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \
    trainer.total_epochs=4 \
    trainer.logger='["console","wandb"]'

如果您使用 AMD GPU（ROCm 内核），需要在运行脚本中添加以下环境变量：

export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES
export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES

步骤 4：在 GSM8k 数据集上使用您的模型进行 PPO 训练

准备您自己的 run.sh 脚本。以下是针对 GSM8k 数据集和 deepseek-llm-7b-chat 模型的示例。
用户可以根据自己的环境替换 data.train_files、data.val_files、actor_rollout_ref.model.path 和 critic.model.path。
有关每个配置字段的详细解释，请参阅 ppo_trainer.yaml 用于 RL FSDP 后端。

奖励模型/函数

我们使用一个基于规则的奖励模型。我们强制模型产生以 4 个 “#” 开头的最终答案，如解决方案所示。我们使用正则表达式匹配从解决方案和模型输出中提取最终答案。我们比较它们，并为正确答案分配奖励 1、不正确答案分配 0.1、无答案分配 0。

训练脚本

FSDP 和 Megatron-LM 后端的训练脚本示例存储在 examples/ppo_trainer 目录中。

cd ../ppo_trainer
bash run_deepseek7b_llm.sh

run_deepseek7b_llm.sh 的脚本内容

set -x

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=1024 \
   data.max_prompt_length=512 \
   data.max_response_length=512 \
   actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \
   actor_rollout_ref.actor.optim.lr=1e-6 \
   actor_rollout_ref.model.use_remove_padding=True \
   actor_rollout_ref.actor.ppo_mini_batch_size=256 \
   actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
   actor_rollout_ref.actor.fsdp_config.param_offload=False \
   actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
   actor_rollout_ref.model.enable_gradient_checkpointing=True \
   actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \
   actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
   actor_rollout_ref.rollout.name=vllm \
   actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \
   actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \
   actor_rollout_ref.ref.fsdp_config.param_offload=True \
   critic.optim.lr=1e-5 \
   critic.model.use_remove_padding=True \
   critic.model.path=deepseek-ai/deepseek-llm-7b-chat \
   critic.model.enable_gradient_checkpointing=True \
   critic.ppo_micro_batch_size_per_gpu=32 \
   critic.model.fsdp_config.param_offload=False \
   critic.model.fsdp_config.optimizer_offload=False \
   algorithm.kl_ctrl.kl_coef=0.001 \
   trainer.critic_warmup=0 \
   trainer.logger='["console","wandb"]' \
   trainer.project_name='verl_example_gsm8k' \
   trainer.experiment_name='deepseek_llm_7b_function_rm' \
   trainer.n_gpus_per_node=8 \
   trainer.nnodes=1 \
   trainer.save_freq=-1 \
   trainer.test_freq=1 \
   trainer.total_epochs=15 $@

如果您使用 AMD GPU（ROCm 内核），需要在运行脚本中添加以下环境变量：

export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES
export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES

如果您在使用 AMD GPU 运行 VeRL 时遇到任何问题，请随时联系我 - Yusheng Su。