执行摘要
- 一句话:V1 separate_async 支持解耦 PPO
- 推荐动作:值得精读。该 PR 展示了如何在现有异步训练框架中支持高级算法(Decoupled PPO),特别是通过 CPU 权重快照实现多步旧策略一致性的设计模式。评审过程展现了好的代码质量实践——将状态追踪移入基类、统一索引约定。建议关注后续对 DetachActorWorker 的测试覆盖和稳定性。
功能与动机
根据 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 等特性。
实现拆解
-
基类注入 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 循环中的位置。
-
替换 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 支持的核心前
提。
-
版本感知的 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。
-
移除强制 bypass 并清理 TODO(trainer_separate_async.py):在 init 中删除了原来强制设置 self.config.algorithm.rollout_correction.bypass_mode = True 的代码和关联的 TODO 注释,使配置不再被硬编码覆盖。
-
修复阈值校验 bug(rollout_corr_helper.py):在 _parse_rollout_rs_thresholds 函数中,当解析 k1 类阈值时,增加了 if lower > upper: raise ValueError(...) 检查,防止用户传入无效的上下界顺序。
关键文件:
verl/trainer/ppo/v1/trainer_separate_async.py(模块 训练器;类别 source;类型 core-logic;符号 _init_resource_pool_mgr, _compute_old_log_prob): 核心变更文件:实现解耦 PPO 的主要逻辑——替换 worker、移除强制 bypass、新增版本感知的 old_log_prob 计算,包含 38 行新增和 3 行删除,是最关键的改动。
verl/trainer/ppo/v1/trainer_base.py(模块 训练器;类别 source;类型 core-logic): 基类改动:添加 local_trigger_step 属性并在 step 循环中赋值,为解耦 PPO 提供位置追踪基础设施,是协作关键。
verl/trainer/ppo/rollout_corr_helper.py(模块 训练器;类别 source;类型 core-logic): 修复阈值解析的 bug:增加 lower <= upper 校验,防止无效配置,属于防御性编程。
关键符号:_compute_old_log_prob, _init_resource_pool_mgr, step, _parse_rollout_rs_thresholds
关键源码片段
verl/trainer/ppo/v1/trainer_separate_async.py
核心变更文件:实现解耦 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
评论区精华
Reviewer wuxibin89 提出了两条关键建议:
-
将 local_trigger_step 移至基类以避免子类重写 step 方法。评审意见指出直接在子类中重写 step 会增加代码冗余和潜在的不一致风险,建议在基类 step 循环中注入 trigger_idx。
结论:已采纳,通过两次 commit 实现:第一个 commit 在子类中实现了 step 重写,第二个 commit 移除了重写,改为在基类中添加 local_trigger_step 并在循环中赋值。
-
对索引基数提出质疑:'Why use 1-indexed, 0-indexed is more intuitive.' 评审者认为 0-based 索引更符合直觉。
结论:从最终代码看,采用了 0-based 索引(for trigger_idx in range(...)),评审意见被采纳。
此外,CLAssitant 检查显示有 1 位 committer(zhangshuai122)未签署 CLA,但不影响代码层面的逻辑。
- local_trigger_step 的位置 (design): 已采纳:第二个 commit 移除了子类 step 重写,改为在基类 PPOTrainer 中注入 trigger_idx。
- 索引基数(1-indexed vs 0-indexed) (style): 最终实现采用 0-indexed(
for trigger_idx in range(...)),评审意见被采纳。
风险与影响
关联脉络
参与讨论