这篇博客主要讲解 RLHF 具体训练的框架 (DeepSpeedChat，OpenRLHF，verl) 的具体细节，包括每个框架的整体架构，架构内的各部分细节 (包括逻辑细节和代码细节)。(建议先阅读我之前关于 RLHF 的博客 The Basic Knowledge of RLHF (Reinforce Learning with Human Feedback))

RLHF 的算法流程

在之前的博客中，我们讲解了 RLHF 的三个阶段：SFT (预训练 LLM 模型 $M_\theta$)，Reward Modeling (预训练奖励模型 $r_\theta$) 和最后的 RL 训练 (使用 PPO 微调 $M_\theta$)。对于前两个阶段而言，其只存在一个模型，因此可以使用 Deepspeed，FSDP，Megatron，甚至是 Transformers 的内置 Trainer 等单模型训练框架直接进行分布式训练 (关于单模型训练框架，可以参考我之前的博客)。对于第三个阶段而言，其包含多个模型，同时不同模型的所处状态也不尽相同 (例如，reference model 和 reward model 只用于 infer，policy model 和 value model 用于 train，同时 policy model 还用于 rollout)。因此，需要在 Deepspeed/FSDP/Megatron 这种单模型的训练框架上再进行进一步的搭建以构建多模型训练框架。因此，本文的所有 RLHF 框架其实主要是聚焦于构建第三阶段的多模型训练框架。

下面，我们以 PPO 为例来了解每个框架的整体架构和每个部分的具体模块 (其他的 RL 算法如 GRPO，REINFORCE++ 等基本上都是在 PPO 的基础上减少某些模块)。首先，如图 1 所示，我们先逻辑化整理一下 PPO 的算法流程：

A. PPO 的生成阶段：即通过给定的输入，生成一系列 PPO 所训练的必要的元素，在经典 RL 中也被称作环境交互。

1. 给定 SFT 后得到的 model，将其复制为 reference model $\pi_{SFT}$ 和需要进一步训练的 actor model $\pi_{RL}$；给定 Reward Modeling 后得到的 model，将其复制为 reward model $R$ 和 critic model $V$。

2. 给定 prompt $x$，将其输入给 actor model $\pi_{RL}$ 生成对应的 response $y$，得到完整的 sequence $x + y$。($\pi_{RL}$ rollout)

3. 给定 sequence $x + y$，将其输入给 actor model $\pi_{RL}$ 和 reference model $\pi_{SFT}$ 分别生成 action logits $p_{RL}$ 和 sft logits $p_{SFT}$，并进一步计算 KL divergence $KL$。($\pi_{RL}$ 和 $\pi_{SFT}$ infer)

4. 给定 sequence $x + y$，将其输入给 reward model $R$ 和 critic model $V$ 分别生成 reward $r$ 和 value $v$。($R$ 和 $V$ infer)

5. 给定 $KL$ 和 $r$，计算得到 PPO 的 return；并通过给定 $v$，计算得到 PPO 的 advantage $A$。

B. PPO 的训练阶段：即通过 PPO 的生成阶段所得到的元素，进行 PPO 的训练，在经典 RL 中也被称作奖励学习。由于 PPO 的生成阶段的时间成本较高，因此通常对生成阶段得到的元素进行缓存，并进行多次训练。对于第 $t$ 次训练，其具体流程如下：

1. 给定 sequence $x + y$，将其输入给第 $t-1$ 次训练完的 actor model $\pi_{RL}$ 生成 new action logits $p^{t}_{RL}$，并和 action logits $p_{RL}$ 计算 ratio $r^{t}(\theta)$。接着和给定的 advantage $A$ 计算 Actor loss 用于 actor model 的训练。

2. 给定 sequence $x + y$，将其输入给第 $t-1$ 次训练完的 critic model $V$ 生成 new value $v^{t}$，并和 value $v$ 计算 clipped value $v_{clip}$。接着和给定的 return 计算 Critic loss 用于 critic model 的训练。

DeepSpeedChat

