图解 OpenRLHF 中基于 Ray 的分布式训练流程

本文着重分析 OpenRLHF 中的PPO-Ray 训练架构设计，没有使用过 Ray 的朋友也可以通过本文快速上手，本文共分成四块：

1. 为什么用 Ray
2. 使用图例抽象出整体训练流程
3. Ray 核心知识速过
4. 使用图例，进一步抽象出核心代码细节，包括：

• 训练入口
• 部署 PPO-Actor/Ref/Critic/RM 实例
• 部署 vllm_engines 实例
• PPO-Actor 与 vllm_engines 之间的通讯
• PPO-Actor/Critic 训练

一、为什么要使用 Ray

对于通常的 rlhf 框架，在训练时会在单卡上同时部署 actor/ref/reward/critic 四类模型，这种单一的部署方式可能存在如下问题：

• 难以突破单卡显存的限制。
• 无法实现更多的并行计算。例如在收集 exp 阶段，拿到(prompt, responses)结果的四类模型其实可以做并行推理；在训练阶段，拿到 exp 的 actor 和 critic 也可以做并行训练。但受到单卡显存等因素影响，通常的 rlhf 框架中使用更多的是串行。
• 无法独立优化训练和推理过程。诸如 vllm 之类的框架，是可以用来提升 actor 生成(prompt, responses)的速度的，而对于其它模型，我们也可能会视算法需要有不同的推理需求。因此我们期望能更加灵活地设计训练、推理过程

而解决以上问题，需要开发者能设计一套较为灵活的分布式计算框架，能够实现资源定制化分配、分布式调度、节点内外通信等目标，同时相关的代码不能太复杂，能够让使用者更专注于算法部分的研发。而 Ray 天然可以帮我们做这件事：我们只需提供自己的资源分配方案，告诉 Ray 我想怎么部署这些模型，不管是分开还是共同部署 Ray 都可以帮我们实现。而复杂的调度策略和通信等事项，就由 Ray 在后台去做，我们无需关心这个过程。

二、整体流程

本节我们将提供 2 个例子，帮助大家更好理解使用 Ray 可以做什么样的“定制化”部署。注意，这些例子只做讲解用，不代表它们一定是训练的最优配置。

2.1 非共同部署

这个例子展示如何完全独立部署各个模型。假设我们有 3 台 node，每台 node 8 张卡。以下展示其中一种可行的部署方式：

（1）部署 4 类模型

在这个例子中，4 类模型分开部署在 node0 和 node1 上。以 Actor 为例，它分布在“node0 的 gpu0/1 + node1 的 gpu0/1”上。这一点是由 Ray 实现的：我们自己定制化资源分配的方案，进而管控模型的分配方式

而当实际训练时，我们还可进一步引入 Deepspeed zero 做优化：以 Actor 为例，上图中的 4 个 Actor 构成 zero 中的数据并行组（world_size = 4），根据 zero 的配置，我们可以在这 4 张卡间做 optimizer/gradients/weights 的切片。

（2）部署 vllm_engines

前文说过，对于 Actor 模型，在收集 exp 阶段我们可以采用 vllm 之类的框架加速(prompt, responses)的生成。在这个例子中：

• 1 个 vllm_engine 维护着一个 vllm 实例，每个 vllm 实例下维护一个完整的 Actor 模型，这里我们还假设一个 vllm 实例按 tp_size = 2 的方法切割模型。
• 在 node2 中，共有 4 个 vllm_engines（也即 4 个 vllm 实例），这种分配方式是通过 Ray 实现的。而每个 vllm 实例内的分布式推理则是由 vllm 自己管控。

（3）Actor 与 vllm_engines 之间的通讯

我们称：

• vllm_engines 中的 actor 为 vllm_actor
• node0/1 中的 actor 为 ds_actor

在整个训练过程中，vllm_actor 需要和 ds_actor 保持权重一致。我们来看这个一致性是如何维护的：

1. 初始化阶段

假设 pretrain 路径下存储着 sft 模型，当我们首次开始训练时，ds_actor 和 vllm_actor 都直接从 pretrain 中加载权重，两者互不影响，独立加载。

2. 训练中

在 1 个 step 结束后，ds_actor 需要把更新后的权重 broadcast 给 vllm_actor，具体步骤如下：

