Skip to content

大模型 RL 框架的演进与发展趋势 ​

1. 从 SFT 到强化学习:模型训练范式的转变 ​

在 2024 年 OpenAI 发布 O1 系列模型之前,主流的机器学习训练方式主要依赖于有监督微调(Supervised Fine-Tuning, SFT)。该方法通过让模型学习“标准答案”,并根据预测与真实标签之间的损失(loss)来更新模型参数。训练流程相对简单, PyTorch 和 TensorFlow 等深度学习框架也围绕这一范式构建了丰富的训练加速工具。

然而,随着 O1 系列模型的发布,模型训练的重心逐渐从 SFT 向强化学习(Reinforcement Learning, RL)转移。SFT 逐渐被视为训练过程中的“预热”阶段,其作用被弱化为参数初始化或策略引导。取而代之的是, RL 在模型能力提升中扮演了越来越关键的角色。

1.1 RL 算法的演进与多样化 ​

RL 算法本身也在不断迭代与优化。从早期的DPO (Direct Preference Optimization),到经典的PPO (Proximal Policy Optimization),再到近年来涌现出的GRPO、RLOO、Reinforce++、DAPO等新方法, RL 算法在策略更新方式、稳定性、样本效率等方面持续优化。

尽管 DPO 因其简洁性曾一度流行,但随着任务复杂度和模型规模的提升,其局限性逐渐显现,目前在实际工程中已较少被采用。尽管如此,主流 RL 框架的整体结构保持相对一致,核心流程主要包括以下几个阶段:

1.2 RL 训练流程的三大模块 ​

1.2.1 模块一:策略生成(Rollout) ​

对应“学生自己寻找答案”的过程。这是 RL 训练中的推演阶段(Rollout),模型基于当前策略生成响应(action),模拟与环境的交互过程。该阶段是模型推理过程的扩展,通常需要大量采样以获取多样化的行为轨迹。

1.2.2 模块二:奖励评估(Reward Evaluation) ​

对应“给学生答案打分”的过程。传统上,这一阶段依赖于奖励模型(Reward Model),用于评估生成结果的质量。在当前阶段,由于任务复杂度提升,奖励评估的实现方式也趋于多样化:

  • 基于规则的评估(Rule-based):在数学、物理、代码等领域,通过结果与规则的匹配度进行打分。
  • 轻量级奖励模型:训练一个小型模型(如 7B 参数)进行打分,成本可控,且效果良好。

在许多研究项目中,这一模块甚至被简化为 Rollout 的一部分,未被单独重视。然而,随着Agent 行为模拟的兴起,尤其是在商业应用场景(如电商、客服等)中,奖励评估的复杂性显著上升,未来该模块的重要性将不断提升。

1.2.3 模块三:策略更新(Policy Update) ​

对应“学生根据打分来学习”的过程。这是 RL 训练的核心阶段,基于传统训练框架(如 PyTorch、DeepSpeed 等),通过修改损失函数实现策略更新。不同算法(如 PPO、DPO、RLOO 等)在此阶段的实现逻辑有所不同,但整体结构保持一致。

总结 ​

从 SFT 主导的训练范式到 RL 驱动的能力提升,大模型的训练流程正经历深刻的变革。RL 框架的结构虽然保持稳定,但其各模块的功能、实现方式和重要性正在不断演化。

  • Rollout 模块:面临长上下文、异构任务带来的性能挑战;
  • Reward Evaluation 模块:从简单规则向复杂评估演进,未来可能成为 RL 训练中的关键瓶颈;
  • Policy Update 模块:依赖于底层训练框架的性能优化与算法迭代。

随着 Agent 行为模拟、复杂任务建模、多模态交互等方向的发展, RL 框架的设计将更加注重模块间的协同、资源调度的高效性以及算法与工程实现的统一性。

2. RL 训练框架设计与性能优化挑战 ​

当前,主流的强化学习(Reinforcement Learning, RL)训练框架通常被划分为两个核心模块:训练(Training) 和 Rollout (推演)。

