跳转至

VeRL Training Flow

导言

这篇文章聚焦 verl 的训练链路：RayPPOTrainer.fit() 如何组织 rollout、reward、logprob、ref 和 actor update，以及这些阶段如何通过 worker 和 DataProto 串起来。

1. 主问题

  • RayPPOTrainer 的 step 内部到底做了哪些阶段。
  • DataProto 在每个阶段新增、删除或重排哪些字段。
  • actor、ref、critic、reward worker 的职责如何划分。
  • Megatron / FSDP backend 在 shape、显存和通信上有什么差异。

2. RayPPOTrainer Main Loop

for batch_dict in train_dataloader:
    gen_batch_output = actor_rollout_wg.generate_sequences(batch)
    old_log_prob = actor_rollout_wg.compute_log_prob(batch)
    ref_log_prob = actor_rollout_wg.compute_ref_log_prob(batch)
    actor_output = actor_rollout_wg.update_actor(batch)

3. Worker 角色

  • Actor：计算 new logprob，执行 policy update。
  • Rollout：生成 response，维护推理 engine。
  • Reference：计算 ref logprob，提供 KL 约束基准。
  • Reward：计算 reward 或 score。
  • Critic：如果算法需要 value function，则计算 value / return。

4. DataProto 字段演化

阶段 输入字段 输出字段 常见风险
rollout prompt ids / mask response ids / response mask sample 与 group 展开不一致
reward full sequence / response rewards / scores reward 标量与 token loss 对不齐
old logprob full sequence old_logprob prompt token 与 response token 口径混淆
ref logprob full sequence ref_logprob ref model dtype / mask 口径不一致
update actor logprob / reward / advantage loss / metrics micro batch 和 dynamic batch shape 不一致

5. Megatron Backend

  • 初始化 actor / ref module。
  • 构造 pipeline / tensor / data parallel group。
  • 调用 forward-backward 函数。
  • 处理 micro batch、dynamic batch、sequence parallel。

待验证

当前只从现有笔记中知道 use_dynamic_bsz 和 use_fused_kernels 是待验证选项。后续要确认它们的配置路径、默认值、后端限制和对 shape / memory 的影响。

6. FSDP Backend

  • 参数 shard 方式不同。
  • optimizer state 管理不同。
  • activation checkpoint 与通信重叠策略不同。
  • 和 Megatron 的主要差异应放在显存、通信、启动复杂度、调试难度四个维度比较。

7. Loss 路径

  • policy loss
  • KL loss
  • entropy loss
  • loss aggregation
  • response mask

需要重点回答：token-level loss 如何从 sample-level reward / advantage 映射而来。

8. 常见失败模式

  • shape mismatch
  • actor / ref logprob 对不齐
  • response mask 错误
  • KL / grad norm spike
  • dynamic batch 造成单卡峰值异常
  • pipeline bubble 或 rank 间 token 不均

9. 待验证项

  • RayPPOTrainer.fit() 在 release/v0.7.1 的准确代码路径。
  • Megatron worker 与 FSDP worker 的类名、方法名和配置路径。
  • dynamic batch 在训练阶段和推理阶段是否共享口径。
  • SP / CP 在 Megatron backend 中的开关与限制。

评论