• 首先，对ds_rank0 + all_vllm_ranks创建一个通讯组。在本例中:
• node0/gpu0 上的 actor 是 ds_rank0
• node2 中所有的 gpu 构成 all_vllm_ranks。
• 我们就是把这两者纳入一个通讯组内，这个通讯组的 world_size = 9。如果我们多一台 node3 来做 vllm_engines，那么这个通讯组的 world_size = 19，以此类推。
• 若我们使用 ds_zero1/2，则 ds_rank0 上维护的是完整的 actor 权重，我们把 ds_rank0 上的权重 broadcast 到每一个 vllm_rank，如有设置 tp，vllm 会自动帮我们完整接下来的模型切割。
• 若我们使用 ds_zero3，则 ds_rank0 上只维护部分 actor 权重，那么：
• ds_rank0 先从 ds_actor 组内 all gather 回完整的模型权重
• 再将完整的模型权重 brocast 给每一个 vllm_rank

3. 从检查点恢复训练（load_checkpoint）

当我们需要从检查点恢复训练时，ds_actor 会负责把检查点权重 broadcast 给 vllm_actor，方式同 2。

（4）整体运作流程

结合 2.1 开头的图例，我们来简述一下整体运作流程。

• 首先明确一些表达。例如，node0中的Actor0/1 + node1中的Actor0/1属于相同的数据并行组，所以接下来我们会用它们在 dp 组中的 rank 来描述它们，也就是分别改称 Actor0/1/2/3。对于其余三类模型也是同理。
• 接着进行分组：
• Actor0 / Ref0 / RM0 / Critic0 / vllm_engine0为一组
• Actor1 / Ref1 / RM1 / Critic1 / vllm_engine1为一组
• Actor2 / Ref2 / RM2 / Critic2 / vllm_engine2为一组
• Actor3 / Ref3 / RM3 / Critic3 / vllm_engine3为一组
• 你可以把每一组想象成原来的一张单卡，那么它的作用就是负责一个 micro_batch 的训练，这样我们就能大致想象到它们之间是如何配合运作的了。需要注意的是，在我们的例子中，这些实例都是一一对应的（各自有 4 个实例），但在实际操作中，根据不同用户的资源配置，不一定存在这个一一对应的关系。例如你可能用 4 卡部署 Actor，2 卡部署 Critic，8 个 vllm_engines...以此类推。不管怎样，我们应该尽量在处理 micro_bathes 的各个组间均匀分配负载，在代码里相关的操作如下：

1. 为每个 actor 分配其对应的 critic/reward/ref，并启动每个分组的训练：https://github.com/OpenRLHF/O...
2. 为每个 actor 分配对应的 vllm_engine，并使用 vllm_engine 进行推理：https://github.com/OpenRLHF/O...

2.2 共同部署

同样，我们可以按照自己的需求，选择性地在单卡上部署不同种类的模型，例如下面的例子中，actor/ref 共部署，critic/remote 共部署，图例如下，运作流程和 2.1 相似，这里不赘述：

三、Ray 的核心概念

在传统的编程中，我们经常使用到 2 个核心概念：function 和 class。而在分布式系统中，我们希望可以分布式并行执行这些 function 和 class。Ray 使用装饰器@ray.remote 来将 function 包装成 Ray task，将 class 包装成 Ray actor，包装过后的结果可以在远程并行执行。接下来我们就来细看 task/actor，请大家特别关注代码中的注释

（注意，这里的 actor 是 ray 中的概念，不是 rlhf-ppo 中 actor 模型的概念）

3.1 Ray Task

import ray
ray.init()

@ray.remote
def f(x):
    return x * x
# ===================================================================
# 创建driver进程，运行main
# ===================================================================
if __name__ == "__main__":
    # ===================================================================
    # 创建4个worker进程，可以在远端并行执行。
    # 每执行1次f.remote(i)，会发生以下事情：
    # - 创建1个worker进程，它将在远端执行函数f(i)
    # - 在driver进程上立刻返回一个引用（feature）,该引用指向f(i)远程计算的结果
    # ===================================================================
    futures = [f.remote(i) for i in range(4)]
    # ===================================================================
    # 阻塞/同步操作：等待4个worker进程全部计算完毕
    # ===================================================================
    results = ray.get(futures)) 
    # ===================================================================
    # 确保全部计算完毕后，在driver进程上print结果
    # ===================================================================
    print(f"The final result is: {results}") # [0, 1, 4, 9]

3.2 Ray Actor

import ray
ray.init()

@ray.remote
class Counter(object):
    def __init__(self):
        self.x = 0
    
    def inc(self):
        self.x += 1
    
    def get_value(self):
        return self.x