在设计一个高效的 RL 训练系统时,开发者将面临一系列关键挑战。以下是我们在技术选型与框架设计过程中总结出的三大核心问题。

2.1 挑战一: Rollout 与训练模块的协同与资源管理 ​

目前, RL 训练普遍采用On-policy策略,这意味着 Rollout 与训练过程必须顺序执行。然而,随着模型规模的持续增长,分布式多卡训练已成为必然趋势。

  • Rollout 阶段:主要为内存密集型任务,尤其在处理长上下文(如 Chain-of-Thought)时,需要维护大量的 KV Cache (Key-Value Cache)。
  • 训练阶段:则属于计算密集型任务,涉及大规模的参数更新和梯度计算。

这两个阶段各自已有大量优化手段(如内存复用、流水线并行等),但如何在统一框架中高效管理这两类异构资源?如何优化两者之间的参数同步机制?这是构建高效 RL 系统的关键挑战之一。

2.2 挑战二:底层训练与推理框架的多样性 ​

当前存在多种主流的训练框架,例如:

  • Megatron-LM
  • DeepSpeed (FSDP)
  • PyTorch FSDP

同时,推理引擎也呈现多样化趋势:

  • vLLM
  • SGLang

不同训练框架与推理引擎的架构差异显著,导致在参数同步、推理调度等环节的实现逻辑差异较大。例如,仅在参数更新部分,不同组合就可能需要完全不同的实现逻辑,这对系统的可维护性与扩展性提出了较高要求。

2.3 挑战三:异构批次执行带来的不确定性 ​

Rollout 任务通常以批次形式执行,但批次内部任务的复杂度可能存在巨大差异。特别是在引入Agent 行为模拟的场景下,这种异构性更加显著,可能导致整体调度效率下降、资源利用率不均衡等问题。

3. 性能优化分析 ​

3.1 初始实现与性能瓶颈 ​

在 RL 训练的早期实现中,整个流程通常分为三个阶段:

  1. 推理阶段(Rollout):模型根据当前策略生成响应。
  2. 评估阶段:通过奖励模型或其他机制对生成结果进行打分。
  3. 训练阶段:基于打分结果更新策略模型。

该流程本质上可以基于 SFT (Supervised Fine-Tuning)框架实现,区别在于需要初始化多个模型实例(如策略模型、奖励模型等)。然而,这种实现方式在实际运行中往往存在显著的性能瓶颈。

3.2 内存优化策略 ​

在大规模模型训练中,显存占用主要包括以下几个部分:

  • 模型参数(Parameters)
  • 梯度(Gradients)
  • 优化器状态(Optimizer States)
  • 激活值(Activations)

以一个 7B 参数模型为例,在 FP32 精度下,仅模型参数和梯度就需要约 28GB 显存,优化器状态则可能额外占用 28GB ×3=84GB,总计高达 112GB。显然,单卡无法承载如此庞大的内存需求。

为此,业界提出了多种分布式训练策略:

  • 数据并行(Data Parallelism, DP):如 DeepSpeed ZeRO-1/2/3,通过 All-Gather 操作动态重建完整参数。
  • 张量并行(Tensor Parallelism, TP)与流水线并行(Pipeline Parallelism, PP):如 Megatron-LM,采用参数切分策略,适用于大规模模型。

根据 NVIDIA 相关论文的研究结论,在千卡以下规模, DP 与 TP/PP 性能相近;但在更大规模下, TP/PP 因避免了 All-Gather 操作的通信开销,性能优势更为明显。

特性数据并行(DP)张量并行(TP)流水线并行(PP)
实现复杂度简单高中等
内存冗余高低低
通信开销中等高低
模型大小限制小大大
计算资源利用率高高中等
调度复杂度低高高
适用场景数据量大、模型较小模型超大、计算密集模型深度大、长序列

这个表格比较了数据并行(DP)、张量并行(TP)和流水线并行(PP)三种并行策略在不同特性上的表现。