可以看到，上述的流程涉及到 actor model $\pi_{RL}$ 的 rollout，actor model $\pi_{RL}$ 和 $\pi_{SFT}$ 的 infer，reward model $R$ 和 critic model $V$ 的 infer，以及 actor model $\pi_{RL}$ 和 critic model $V$ 的 train。最直接的实现方式是，按照上述流程的逻辑编写 PPO 训练的架构，通过简单扩展单模型训练框架得到多模型训练框架。如图 2 所示，DeepSpeedChat 就是按照这种思路扩展 DeepSpeed 框架来实现 PPO 的训练的。(下面讲解的 DeepSpeedChat 的版本为 bd47e5bc38d292f44bf183e7bda992cde36a769b)

接下来，我们讲解 DeepSpeedChat 的每个模块的逻辑和代码细节：

与之前的博客中讲解的 DeepSpeed 的分布式训练一致，DeepSpeedChat 通过 deepspeed 命令启动分布式训练，并在每张卡上运行 main.py 文件。

1. 第一阶段：在 main.py 文件中，首先是初始化分布式环境，加载 tokenizer，加载数据集。接着，初始化 DeepSpeedRLHFEngine，包括初始化 actor model，ref model，reward model 和 value model，其中 actor model 和 critic model 是通过get_train_ds_config()获取 train 的 config 实现 DeepSpeedEngine 的初始化，而 ref model 和 critic model是通过get_eval_ds_config()获取 inferinfer的 config 实现 DeepSpeedEngine 的初始化。然后，初始化 DeepSpeedPPOTrainer，这是一个训练类，里面包含了所有关于 PPO 的计算函数。其初始化的过程就是简单的对 PPO 所需的各个 model 和系数进行赋值。最后，初始化 MiniDataset，其用于缓存后续 PPO 生成阶段的结果，并提供给 PPO 训练阶段使用。

2. 第二阶段：在完成初始化和数据的准备后，下一步便开始 PPO 的生成和训练阶段。在 PPO 的生成阶段，给定准备好的 prompt，首先生成 response，DeepSpeedChat 使用了最原始的 model.generate() 的方式。接着，将生成好的 seq (prompt + response) 输入给 actor model 生成 $p_{RL}$，输入给 ref model 生成 $p_{SFT}$，输入给 reward model 生成 $R$，最后输入给 critic model 生成 $V$，最终将生成的所有结果缓存到 MiniDataset 中。

3. 第三阶段：完成 PPO 的生成阶段，便是 PPO 的训练阶段。可以发现，为了更好地利用 PPO 生成阶段所生成的结果，一般会使用其进行多次训练，对应代码中的 args.ppo_epochs。对于第 $t$ 次训练，首先是根据 PPO 生成阶段的 $p_{RL}$，$p_{SFT}$ 和 $R$ 计算 old_rewards (这里可能会有些奇怪，在上述的RLHF 的算法流程中似乎没有这个流程，其实这里就是计算 $KL$，并将其融入到 $R$ 中进行后续的计算而已)。接着便是标准的使用 GAE 生成 $A$ 和 return。对于 actor model，接着计算 $p^t_{RL}$，并和 $p_{RL}$ 计算 $r^t(\theta)$ (这里需要注意，因为 actor model 已经经过了 $t-1$ 次的更新，因此计算得到的 $p^t_{RL}$ 和 $p_{RL}$ 不相等，如果 args.ppo_epochs 等于 $1$，那么其二者会一直相等，即 $r^t(\theta)$ 恒等于 $1$)。最后计算 actor loss，并进行 backward 和 actor model 更新。而对于 critic model，接着计算 $v^t$，并和 $v$ 计算 $v_{clip}$ (这里需要注意，因为 critic model 也已经经过了 $t-1$ 次的更新，因此计算得到的 $v^t$ 和 $v$ 不相等)。最后和 return 计算 critic loss，并进行 backward 和 critic model 更新。

至此，DeepSpeedChat 的逻辑和代码细节便已讲解完毕。可以发现，DeepSpeedChat 的 PPO 训练框架的构建和上述的 PPO 的算法流程的逻辑是一致的，因此，DeepSpeedChat 的 PPO 训练框架的构建是相对容易的。同时，其将所有 model 都分配到了相同的设备上，即共用相同的 GPU 资源，如图 3 所示。这种做法的好处是可以简化代码，但是其缺点是无法实现不同 model 的并行计算，即每类 model 只能顺序执行计算，如 actor model infer 完后再进行 ref model infer。其次，对于每个 model 的不同阶段，其所需要的分布式优化不同，例如 actor model 在 rollout 时可能需要 vllm 的分布式优化，而在 train 时需要 DeepSpeed 的分布式优化。而 DeepSpeedChat 统一使用 DeepSpeed (虽然其存在 HybridEngine)，导致在 rollout 时 GPU 资源利用率不高。这对于单机多卡而言，由于卡间通信较快，其可以一定程度上缓解，但是对于多机多卡而言，其资源利用率会非常低。