# ===================================================================
# 创建driver进程，运行main
# ===================================================================
if __name__ == "__main__":
    # ===================================================================
    # 创建1个worker进程，具体做了以下事情：
    # - 在远端创建Counter实例
    # - 在driver端即刻返回对该实例的引用c（称为actor handle）
    # - 我们可以在Ray集群的任何结点上传递和使用这个actor handle。即在任何地方，
    #   我们可以通过c来invoke对应Counter实例下的各种方法
    # ===================================================================
    c = Counter.remote()

    # ===================================================================
    # 阻塞/同步：通过c来invoke远端Counter实例的get_value()方法，并确保方法执行完毕。
    # 执行完毕后才能接着执行driver进程上剩下的代码操作
    # ===================================================================
    print(ray.get(c.get_value.remote()))  # 0
    
    # ===================================================================
    # Increment the counter twice and check the value again.
    # 道理同上，不赘述
    # ===================================================================
    c.inc.remote()
    c.inc.remote()
    print(ray.get(c.get_value.remote()))  # 2

3.3 Ray cluster 架构简图

现在我们已经通过以上例子对 Ray 运作原理有了一些基本感知，我们来进一步探索一个 ray cluster 的组成：

• 在一个 ray cluster 中，会有一台 head node 和若干 worker node
• Driver process是一种特殊的 worker process，它一般负责执行 top-level application（例如 python 中的__main__），它负责提交想要执行的任务，但却不负责实际执行它们。理论上 driver process 可以运行在任何一台 node 内，但默认创建在 head node 内。
• Worker process负责实际任务的执行（执行 Ray Task 或 Ray Actor 中的方法）。
• 每台 node 中还有一个 Raylet process，它负责管控每台 node 的调度器和共享资源的分配。
• Head node 中的 GCS将会负责维护整个 ray cluster 的相关服务。

四、代码细节

本章将解读更多代码实践上的重要细节。我们通过图例的方式抽象出代码运行的过程，而具体代码可参考文中给出的相关链接

4.1 训练入口

ppo_ray 相关的训练入口在：https://github.com/OpenRLHF/O...。

在 main 中我们启动了 driver 进程，并执行训练函数 train(args)，这里主要做了如下几件事：

• 在 ray 集群上部署 Actor/Ref/Critic/RM 实例
• 在 ray 集群上部署 vllm_engines 实例
• 训练 Actor 和 Critic 模型

我们依次来解读这三个步骤。同时为了在表述上消除歧义，我们接下来谈到“Actor”时，会使用 Ray-Actor 和 PPO-Actor 来做区分，从之前的介绍中可知，Ray-Actor 是指部署在 Ray 集群中的远端 class，PPO-Actor/Ref/Critic/RM 都属于 Ray-Actor。

4.2 部署 Actor/Ref/Critic/RM 实例

（1）非共同部署

针对图 2.1 的情况，我们以 PPO-Actor 为例，看代码是如何将其部署到 Ray 集群上的。

• PPORayActorGroup：创建在 driver 进程上，可将它理解成一种部署方案，专门负责部署 PPO 中的 4 类模型。
• PPORayActorGroup中维护着self._actor_handlers，它是一个List[ray.actor.ActorHandle]，列表中每个元素表示某个远端 Ray-Actor 的引用，而这个远端 Ray-Actor 可以是 PPO-Actor/Ref/Critic/RM 实例。如前文所说，我们可以在 ray 集群中的任何位置调用这个 handler，来对相应的远端 Ray-Actor 执行操作。
• 在本例中，我们创建了 4 个 Ray-Actor（1 个 master-actor，3 个 worker_actor）。每个 Ray-Actor 都运行在一个 worker 进程中。在创建 Ray-Actor 的同时，我们也会去修改 worker 进程的环境变量。后续当我们在这些 worker 进程中启动 ds_zero 相关的分布式配置时，ds 会读取这些环境变量信息，这样我们就知道哪些 Ray-Actor 同时由构成 ds 中的数据并行组。
• 使用PPORayActorGroup部署模型实例的代码如下：

