简介(未完成)
rlhf 和 llm-reasoning 这两个方向有什么区别吗?prompt 的难度和 response 的长度。以往我们使用 rlhf 的场景主要是:安全问题、诗歌创作修复韵脚问题、简单的代码数学问题等,往往都是几百个 token 就搞定的任务。现在不一样了,模型的一条回复就达到了上万个 token,这会让显存压力和解码时长陡增。作为应对,vllm 或 sglang 已经成为 rlhf 框架的标配。最初我也认为 grpo 省掉 critic_model 这一点并不关键,现在看来我还是只考虑了算法和数据的视角,并没有充分理解到额外维护一个和 actor_model 相同规模的 critic_model,对训练框架的稳定性有多大的挑战。当模型 size 和 response_len 逐渐增大,“训练效率和预防训练过程中莫名其妙的 OOM ”就是复现 r1 工作中最大的难点(对,就是莫名其妙,在程序断了之前,你不会知道明明限制了最大 response_len,为啥它还会 OOM)。
理念
豆包大模型团队发布全新 RLHF 框架,现已开源! 待细读。
在深度学习中,数据流(DataFlow)是一种重要的计算模式抽象,用于表示数据经过一系列复杂计算后实现特定功能。神经网络的计算就是典型的 DataFlow ,可以用计算图(Computational Graph)来描述,其中节点代表计算操作,边表示数据依赖。大模型 RL 的计算流程比传统神经网络更为复杂。在 RLHF 中,需要同时训练多个模型,如 Actor 、Critic 、参考策略(Reference Policy)和奖励模型(Reward Model),并在它们之间传递大量数据。这些模型涉及不同的计算类型(前向反向传播、优化器更新、自回归生成等),可能采用不同的并行策略。传统的分布式 RL 通常假设模型可在单个 GPU 上训练,或使用数据并行方式,将控制流和计算流合并在同一进程中。这在处理小规模模型时效果良好,但面对大模型,训练需要复杂的多维并行,涉及大量分布式计算,传统方法难以应对。
大模型 RL 本质上是一个二维的 DataFlow 问题:high-level 的控制流(描述 RL 算法的流程)+ low-level 的计算流(描述分布式神经网络计算)。近期开源的 RLHF 框架,如 DeepSpeed-Chat、OpenRLHF采用了统一的多控制器(Multi-Controller)架构。各计算节点独立管理计算和通信,降低了控制调度的开销。然而,控制流和计算流高度耦合,当设计新的 RL 算法,组合相同的计算流和不同的控制流时,需要重写计算流代码,修改所有相关模型,增加了开发难度。与此前框架不同,HybridFlow 采用了混合编程模型,控制流由单控制器(Single-Controller)管理,具有全局视图,实现新的控制流简单快捷,计算流由多控制器(Multi-Controller)负责,保证了计算的高效执行,并且可以在不同的控制流中复用。
veRL:All in RL元年的必修课RL 训练任务可被视为一种数据流(DataFlow),比如 RLHF的DataFlow主要可以分解为以下三个阶段:
- 第一阶段:回复生成(Generation/Rollout):Actor 模型(一个预训练或微调的 LLM)使用一批提示(prompts)自动回归地生成回复(responses),这一步执行 LLM 推理。这一步步骤通常称为 Rollout,也是常常是 RL 中时间占比 80% 以上的热点部分。
- 第二阶段:经验准备(Preparation):使用提示和生成的回复,通过各自模型的单次前向计算,对生成的回复进行评分。这个阶段通常涉及以下模型:
- Critic Model:计算生成回复的值(values)。
- Reference Model:计算生成回复的参考对数概率,它通常是 Actor 模型在 RLHF 之前的版本,用于限制 Actor 模型在训练过程中偏离过远。
- Reward Model:计算生成回复的奖励(rewards)。奖励模型通常是一个基于人类偏好数据进行微调的 LLM,其语言建模头被替换为标量输出头。
- 第三阶段:训练(Training):Actor 模型和 Critic 模型使用前两个阶段产生的数据以及相应的损失函数进行更新。这个过程就是 LLM 训练过程。Actor 模型的训练目标是最大化奖励模型给出的奖励,同时避免偏离参考策略过远 。Critic 模型的训练目标是准确预测状态的Value。
R1 出现后,RL 从 Human Alignment 向 Reasoning 转变,除对齐任务外,还可用于代码生成、数学推理等领域。这些任务有明确 ground truth,可通过评估代码在测试用例中的正确性或验证数学结果准确性确定。奖励模型可用非神经网络的奖励模块替代,如代码生成任务用沙箱环境评估执行结果,数学推理任务用奖励函数验证结果正确性。
将工作流抽象成DataFlow 并不稀奇,比如在线推理服务微服务MLOps pipeline,每个节点都是一个在线的微服务,每个微服务的实例也可以运行在多个 GPU 上,例如 vLLM/SGLang的并行推理服务。RL DataFlow 有几个独特的特点,使得为其特定设计一个 RL 训练变的必要。
- RL 训练是离线任务,这一点比在线任务挑战更小,不需要考虑并发请求、扩缩容、负载均衡等问题。
- RL Actor 的 rollout 和 train共享权重。这也是 RL DataFlow的特色,同一个模型训练又生成,是有状态的。这也是设计 RL 定制训练框架的主要动机,很多重要功能都是围绕这一特点进行设计。
RL 训练框架要解决的问题是:
- 灵活定义 DataFlow
- 将定义出来的 DataFlow 在 GPU 集群上高效执行
一个完美的 RL 训练框架,希望为 DataFlow 的每个节点灵活定义 Placement 和 Parallelism。Placement 表示模型放置在哪些卡上,如下图所示,可以 colocate 所有模型相同 GPU 上,每个模型放在单独设备上,还有混合的方式。Parallelism 则是模型并行方式,有 ZeRO、TP(Tensor Parallel)、PP(Pipeline Parallel) 等策略。在pre-LLM时代,RL 的每个节点在单个 CPU/GPU 上进行。但是 LLM 时代,RL 训练 DataFlow 中每个节点都需要多 GPU 上分布式执行。这导致在 veRL 出现之前,面向 LLM 的 RL 训练框架难以平衡编程灵活性和分布式执行效率的问题。比如,一些实现限制了编程灵活性,每个 DataFlow 节点是一个运行在独立资源上的的并行程序,不同节点通过定制化的 P2P 通信串联起来,构成完整 DataFlow,这导致可编程性很差,比如 NemoAligner。在另一些实现限制了DataFlow 的 Placement 和 Parallelism方式,比如 DeepSpeed-Chat,不同DataFlow 节点在同一个程序中运行在同一份资源上,互相抢占显存,导致效率很低。
single-controller + multi-controller:
- veRL 借助混合控制方式,在一定程度上解决了数据流(DataFlow)灵活定义与高效执行的问题。具体而言,它在不同层级分别采用了单控制器(single -controller)和多控制器(multi - controller)两种模式。其中,单控制器负责控制,多控制器负责计算。具体来说:在模型(DataFlow 的节点)之间(Inter-node level),HybridFlow 采用了single-controller模式。有中央化的控制器负责协调不同的模型在 RLHF 数据流中的执行顺序和数据传输。这个单控制器运行在一个独立的进程中,类似 master。 它保证了 DataFlow 定义的灵活性,而且在模型数量不多的 RLHF 数据流中,控制分发的开销可以忽略不计。veRL 的 single controller 使用 Ray 实现,也是发挥大家常说的 ray 适合做胶水层的作用。在每个模型内部的分布式计算(DataFlow 的一个节点内)中(Intra-node level),HybridFlow 采用了multi-controller模式。这意味着每个计算设备(例如 GPU)都有自己的控制器,独立管理其上的计算任务,就是的 SPMD 方式运行的,也就是复用sglang/vllm/torchdpp 启动推理/训练的方式。
- 灵活 Placement+Parallelism:用户定义了一个 DataFlow,veRL能够高效执行出来,并行保持,veRL 还需要相对自动地帮用户解决如下事情:
- DataFlow node 之间如何传递 tensor:实际的 tensor 数据传输通常发生在计算 GPU 之间,而不通过中央master 节点。因为不同 node 的 Parallelism 方式不同,导致输出 tensor 的 sharding 方式是不同的。比较呆板实现就是在每个 node 算完 gather,然后 node 计算前去 shard。veRL 定义数据依赖和传输协议简化用户的编程:对于节点之间的 tensor 传递,会通过 @register 装饰器将节点的操作与预定义的若干种传输协议关联起来。传输协议定义了如何收集发送节点的输出以及如何分发到接收节点的输入。
- Async DataFlow Execution:通过 Ray 的 future 机制 实现异步数据流执行,使得没有依赖关系的节点可以并发执行。
- 3D-HybridEngine,尽管我们期望能够任意配置 Placement,但对于 Actor 而言,将 Generation 和 Train 进行并置(colocate)的方式最为有效。veRL 最核心的功能是 3D-HybridEngine。
- Auto-Mapping,single-controller 做控制,实现了对 DataFlow编程的灵活性。给定一个模型如何找到适合的Placement+Parallelism 方法?论文中提出了 Auto - Mapping 算法来搜索每个 node 最佳 Parallelism 和 Placement 方式,但我认为在实际应用中仍需依赖基于经验的试验。最后试验结果部分也表明:在较小的 GPU 集群(例如 16 到 64 块 GPU)上,将所有模型共置 (colocate) 在同一组设备上通常能获得最佳性能。在较大的 GPU 集群(例如 96 到 128 块 GPU)上,分割 (split) 策略(例如将 Actor 和 Reference Policy 放在一组 GPU 上,Critic 和 Reward Model 放在另一组 GPU 上)有时会成为最优选择。
工程
分布式RLHF武庙十哲上 未读。 梳理了rhlf 框架