Prhub

#7188 [trainer] feat: support decouple ppo for v1 separate_async

原始 PR 作者 zpltys 合并时间 2026-07-29 21:20 文件变更 3 提交数 2 评论 6 代码增减 +44 / -4

执行摘要

V1 separate_async 支持解耦 PPO

根据 PR 标题 'feat: support decouple ppo for v1 separate_async',旨在实现 Decoupled PPO(https://arxiv.org/abs/2505.24298)算法。Decoupled PPO 要求当 parameter_sync_step > 1 时,每个 mini-batch 的 old_log_probs 必须基于同一个旧策略计算,而不是使用被当前 mini-batch 更新后的权重。原有实现通过 TODO 注释标记了需求('# TODO: Support Decoupled PPO'),且强制开启了 bypass_mode 以避免旧策略不一致问题,但这也失去了 PPO 的 clipping 等特性。

值得精读。该 PR 展示了如何在现有异步训练框架中支持高级算法(Decoupled PPO),特别是通过 CPU 权重快照实现多步旧策略一致性的设计模式。评审过程展现了好的代码质量实践——将状态追踪移入基类、统一索引约定。建议关注后续对 DetachActorWorker 的测试覆盖和稳定性。

讨论亮点

Reviewer wuxibin89 提出了两条关键建议:

  1. 将 local_trigger_step 移至基类以避免子类重写 step 方法。评审意见指出直接在子类中重写 step 会增加代码冗余和潜在的不一致风险,建议在基类 step 循环中注入 trigger_idx。
    结论:已采纳,通过两次 commit 实现:第一个 commit 在子类中实现了 step 重写,第二个 commit 移除了重写,改为在基类中添加 local_trigger_step 并在循环中赋值。

  2. 对索引基数提出质疑:'Why use 1-indexed, 0-indexed is more intuitive.' 评审者认为 0-based 索引更符合直觉。
    结论:从最终代码看,采用了 0-based 索引(for trigger_idx in range(...)),评审意见被采纳。

此外,CLAssitant 检查显示有 1 位 committer(zhangshuai122)未签署 CLA,但不影响代码层面的逻辑。

实现拆解

  1. 基类注入 local_trigger_step(trainer_base.py):在 PPOTrainer.init 中初始化 self.local_trigger_step = 0;在 step() 方法循环中将 for _ in range(...) 改为 for trigger_idx in range(...), 并将 trigger_idx 赋值给 self.local_trigger_step,供子类追踪当前 mini-batch 在 parameter_sync_step 循环中的位置。

  2. 替换 worker 为 DetachActorWorker(trainer_separate_async.py):新增 _init_resource_pool_mgr 方法,调用父类初始化后,检查 role_worker_mapping 中是否存在 ActorRolloutRef 或 ActorRollout,若存在则将对应的 worker 类替换为从 experimental.separation 导入的 DetachActorWorker。DetachActorWorker 提供 save_model_to_cpu / restore_model_from_cpu 能力,这是解耦 PPO 支持的核心前
    提。

  3. 版本感知的 old_log_prob 计算(trainer_separate_async.py):新增 _compute_old_log_prob 方法。如果配置中 rollout_correction.bypass_mode 为 True,则直接委托父类实现(旧行为);否则进入 Decoupled 模式:

    • local_trigger_step == 0: 当前权重即为 π_old,先 save 到 CPU slot 0,然后直接计算。
    • local_trigger_step >= 1: 先将当前权重 save 到对应 slot,再从 slot 0 restore 回 π_old,计算完毕后 restore 回当前权重并清理临时 slot。
  4. 移除强制 bypass 并清理 TODO(trainer_separate_async.py):在 init 中删除了原来强制设置 self.config.algorithm.rollout_correction.bypass_mode = True 的代码和关联的 TODO 注释,使配置不再被硬编码覆盖。

  5. 修复阈值校验 bug(rollout_corr_helper.py):在 _parse_rollout_rs_thresholds 函数中,当解析 k1 类阈值时,增加了 if lower > upper: raise ValueError(...) 检查,防止用户传入无效的上下界顺序。

文件 模块 状态 重要度
verl/trainer/ppo/v1/trainer_separate_async.py 训练器 modified 7.58
verl/trainer/ppo/v1/trainer_base.py 训练器 modified 5.13
verl/trainer/ppo/rollout_corr_helper.py 训练器 modified 4.23

关键符号

_compute_old_log_prob _init_resource_pool_mgr step _parse_rollout_rs_thresholds

关键源码片段

verl/trainer/ppo/v1/trainer_separate_async.py core-logic

核心变更文件:实现解耦 PPO 的主要逻辑——替换 worker、移除强制 bypass、新增版本感知的 old_log_prob 计算,包含 38 行新增和 3 行删除,是最关键的改动。

def _compute_old_log_prob(self, batch: KVBatchMeta, metrics: dict) -> KVBatchMeta:
    """Version-aware old_log_probs computation for Decoupled PPO.    In bypass mode, delegates to the base class (copies rollout_log_probs directly).
    In Decoupled mode, uses save_model_to_cpu / restore_model_from_cpu to ensure
    all mini-batches within a parameter_sync_step cycle use the same stable π_old.
    """
    rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
    bypass_recomputing_logprobs = (
        rollout_corr_config and rollout_corr_config.get("bypass_mode", False)
    )
    if bypass_recomputing_logprobs:
        # 旧行为:直接拷贝 rollout_log_probs,不重新计算
        return super()._compute_old_log_prob(batch, metrics)
​
    # Decoupled 模式:确保每个 mini-batch 使用相同的 π_old
    if self.local_trigger_step == 0:
        # 当前权重即为 π_old,保存到 CPU slot 0 供后续 step 恢复
        self.actor_rollout_wg.save_model_to_cpu(0)
        return super()._compute_old_log_prob(batch, metrics)
    else:
        # 保存当前权重到临时 slot,恢复 π_old(slot 0)
        self.actor_rollout_wg.save_model_to_cpu(self.local_trigger_step)
        self.actor_rollout_wg.restore_model_from_cpu(0)
        result = super()._compute_old_log_prob(batch, metrics)
        # 恢复当前权重并清理临时 slot
        self.actor_rollout_wg.restore_model_from_cpu(self.local_trigger_step)
        self.actor_rollout_wg.clear_cpu_model(self.local_trigger_step)
        return result

评论区精华

local_trigger_step 的位置 设计

wuxibin89 建议将 local_trigger_step 移至基类以避免子类重写 step 方法,并保持一致性。

结论:已采纳:第二个 commit 移除了子类 step 重写,改为在基类 PPOTrainer 中注入 trigger_idx。 · 已解决

索引基数(1-indexed vs 0-indexed) style

wuxibin89 提问为何使用 1-indexed,认为 0-indexed 更直观。

结论:最终实现采用 0-indexed(`for trigger_idx in range(...)`),评审意见被采纳。 · 已解决

风险与影响

  1. CPU 权重保存/恢复的资源开销:每次 mini-batch 都需要 save/restore 模型权重到 CPU,会引入额外显存拷贝和 CPU 内存占用。当 parameter_sync_step > 1 时,save_to_cpu 会为每个 step 占用一个 CPU slot,可能增加内存压力。
  2. 正确性风险:_compute_old_log_prob 中 save/restore 时序复杂,若出现异常(如 OOM 或进程 crash),可能导致权重状态不一致。建议增加异常处理确保 restore 在异常时也能执行。
  3. DetachActorWorker 依赖:功能依赖 experimental 模块下的 DetachActorWorker,该模块可能尚未充分测试或 API 不稳定,存在兼容性风险。
  4. 现有行为变更:删除了强制 bypass_mode = True,如果用户未显式配置 bypass_mode,默认行为可能从'enabled'变为'disabled',影响已有训练脚本的日志概率计算逻辑。建议在 release notes 中说明此行为变更。
  5. 缺少测试覆盖:本次 PR 未新增测试文件,仅修改了源代码。解耦 PPO 涉及多条复杂控制流路径(bypass/decoupled,trigger_step 0/>0),建议至少添加单元测试验证 _compute_old_log_prob 各分支。

影响范围:仅影响使用 V1 separate_async 训练器的用户,且只有在关闭 rollout_correction.bypass_mode 时才会启用新的 Decoupled PPO 路径。基类 PPOTrainer 中 local_trigger_step 的修改对同步和其他异步模式无副作用。
影响程度:中等。对于需要 parameter_sync_step > 1 的用户,该特性是使 PPO 训练正确的关键修复;对于其他用户,代码透明。rollout_corr_helper 的阈值校验修复是一个安全补丁,可防止用户配置错误。

缺少测试覆盖 核心路径变更 依赖 experimental/DetachActorWorker

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论