model = PPORayActorGroup(
        # 为部署该模型的全部实例，我们想用多少台node，例如本例中为2
        args.actor_num_nodes,
        # 为部署该模型的全部实例，我们每台node上想用多少gpu，例如本例中为2
        args.actor_num_gpus_per_node,
        # Actor/Critic/Reward/ReferenceRayActor
        ActorModelRayActor, 
        # pg可理解为，在ray cluster中锁定/预留一片资源，然后只在这片资源上部署该模型全部实例。
        # （pg维护在Head Node的GCS上，参见3.3）
        # 例如本例中，pg锁定的资源为node0 gpu0/1, node1 gpu0/1，
        # 我们只在上面部署ActorModelRayActor全部实例
        pg=pg,
        # 当我们在pg指向的预留资源中分配模型实例时，再进一步指定每个实例占据一张gpu的多少部分
        # 等于1说明每个实例占满一张gpu，即“非共同部署”
        # 小于1说明每个实例只占部分gpu，即“共同部署”，例如PPO-Actor/Ref共同部署在一张卡上
        num_gpus_per_actor=0.75 if pg else 1,
    )

• ActorModelRayActor：创建在远端 worker 进程上，是 Ray-Actor。它包含了设置 ds_zero 分布式环境、加载模型权重、数据集准备、optimizer/scheduler 准备、训练等一系列操作。

PPORayActorGroup 代码参见：https://github.com/OpenRLHF/O...

根据这份代码，大家可自行去找 Actor/Critic/Reward/ReferenceRayActor 的相关实现。

（2）共同部署

针对图 2.2 的情况，我们以 PPO-Actor 为例，看代码是如何将其部署到 Ray 集群上的。

• PPORayActorGroup：在 driver 进程上创建 2 个 PPORayActorGroup，分别管理 PPO-Actor，PPO-Ref 的部署
• 使用actor_model = PPORayActorGroup(..., pg = pg, num_gpus_per_actor=0.75)创建 PPO-Actor 部署方案实例；使用ref_model = PPORayActorGroup(..., pg = pg, num_gpus_per_actor=0.25)创建 PPO-Ref 部署方案实例
• 这里，两个方案实例使用的 pg 都是同一个，即这个 pg 都指向“1 台 node，每台 node 8 张卡”这片预留好的资源。
• num_gpus_per_actor = 0.75/0.25 是一种创建 trick，虽然我们的最终目的是为了让 PPO-Actor 和 PPO-Ref 对半分一张卡，但是：
• 假设设置为 0.5，当我们实际部署 ActorModelRayActor 时，Ray 先在单卡上部署 1 个 ActorModelRayActor 实例，当它准备部署第二个 ActorModelRayActor 实例时，它发现由于每个实例只占 0.5 块卡，因此完全可以把第二个实例接着第一个实例部署，这样就导致最终无法让 PPO-Actor 和 PPO-Ref 共享一张卡
• 假设设置 0.75，当我们在单卡上部署完 1 个 ActorModelRayActor 实例后，ray 发现单卡剩下的空间不足以部署第 2 个 ActorModelRayActor 实例，所以就会把第二个实例部署到别的卡上，这样最终实现 PPO-Actor 和 PPO-Ref 共享一张卡
• 所以，这个设置是为了达到不同类型模型的实例共享一张卡的目的，而并非真正指模型实际占据的单卡显存空间。
• 最后，在这一步中，我们对全部 ActorModelRayActor 共创建 8 个 worker 进程，对全部 RefenreceModelRayActor 共创建 8 个 worker 进程，一共创建 16 个工作进程。

相关代码依然在：https://github.com/OpenRLHF/O...

4.3 部署 vllm_engines 实例

• create_vllm_engines：在 driver 端，我们通过运行该函数来创建 vllm_engines，过程相似于 4.2 节中的介绍，信息都在图中，这里不赘述。
• LLMRayActor：worker 端 Ray-Actor，它主要是把 vllm 实例进行了一些包装，包装的目的是为了让 ds_rank0 和 all vllm ranks 间可以进行 PPO-Actor 的权重通讯（参见 2.1（3））
• 在上面的例子中，我们会创建 4 个 worker 进程，用于运行管理 4 个 vllm_engine。在每个 worker 进程内，vllm 实例还会创建属于自己的 worker 进程做分布式运行。

相关代码参见：

• https://github.com/OpenRLHF/O...
• https://github.com/OpenRLHF/O...

4.4 ds_rank0 与 vllm_ranks 之间的通讯

在 2.2 中，我们说过，PPO-Actor 的 ds_rank0 需要和 all_vllm_ranks 进行通讯，传递最新的 PPO-Actor 权重，例如以下 ds_rank0 要把完整的权重 broadcast 给 16 个 vllm_ranks：

我们分成如下几步实现这个目标：

（1）创建通信组

如上图所示，创建通信组实际包含了 2 步。

Step1：

代码来自：https://github.com/OpenRLHF/O...