3.3 推理速度优化与引擎选型 ​

当前主流推理引擎(如 vLLM 和 SGLang)在 KV Cache 复用、底层算子优化等方面已实现显著性能提升。尽管如此,训练与推理引擎之间的参数同步仍存在一定挑战:

  • 推理引擎生成的输出与训练引擎在精度上存在差异;
  • 当前主流做法是:在 Rollout 阶段使用推理引擎加速生成,训练阶段再由训练引擎重新计算 logits (仅需 prefill 阶段,计算效率高)。

因此,将高性能推理引擎与训练框架进行集成,是提升整体 RL 训练效率的有效路径。但如何高效地实现训练与推理模块的拼接与协同,仍是值得深入研究的问题。

4. 训练框架与推理引擎的整合 ​

4.1 SPMD 和 MPMD 概念解析 ​

在讨论训练框架和推理引擎的集成时,首先需要理解两种并行处理模式:SPMD (Single Program, Multiple Data) 和 MPMD (Multiple Programs, Multiple Data)。这两种模式也可以被描述为单一控制器与多控制器架构。

  • 单一控制器(SPMD):所有工作节点执行相同的程序逻辑,适用于数据量大但模型规模较小的场景。
  • 多控制器(MPMD):每个工作节点可以执行不同的程序,增加了实现复杂度,但无需集中控制,适合特定应用场景。

主流的深度学习训练框架如 DeepSpeed 和 Megatron 都采用了 SPMD 模式,保证所有进程遵循相同的代码逻辑进行运算。然而,对于推理引擎(例如 SGlang 和 vLLM),情况则有所不同。尽管推理引擎(例如 SGLang 和 vLLM)在计算过程中遵循 SPMD 原则,但在决定下一个 token 来源或如何处理 KV 缓存等方面,则不完全适用 SPMD/MPMD 分类。对于这些情况, Google Pathway 等系统提供了更灵活的解决方案。

考虑到上述背景,我们更应关注的是训练框架与推理引擎之间关于训练数据和模型参数的通信机制,而非局限于是否采用单一控制器或多控制器架构。

4.2 SLIME 的具体实现方法 ​

训练框架与推理引擎之间的核心挑战在于训练数据与模型参数的通信机制。为了更好地理解这一点,我们可以通过分析 slime 和 roll 项目来探讨具体实现方案。

SLIME 是一个专注于强化学习扩展的后训练框架,它定义了两个主要组件: RayTrainGroup 用于训练框架, RolloutGroup 用于推理引擎。

4.2.1 数据传输机制 ​

SLIME 通过定义一个中间件类——Buffer,实现了推理引擎与训练模块间的数据传输。所有的数据都会被存储在这个 Buffer 中(甚至可以写入磁盘),并通过 rollout ID 进行指定访问。此外, Buffer 类中的数据处理函数以及 rollout/eval 函数均可以通过命令行参数灵活配置,极大地提高了系统的适应性。

python
self.generate_rollout = load_function(self.args.rollout_function_path)
self.eval_generate_rollout = load_function(self.args.eval_function_path)

这种设计使得应对业务需求时更加灵活高效,尤其是面对各种特殊需求和数据格式时尤为重要。

Rollout 的 generate 函数是通过 Buffer。

python
def async_generate(self, rollout_id, evaluation=False):
     return self.data_buffer.generate.remote(rollout_id, evaluation=evaluation)

获取训练框架所需的数据同样依赖于这个 Buffer:

python
def get_rollout_data(self, rollout_id):
    megatron_utils.process_rollout_data(rollout_id, self.args, self.data_buffer)

同步 rollout 的 buffer 给 actor 的过程如下所示:

python
 def async_init_weight_update_connections(self, rollout):
        """
        Connect rollout engines and actors, e.g. initialize the process group between them
        to update weights after each training stage.
        """
        self.rollout = rollout
        ray.get([actor.set_data_buffer.remote(rollout.data_buffer) for actor in self._actor_handlers])

