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：准备数据集
------------------

.. code:: bash

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

步骤 2：下载模型
----------------

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

- 从 huggingface 或 modelscope 下载所需模型

.. code:: bash

   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 <https://github.com/volcengine/verl/blob/main/verl/trainer/fsdp_sft_trainer.py>`_ 中提供了一个使用 PyTorch FSDP 的 SFT 训练器。 
用户可以使用我们的 FSDP SFT 训练器自定义自己的 SFT 脚本。

我们还在 `gsm8k sft directory <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/>`_ 中为 GSM8k 数据集提供了各种 SFT 训练脚本。

.. code:: shell

   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 内核），需要在运行脚本中添加以下环境变量：

    .. code-block:: bash

        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``。
- 有关每个配置字段的详细解释，请参阅 :doc:`config`。

**奖励模型/函数**

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

**训练脚本**

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

.. code:: bash

   cd ../ppo_trainer
   bash run_deepseek7b_llm.sh

run_deepseek7b_llm.sh 的脚本内容

.. code:: bash

   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 内核），需要在运行脚本中添加以下环境变量：

    .. code-block:: bash

        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 <https://yushengsu-thu.github.io/>`_。