GSM8K 示例

最后更新：2025年3月25日。

简介

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

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

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

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

数据集简介

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

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

一个示例

提示

Katy 用糖（茶匙）和水（杯）以 7:13 的比例制作咖啡。如果她总共使用了 120 茶匙的糖和杯水，请计算她使用了多少茶匙的糖。

解答

她制作咖啡所用配料的总比例为 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：下载模型

有三种方法可以准备用于后续训练的模型检查点：

从 huggingface 或 modelscope 下载所需模型

huggingface-cli download deepseek-ai/deepseek-math-7b-instruct --local-dir ~/models/deepseek-math-7b-instruct --local-dir-use-symlinks False
# 或
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和reward模型的检查点。

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

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

我们还在 gsm8k sft 目录中提供了各种用于在 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 kernel），则需要在运行脚本中添加以下环境变量：

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。
有关每个配置字段的详细说明，请参阅配置说明。

奖励模型/函数

我们使用基于规则的奖励模型。我们强制模型在解决方案中显示的 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 kernel），则需要在运行脚本中添加以下环境变量：

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

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