4.2.2 模型参数同步机制 ​

为了让 rollout 引擎能够在适当的时候正确地同步参数, SLIME 将 actor 的配置信息传递给 rollout。这部分涉及到初始化过程组以便在每个训练阶段之后更新权重。

python
 def async_init_weight_update_connections(self, rollout):
        """
        Connect rollout engines and actors, e.g. initialize the process group between them
        to update weights after each training stage.
        """
        self.rollout = rollout
        ray.get([actor.set_data_buffer.remote(rollout.data_buffer) for actor in self._actor_handlers])
        actor_parallel_configs = ray.get([actor.get_parallel_config.remote() for actor in self._actor_handlers])
        parallel_config = {}
        for rank, config in enumerate(actor_parallel_configs):
            assert config["rank"] == rank and config["world_size"] == len(self._actor_handlers)
            config.pop("rank")
            for key, value in config.items():
                if "size" in key and key:
                    if key not in parallel_config:
                        parallel_config[key] = value
                    else:
                        assert (
                            parallel_config[key] == value
                        ), f"mismatch {key} on rank {rank}: {parallel_config[key]} != {value}"
        parallel_config["actors"] = actor_parallel_configs
        ray.get(rollout.async_set_parallel_config(parallel_config))

        return [
            actor.connect_rollout_engines.remote(
                rollout.rollout_engines,
                rollout.rollout_engine_lock,
            )
            for actor in self._actor_handlers
        ]

上述过程不仅包括数据缓冲区的同步,还涵盖了 actor 间并行配置的协调,保证了参数更新的一致性和准确性。

4.3 ROLL 的具体实现方法 ​

ROLL 通过集群(Cluster)的方式定义了多个角色,每个角色负责不同的任务。这种设计方式与算法层面的认知较为一致,因为从算法角度来看,训练框架和推理引擎之间的差异并不明显,而使用集群封装则很好地隐藏了这些复杂性。

python
self.actor_train = Cluster(
    name=self.pipeline_config.actor_train.name,
    worker_cls=self.pipeline_config.actor_train.worker_cls,
    resource_manager=self.resource_manager,
    worker_config=self.pipeline_config.actor_train,
)
self.actor_infer = Cluster(
    name=self.pipeline_config.actor_infer.name,
    worker_cls=self.pipeline_config.actor_infer.worker_cls,
    resource_manager=self.resource_manager,
    worker_config=self.pipeline_config.actor_infer,
)
self.reference = Cluster(
    name=self.pipeline_config.reference.name,
    worker_cls=self.pipeline_config.reference.worker_cls,
    resource_manager=self.resource_manager,
    worker_config=self.pipeline_config.reference,
)
if self.pipeline_config.adv_estimator == "gae":
    self.critic = Cluster(
        name=self.pipeline_config.critic.name,
        worker_cls=self.pipeline_config.critic.worker_cls,
        resource_manager=self.resource_manager,
        worker_config=self.pipeline_config.critic,
    )

4.3.1 数据传输机制 ​

类似于 Megatron, ROLL 允许按照领域(domain)分开采样,并在pipeline.py文件中进行配置。这使得如果用户不想编写数据生成器, ROLL 提供了一种更为便捷的解决方案。特别是对于奖励(reward)模型,理想的状况是有一个统一的模型,但由于训练难度大,目前更倾向于针对不同领域使用不同的奖励模型,并最终进行聚合处理。ROLL 支持对不同领域、批次以及查询进行自定义配置,以适应多样的应用场景。

4.3.2 模型参数同步机制 ​

ROLL 中的模型更新逻辑结合了点对点通信和集体通信两种方式:

python
def model_update(self, tgt_workers, broadcast_tgt_devices, p2p_tgt_devices):
    # 更新逻辑代码...
  • 点对点通信:用于同一设备上的参数更新,直接通过 worker 的 node_rank 和 gpu_rank 来判断是否在同一设备上,从而进行高效的数据交换。
  • 集体通信:通过广播参数到目标集群,只在主进程(rank 0)执行广播操作,适用于跨设备间的参数同步。