我们将 DeepSpeedChat 这种将所有 model 都分配到相同的 GPU 资源上的结构称为 collocate all models。而在理想的情况下，各个 model 在 GPU 资源上的结构应该如图 4 所示。首先，将 actor model 复制为 $2$ 份，一份用于 train，使用 TrainEngine (如 DeepSpeed, FSDP, Megatron) 进行优化，称为 $\pi_{train}$；而另一份用于 rollout，使用 InferEngine (如 vllm, sglang) 进行优化，称为 $\pi_{rollout}$，并在每次 PPO 训练阶段完成后，下一次 PPO 生成阶段开始前，将更新后的 $\pi_{train}$ 的参数同步给 $\pi_{rollout}$。这样做的目的是可以更好地利用目前开源的各个 train/infer engine，提升各个阶段的效率。其次，将每个 model 分配到不同的 GPU 资源上，使其独占给定的 GPU 资源，这样，图 1 中那些没有数据依赖关系的计算模块就可以并行，从而节省整体的时间开销。

在将每个 model 都分配到不同的 GPU 资源之后，RLHF 的整个流程就可以引入 model 并行推理，形成如图 5 所示的逻辑流程。可以看到，大多数的计算模块都可以并行进行，与图 1 相比，其可以节省大量的时间开销。下面要讲的 OpenRLHF 和 verl 都是使用这种并行的逻辑流程来编写代码的。

与 DeepSpeedChat 一开始就使用 deepspeed 命令启动分布式，并在每个子进程中运行 main.py 不同。关于图 5 所示的逻辑流程的代码编写，由于其需要模块并行，即每个 model 的分布式进程组执行的模块不同 (例如 actor model 的分布式进程组在生成 action logits 时，ref model 的分布式进程组在同时生成 sft logits)。因此最直观，也是最具扩展性的方式是使用一个主进程来编写 PPO 的整体计算逻辑 (这个主进程也被称为 single controller)，在遇到分布式初始化/计算时，则异步启动/调用各个 model 的分布式进程组，然后继续主进程的下一步计算逻辑，并在之后需要原先分布式进程组结果的时候获取它。因此，整体的代码训练框架如图 6 所示 (由于篇幅限制，这里只展示一小部分代码逻辑)。

题外话：原本的 OpenRLHF 的代码不是 single controller 的模式，而是将 PPO 的计算逻辑分散到各个 model 的分布式进程中，导致其很难扩展。不过好在现在已经重构为 single controller 的模式了。

那么如何异步地启动/调用不同 model 的分布式进程组呢？目前 OpenRLHF 和 verl 都采用了 ray 来实现这一目的。ray 有些类似于计算集群管理和调度的软件，通过 ray start 或者 ray.init() 来启动 ray，并指定集群所拥有的 CPU 数，GPU 数等计算资源。接着使用装饰符 @ray.remote() 将某个函数/类装饰为一个 Task/Actor (可以初略地理解为任务)，则在后续调用该任务时，ray 会自动将其异步地调度到目前可用的计算资源上，从而减轻我们编写异步代码的难度。

由于 OpenRLHF 和 verl 都是基于 ray 来构建，同时构建逻辑遵循图 5，因此其代码结构有些相似。但是不同的是 OpenRLHF 的分布式进程组的后端使用的是 DeepSpeed 和 vllm；而 verl 的分布式进程组的后端使用的是 DeepSpeed/FSDP 和 vllm/sglang。同时，OpenRLHF 使用 PPORayActorGroup 封装每个 model 的分布式进程组；接着通过统一的 async_run_method_batch() 来调用每个 model 的统一接口 execute_batch(method_name, ...)，根据提供的 method_name 的不同来调用不同 model 的具体方法。而 verl 则是使用 WorkerDict 封装所有 model 的分布式进程组，并使用 _bind_workers_method_to_parent() 将所有 model 的特有方法 (含有 MAGIC_ATTR 属性的方法)绑定到 WorkerDict 自身上，从而通过直接调用自身的方法来调用不同 model 的具体方法。