这段代码执行在 PPO-Actor0（ds_rank0）所在的 worker 进程中。这个 worker 进程将通过 handler 引用，触发远端每个 vllm_engine 上的 init_process_group 操作，并将 ds_rank0 纳入通讯组

        # Create torch group with deepspeed rank 0 and all vllm ranks
        # to update vllm engine's weights after each training stage.
        #
        # Say we have 3 vllm engines and eache of them has 4 GPUs,
        # then the torch group is:
        # [    0,      1, 2, 3, 4,  5, 6, 7, 8,  9, 10, 11, 12]
        # |ds rank 0 |  engine-0  |  engine-1  |   engine-2   |
        #
        # For ZeRO-1/2:
        #   1. Broadcast parameters from rank 0 to all vllm engines
        # For ZeRO-3:
        #   1. AllGather paramters to rank 0
        #   2. Broadcast parameters from rank 0 to all vllm engines
        if self.vllm_engines is not None and torch.distributed.get_rank() == 0:
            ...
            # world_size = num_of_all_vllm_ranks + 1 ds_rank0
            world_size = vllm_num_engines * vllm_tensor_parallel_size + 1
            ...
            # =====================================================================
            # 遍历每个vllm_engines，将其下的每个vllm_rank添加进通讯组中，这里又分成两步：
            # 1. engine.init_process_group.remote(...)：
            #    首先，触发远程vllm_engine的init_process_group方法
            # 2. 远程vllm_engine是一个包装过的vllm实例，它的init_process_group
            #    方法将进一步触发这个vllm实例下的各个worker进程（见4.4图例），
            #    最终是在这些worker进程上执行“将每个vllm_rank"添加进ds_rank0通讯组的工作
            # =====================================================================
            refs = [
                engine.init_process_group.remote(
                    # ds_rank0所在node addr
                    master_address, 
                    # ds_rank0所在node port
                    master_port,
                    # 该vllm_engine的第一个rank在"ds_rank0 + all_vllm_ranks“中的global_rank，
                    # 该值将作为一个offset，以该值为起点，可以推算出该vllm_engine中其余vllm_rank的global_rank
                    i * vllm_tensor_parallel_size + 1, 
                    world_size,
                    "openrlhf",
                    backend=backend,
                )
                for i, engine in enumerate(self.vllm_engines)
            ]
            # =====================================================================
            # 将ds_rank0添加进通讯组中
            # =====================================================================
            self._model_update_group = init_process_group(
                backend=backend,
                init_method=f"tcp://{master_address}:{master_port}",
                world_size=world_size,
                rank=0,
                group_name="openrlhf",
            )
            # =====================================================================
            # 确保all_vllm_ranks都已添加进通讯组中
            # =====================================================================
            ray.get(refs)

Step2:

代码来自：https://github.com/OpenRLHF/O...

这段代码实际运行在每个 vllm_engine（即每个包装后的 vllm 实例）下的 worker 进程内。例如 tp_size=2，那么每个 vllm 实例下就有 2 个 worker 进程，这两个 worker 进程都会运行这段代码。

class WorkerWrap(Worker):
    def init_process_group(self, master_address, master_port, rank_offset, world_size, group_name, backend="nccl"):
        """Init torch process group for model weights update"""
        assert torch.distributed.is_initialized(), f"default torch process group must be initialized"
        assert group_name != "", f"group name must not be empty"
        # =====================================================================
        # torch.distributed.get_rank(): 在当前vllm_engine内部的rank，
        #                               例如在tp_size = 2时，这个值要么是0，要么是1
        # rank_offset：当前vllm_engine中的第一个rank在“ds_rank0 + all_vllm_ranks"中的global_rank
        # 两者相加：最终得到当前rank在“ds_rank0 + all_vllm_ranks"中的global_rank
        # =====================================================================
        rank = torch.distributed.get_rank() + rank_offset
        self._model_update_group = init_process_group(
            backend=backend,
            init_method=f"tcp://{master_address}:{master_port}",
            world_size=world_size,
            rank=rank,
            group_name=group_name,
        )
        ...

（2）_broadcast_to_vllm

构建好通讯组，我们就可以从 ds_rank0 广播 PPO-Actor 权重到 all_vllm_ranks 上了，这里也分成两步。

Step1：PPO-Actor ds_rank0 发送权重

代码在：https://github.com/OpenRLHF/O...