这两种通信策略分别对应于 colocate 和非 colocate 场景,确保了参数同步的灵活性和效率。

4.3.3 跨机器部署时的考量 ​

当所有组件位于同一台机器上时,硬编码实现参数同步相对简单,但当涉及到跨机器部署时,情况变得更加复杂。此时,不仅需要考虑如何有效地管理网络通信带来的延迟和带宽限制,还需要优化分布式环境下的资源分配和负载均衡。此外,单控制器(single controller)模式下,控制器的压力会随着集群规模的扩大而增加,尤其是在处理多媒体数据时,可能需要特别注意性能瓶颈的问题。因此,在跨机器部署的情况下,选择合适的通信策略和优化控制器的工作负载变得尤为重要。不过,从 SLIME 和 ROLL 的设计来看,参数同步的核心在于通知 GPU 进行同步操作,中间的通信过程不依赖于控制器,这为跨机器部署提供了一定的便利性和灵活性。

4.4 Colocation 与 Ray 的应用 ​

将 Actor、Ref、Reward、Critic 等模型放置在同一张 GPU 卡上被称为colocation。然而,正如前文所述,随着模型规模的增大(例如 7B 模型已难以在单张卡上训练),预计下半年会出现多个超过 1000B 参数量级的模型。这使得并行计算带来的开销变得极其显著。当前, Reward 模型普遍较小, 7-30B 的规模即可满足需求,因此分开部署往往更具性价比。

为了应对这种复杂性,项目中引入了 Ray ——一个支持分布式计算的强大框架,它能够帮助开发者减轻底层逻辑管理的负担。有关基于 Ray 的分布式训练流程和 Ray 分布式计算框架的详细介绍,请参阅以下文章:

接下来,我们将比较 SLIME、veRL、ROLL 和 OpenRLHF 四个框架在 colocation 与非 colocation 实现上的差异。

4.4.1 SLIME ​

SLIME 仅定义了两个主要 worker: RayTrainGroup 用于训练, RolloutGroup 用于推理。对于 colocate,训练和推理可以分开部署;而在非 colocate 的情况下,则需要处理分布式通信以同步参数。这种设计抽象层次高,易于理解,并且能够很好地适应训练和推理的不同需求。只需在配置中指定是否 colocate,即可自动在所有关键环节执行相应操作。

4.4.2 ROLL ​

对于非 colocate 场景, ROLL 允许细粒度地指定不同 worker (例如 actor、critic、reward 等)部署在不同的显卡上,甚至可以根据轮次进行配置。若不手动指定, Ray 会自动完成部署。鉴于 RL 任务对资源的高消耗,细粒度的 GPU 资源配置有助于提高资源利用效率,但这同时也对算法侧的资源调度能力提出了更高要求。显然,使用 Ray 来管理这些复杂性更为合适。

4.4.3 veRL ​

veRL 采用了一种独特的方法来实现 colocation 和非 colocation 部署。在非 colocation 模式下,每个 worker (如 actor、critic、reward 等)作为一个独立进程运行,依靠 Ray 来进行调度。而在 colocation 模式下,多个角色共享同一个 Ray actor 实例,在同一进程中实例化多个 worker 类。通过 create_colocated_worker_cls 或 create_colocated_worker_cls_fused 方法动态生成一个多角色类(例如 WorkerDict/FusedWorker),该类内部持有多个 worker 实例。外部可通过统一接口调用不同角色 worker 的方法,内部则自动分发到对应的 worker 实例。这种方式使得同进程内的多角色共存成为可能,并在某些场景下能大幅提高性能,比如减少跨进程通信带来的延迟和内存碎片问题。

4.4.4 OpenRLHF ​

OpenRLHF 提供了灵活的混合部署选项,既支持 vLLM 引擎、Actor、Reference、Reward 和 Critic 模型节点的共置部署,也支持部分混合部署或完全分离部署,以适应异步训练的需求。这种灵活性使其能够应对多样化的应用场景,但也意味着更复杂的管理和优化需求。