OpenRLHF

接下来，我们讲解 OpenRLHF 的每个模块的逻辑和代码细节：(下面讲解的 OpenRLHF 的版本为 494850f50342ed38d5ae76ef45a3207f3523b582)

如图 7 所示 (这里直接盗用 OpenRLHF 的图片🥳)，OpenRLHF 的每个模块的逻辑细节和图 5 类似。在启动时，通过 ray job submit 的方式向 ray 提交运行主进程的 train_ppo_ray.py 文件。在后续的代码分析中，我们先聚焦于每个 model 各自分配不同的 GPU 资源，后面再讲解 collocate all models 的实现。train_ppo_ray.py 包括 $3$ 个阶段。第一阶段：首先，实例化 strategy，其是一个 DeepspeedStrategy 类，主要用于构建 model 的 DeepSpeed 分布式进程组，model 的 backward 和参数更新等一系列与分布式初始化/计算以及操作有关的内容。接着是构建 vllm 封装的 actor model 用于 rollout，以及构建 PPORayActorGroup 封装的 actor model (用于 train 和 infer)，ref model (用于 infer)，critic model (用于 train 和 infer)，和 reward model (用于 infer)。以 actor model 为例，PPORayActorGroup 的初始化主要涉及初始化 actor model 的分布式进程组的预处理。先是使用 ray 为分布式进程组预分配 GPU 等计算资源，接着初始化 master actor，得到 master actor 的 master addr 和 master port 后，基于此继续初始化剩下的 work actor。以 master actor 的初始化为例，主要是初始化 world size，rank，master addr，master port 等环境变量，为后面真正的 model 的分布式进程组初始化做准备。

第二阶段：首先是初始化 PPOTrainer，它是 PPO 训练和生成阶段的主要类，其初始化主要包括初始化 tokenizer，将之前初始化的各个 model 的分布式进程组赋值给自身的变量，以及 SamplesGenerator 的初始化 (用于使用 vllm 封装的 actor model 进行 rollout)，RemoteExperienceMaker 的初始化 (用于 PPO 的生成阶段)，构建数据集等其他参数/日志的初始化。接着是 actor model，ref model，critic model 和 reward model 的分布式进程组的初始化。以 actor model 为例，其主要是调用每个分布式进程的 init_model_from_pretrained() 方法。在该方法中，首先调用 strategy 的 setup_distributed() 来初始化分布式进程组的通信后端和分布式进程组的 device mesh。接着初始化 Actor 类，其主要是加载指定的 hf model，接着初始化 tokenizer，优化器，学习率 scheduler，并通过 strategy 的 prepare() 方法将 model，优化器和学习率 scheduler 使用 deepspeed 后端进行封装。最后是构建 ActorPPOTrainer 用于 actor model 的 train。至此，PPO 的所有初始化便结束了。

第三阶段：这一阶段主要是使用 PPOTrainer 的 fit() 方法执行 PPO 的生成和训练阶段。在该方法中，首先是加载 checkpoint (这里的 checkpoint 指的是整个训练环境，包括 model，dataloader等的 checkpoint state)。接着开始 args.num_episodes 次 epoch 的训练。对于每一次训练，首先是使用 SamplesGenerator 的 generate_samples() 生成 response，其主要是调用 _generate_vllm() 方法，将给定的 prompt 均匀分配给每一个 vllm 包装的 actor model，并通过 vllm 的 add_requests() 和 get_responses() 来请求和返回生成的 response，最后将其整理并使用 Experience 类进行存放。得到 response 后，接着便使用 RemoteExperienceMaker 的 make_experience_batch() 方法执行 PPO 的生成阶段的剩下部分，包括将 response 组成 batch，调用 make_experience() 方法生成 $p_{RL}$，$p_{SFT}$，$V$，$R$ 以及计算 $KL$，最后调用 compute_advantages_and_returns() 方法生成 $A$ 和 return。如前所述，make_experience() 和 compute_advantages_and_returns() 这两个方法是通过调用各个 model 的分布式进程组的方式为使用统一的 async_run_method_batch() 接口，并传入具体需要调用的 model 方法的名字来实现的。在 make_experience() 方法中，首先是调用 reward model 的 forward() 方法生成 reward $R$，然后是调用 actor model 的 forward() 方法获取 action logits $p_{RL}$，接着是调用 critic model 的 forward() 方法获取 value $V$，调用 ref model 的 forward() 方法获取 sft logits $p_{SFT}$，最后利用 compute_approx_kl() 方法 (其中包括 k1，k2，k3 三种方式近似计算 $KL$) 计算 KL divergence $KL$。在 compute_advantages_and_returns() 方法中，首先是对 reward 的后处理 (主要是 PPO 外的其他 RL 算法需要)，接着是使用 compute_reward() 方法将 $KL$ 融入到 reward 中，然后是使用 get_advantages_and_returns() 方法计算 advantage $A$ 和 return，最后对 $A$ 进行归一化处理。而 get_advantages_and_returns() 方法里的内容和 DeepSpeedChat类似，是标准的使用 GAE 计算 advantage $A$ 和 return 的过程。