这段代码运行在 ds_rank0 对应的 worker 进程中

    def _broadcast_to_vllm(self):
        # avoid OOM
        torch.cuda.empty_cache()
        model = self.actor.model.module
        count, num_params = 0, len(list(model.named_parameters()))
        for name, param in model.named_parameters():
            count += 1  # empty_cache at last param

            # Fire all vllm engines for broadcast
            if torch.distributed.get_rank() == 0:
                shape = param.shape if self.strategy.args.zero_stage != 3 else param.ds_shape
                refs = [
                    # 远端vllm_engine的每个rank上，初始化一个尺寸为shape的empty weight张量，
                    # 用于接收广播而来的权重
                    engine.update_weight.remote(name, dtype=param.dtype, shape=shape, empty_cache=count == num_params)
                    for engine in self.vllm_engines
                ]

            # For ZeRO-3, allgather sharded parameter and broadcast to all vllm engines by rank 0
            # ds_rank0发出权重（视是否使用zero3决定在发出前是否要做all-gather）
            with deepspeed.zero.GatheredParameters([param], enabled=self.strategy.args.zero_stage == 3):
                if torch.distributed.get_rank() == 0:
                    torch.distributed.broadcast(param.data, 0, group=self._model_update_group)
                    ray.get(refs) # 确保所有vllm_ranks接收权重完毕

Step2: 各个 vllm_ranks 接收权重

代码在：https://github.com/OpenRLHF/O...

代码运行在每个 vllm_engine(即每个包装后的 vllm 实例)下的各个 worker 进程中。例如 tp_size = 2，那么每个 vllm 实例下有 2 个 worker 进程，这 2 个 worker 进程都会运行这段代码。

    def update_weight(self, name, dtype, shape, empty_cache=False):
        """Broadcast weight to all vllm workers from source rank 0 (actor model)"""
        if torch.distributed.get_rank() == 0:
            print(f"update weight: {name}, dtype: {dtype}, shape: {shape}")

        assert dtype == self.model_config.dtype, f"mismatch dtype: src {dtype}, dst {self.model_config.dtype}"
        # 创建同尺寸空张量用于接收ds_rank0广播来的权重
        weight = torch.empty(shape, dtype=dtype, device="cuda")
        # 接收权重
        torch.distributed.broadcast(weight, 0, group=self._model_update_group)
        # 使用接收到的权重进行更新
        self.model_runner.model.load_weights(weights=[(name, weight)])

        del weight

4.5 PPO-Actor/Critic Training

正如 2.1（4）中所说，我们将部署在 ray 集群上的 PPO-Actor/Ref/Critic/RM 实例们进行分组，每组分别负责一份 micro-batch 的训练，上图刻画了某个组内的训练流程。一组内的训练流程发起自 PPO-Actor 实例（fit 方法），共分成如下步骤执行。

Step1：发送 prompts，并从 vllm_engine 上收集(prompt, response)。

代码参见：https://github.com/OpenRLHF/O...

Step2：从 Ref/Reward/Critic 上收集并处理 exps。

代码参见：https://github.com/OpenRLHF/O...

Step3: 确保将处理后的 exps 传送给 Critic，并行执行 Actor 和 Critic 的训练

• 将 exps 传送给 Critic：https://github.com/OpenRLHF/O...
• Actor 训练：https://github.com/OpenRLHF/O...
• Critic 训练：https://github.com/OpenRLHF/O...

我们在 Actor 实例所在的 worker 进程上出发 Actor 和 Critic 的训练。以上代码只给出了训练入口，更多细节需要顺着入口去阅读。

Step4：vllm_engine 权重更新。

代码参见：https://github.com/OpenRLHF/O...

五、参考

1、OpenRLHF：https://github.com/OpenRLHF/O...

2、Ray official architecture whitepaper: https://docs.google.com/docum...

(建议想看 ray 架构的朋友，直接看这个最新的官方白皮书，不要看 2018 年的那篇 paper 了，那个比较老了)

3、Ray official document：https://docs.ray.io/en/latest...

4、推荐一篇快速了解 Ray 应用层核心概念的 blog：https://towardsdatascience.co...

5、Ray：https://github.com/ray-projec...

6、vllm: https://github.com/vllm-proje...

END

来源：GiantPandaCV

推荐阅读

• 超越 YOLOv10 和 YOLOv7，专为大规模高分辨率图像处理设计 ！
• 视觉 Transformer 与目标检测的完美融合：解读 ViTOC 架构 !
• PyTorch 通讯实践
• CUDA-MODE 课程笔记 第 29 课 Triton 内部机制

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。