结论 ​

综上所述,在非 colocation 情况下, Ray 确实可以帮助我们更加轻松地管理资源,尤其是在处理复杂的 Agent 和多轮交互场景时。然而,根据运维团队的反馈, Ray 的设计理念与现有的 Kubernetes 云原生生产环境存在一定的冲突,导致在实际生产环境中部署时管理成本较高。不过, Ray 团队也在针对这些问题进行优化,例如使 Ray 可以直接通过 NCCL 传输 tensor 数据,从而绕过对象存储,提高效率。未来,我们可以期待更多来自 Ray 的更新和改进。

4.5 不同训练框架与推理引擎的集成 ​

在将不同的训练框架和推理引擎进行集成时,可能会遇到参数转换的问题。例如,如果 vLLM 使用 4 维张量并行(TP),而 DeepSpeed 分布在 8 个 GPU 上,则需要进行适当的参数转换以确保数据传输的一致性。Megatron-LM 也有类似的需求。当存在多个训练框架和推理引擎时,适配的工作量会成倍增加,这可能导致配置错误和性能问题。

4.6 代码解耦设计 ​

以 Slime 为例,其架构分为三层:顶层 RolloutGroup 负责管理推理引擎的整体流程;中层 RolloutRayActor 处理具体的推理请求;底层 SglangEngine 实现具体的推理逻辑。这种分层设计使得替换后端推理引擎变得简单,只需更改底层实现即可,无需修改上层控制逻辑。同样,训练框架也采用了类似的分层结构,保证了系统的灵活性和可维护性。

5. 关于 Agentic RL ​

目前, ROLL、veRL 和 OpenRLHF 等框架对 Agentic RL 提供了良好的支持。尽管这样做可能增加了代码复杂度,但随着技术成熟,预计会有更清晰的设计出现。未来, Agentic RL 有望成为主流,现有的 RL 方法将成为其中的一部分。

6. 框架选择建议 ​

6.1 框架难点分析 ​

快速发展的技术环境意味着旧框架容易过时,因此保持框架简洁和高维护性是关键。新框架由于没有历史负担,可以更容易地适应新技术趋势。

6.2 推荐框架 ​

  • OpenRLHF:一个高性能的开源 RLHF 框架,集成了 Ray、vLLM、ZeRO-3 和 Hugging Face Transformers。
  • slime:新推出的框架,代码简洁,适合想要尝试大胆框架修改的研究者。
  • ROLL:强调数据处理和异步操作的支持,特别适用于深入探索 Agentic RL 的团队。
  • veRL:稳定且优化良好,适合大规模集群部署,尤其适合资源丰富的团队。

根据团队的具体需求和技术背景,可以选择最适合的框架来开展工作。对于有特定需求或希望快速扩展的团队, veRL 可能是更好的选择,因为它已经被多个大厂验证过。而对于追求技术创新和敏捷开发的团队, SLIME 或 ROLL 可能更具吸引力。

6.3 结尾 ​

在过去半年中,我们深入探讨了 RL 训练框架、Agent 框架以及推理引擎框架。总体而言,代码量方面, Agent 框架最为庞大,其次是推理引擎和 RL 训练框架;而在代码难度上,推理引擎居首,随后是 RL 训练框架和 Agent 框架。值得注意的是,如果排除推理引擎底层算子的复杂性, RL 训练框架的挑战主要在于集成各种系统和技术,这要求框架开发者对多种技术和业务逻辑有深刻的理解。

开源框架如 veRL、SLIME、ROLL 及 OpenRLHF 各具特色,展现了作者们的追求与坚持,并且社区活跃度高。可以说,在开源 RL 框架领域,中国在技术实力和认知深度方面处于世界领先位置。虽然算法人才间的差异不大,但在硬件资源(如显卡)方面仍存在一定的差距。

Maintained by Robin