在使用 RemoteExperienceMaker 的 make_experience_batch() 方法完成 PPO 的生成阶段后，接下来便是 PPO 的训练阶段。首先是调用 actor model 的 append() 方法和 critic model 的 append() 方法将 PPO 生成阶段的结果存到每个分布式进程的 NaiveReplayBuffer 中，接着便使用 ppo_train() 方法执行 PPO 训练阶段，其内部是调用了 critic model 的 fit() 方法和 actor model 的 fit() 方法分别进行 actor model 和 critic model 的 train，最后将更新的 actor model 参数使用 _broadcast_to_vllm() 方法同步到 vllm 封装的 actor model 中。在 critic model 的 fit() 方法中，主要是调用第二阶段初始化的 CriticPPOTrainer 的 ppo_train() 方法。而 ppo_train() 方法内主要包括将 NaiveReplayBuffer 中的 PPO 生成阶段的结果封装为 DataLoader，对于 Dataloader 中的每个 batch，调用 training_step() 方法进行 train (包括生成 new value $V^t$，计算 critic loss，调用 strategy 的 backward() 方法和 optimizer_step() 方法进行 backward 和 critic model 的参数更新)。在 actor model 的 fit() 方法中，主要是调用第二阶段初始化的 ActorPPOTrainer 的 ppo_train() 方法。而 ppo_train() 方法内主要包括将 NaiveReplayBuffer 中的 PPO 生成阶段的结果封装为 DataLoader，对于 Dataloader 中的每个 batch，调用 training_step() 方法进行 train (包括生成 new action logits $p^t_{RL}$，计算 actor loss 和可选的其他辅助 loss，调用 strategy 的 backward() 方法和 optimizer_step() 方法进行 backward 和 actor model 的参数更新)。在 _broadcast_to_vllm() 方法中，主要是调用 actor model 分布式进程组的 broadcast_to_vllm() 方法，进而调用 ActorPPOTrainer 的 broadcast_to_vllm() 方法 (这个方法细节见 Appendix A) 来实现两个 actor model 之间的参数同步。

至此，OpenRLHF 的逻辑和代码细节便已讲解完毕。可以发现，通过 ray 来编写整个代码框架非常的便捷，除了前述的资源管理和分配的优势外，在编写代码时只需要在需要调用 model 的分布式进程组的方法时执行 func_handler=function.remote() 得到执行结果的句柄 (注意，此时函数真正的结果可能还没执行完毕)，接着在需要得到执行结果时使用 ray.get(func_handler) 获得结果，即可实现异步的程序执行。同时，由于执行时只获得句柄，如果需要不同 model 的分布式进程组数据转移的逻辑编写，也只需要将在主进程执行 model_1 的 function_1 func_handler=model_1.function_1.remote() 获得句柄，并传递给 model_2 的 function_2 model_2.function_2.remote(func_handler)，并在 function_2 里执行 ray.get(func_handler) 即可直接实现由 model_1 向 model_2 传递结果，而不需要主进程作为中介 (主进程只起到一个传递句柄的作用，而不是传递真正的数据)。

verl

最后，我们讲解 verl 的每个模块的逻辑和代码细节：(下面讲解的 verl 的版本为 78532923368aeb058f62201489546d013df47710)

Appendix A