[YanhuiDua]

Stack 关系

本 PR 属于 Stack #2189，依赖顺序从底到顶为：

层	PR	职责
AgentLoop 数据准备	#2102	统一 RolloutState 的全序列字段和 labels 监督约定
训练数据打包	#2185	controller 生成索引计划，worker 转换并物化训练 pack
Value head	#2188	value-head 模型、HF 权重映射和 critic loss；对应原 #2099
PPO	#2186	在共享样本和 pack 计划上接入 critic worker、GAE 与 actor 训练

本层基于 value-head #2188，是当前栈顶。 这是原 fork PR #2132 在主仓库中的 stack 版本。原 PR 的 packing 依赖 #2055 由 #2185 提供，value-head 依赖 #2099 由 #2188 提供；本层仅保留其上的 PPO、critic worker、GAE 和调度改动。原 PR 保留用于查阅 review 讨论。

概述

本 PR 在索引式 packing 上接入 PPO。actor Ray worker 承载一个独立的 critic 对象；两者共享进程、训练样本缓存和 pack 计划，分别持有自己的模型与优化器。controller 收集 critic value，在原始样本顺序上计算 GAE，再将 returns 和 advantages 分发给 critic 与 actor。

Value-head 依赖

value-head 模型、HF key 映射、CriticLossConfig、fp32_head 配置适配及 critic 模型加载测试属于下层 #2188。本层使用这些能力接入完整 PPO 训练，不再重复包含下层的实现 diff。

Controller 训练流程

rollout groups（已通过训练入口检查）
  → _prepare_rollout_items：保留原始样本与 reward/session 映射
  → _build_pack_plan：生成与训练目标无关的 dp_plans
  → load_train_batch：各 worker 获取并缓存本 DP rank 的 TrainSample 和本地 plan
  → 切换到 critic
  → critic.collect_values：value forward，返回 TokenTensorBatch
  → controller：按 global_indices 恢复原样本顺序，计算 GAE advantages / returns
  → critic.fit：按同一份样本映射分发 returns，训练 value model
  → 切换到 actor
  → actor.fit：分发对应样本的 token advantages，训练 policy
  → 汇总 actor、critic、reward、advantage 和 packing 指标

未配置 critic 且使用组级 estimator 时，继续按下层的 reward/group 映射计算标量 advantage。GAE 需要 critic；配置与角色不匹配时直接报错。

共享样本与 packing 契约

对象	职责
TrainSample	缓存已 shift 的 input/label/logprob、seq_ctx 和 teacher targets；后续填充 values、returns、advantages
_PackPlan.dp_plans	保存每个 DP rank 的全局样本索引和本地 pack 计划，不保存 advantage 或 returns
TokenTensorBatch	传输连续 token 张量及样本边界，支持拆回逐样本的 value/advantage/return
actor worker	获取并缓存样本，提供角色切换；根据标量或 token advantage 准备原始 loss inputs，最终 pack 后构造 loss context
attached critic worker	复用缓存中的样本和计划，独立执行 value forward、critic loss 和优化器更新

• 同一 DP rank 的 replicas 共享 rollout ObjectRef 和相同的样本映射；critic 与 actor 后续复用 worker 内缓存。
• value 收集结果通过 global_indices 对应回原始样本，GAE 的结果再按该映射分发，避免 packing 重排改变目标归属。
• actor 继续使用 [Refactor] Add index-based data packing and worker-side materialization #2185 的 CPU 原始 loss inputs、最终 pack padding 和 loss-context 构造流程。
• 保留下层的训练入口校验、greedy candidate-buffer 配置、MRoPE global-amax 续写、routed-expert packed 表示和有效监督 token 的 advantage 统计。

GAE 与优化调度

• 支持 gamma、lambda 和 actor advantage normalization。
• GAE 根据 labels 区分有效 action 和 observation；轨迹边界不会互相 bootstrap。
• 支持按 rollout 或 session 放置终局 reward；归一化 actor advantage 时保留 critic return 的数值。
• actor/critic 角色切换分别管理模型资源；critic 支持多 pass、optimizer-step 分组和学习率调度。

LPT packing

新增 pack_strategy="lpt"：按长度从大到小，把样本放入当前负载最低且能容纳该样本的 pack。容量不足时按 DP size 增加 pack 数并重新分配，随后完成 DP 分发和 optimizer-step 分组。既有 legacy、greedy、balance、native 策略继续保留。

配置与文档

• examples/v1/config/rl_ppo_gsm8k.py：PPO GSM8K 示例。
• recipe/ppo/config/rl_ppo_math.py、rl_ppo_qwen35_35b_a3b.py：PPO 训练配置。
• 中英文 loss 文档及中文 PPO 教程，说明 value head、critic 和 GAE 的使用方式。
• 配置中的 FSDP head 精度选项同步使用 fp32_head。

测试与验证

• test_critic_sft.py 的 value-head 和 critic-loss 覆盖位于下层 [feat] Support value head and critic loss #2188。
• test_ppo_role_wiring.py：角色注册/切换、GAE reward/action/trajectory/session 语义、advantage normalization 和调度。
• test_pack.py：LPT 策略及 packing 不变量。
• 转换测试使用共享 TrainSample 接口，并保留下层的多字段 pack 拆回一致性覆盖。

PPO 原始提交已重放到最新 packing 层，冲突和调用关系经过源码核对；本层当前 head 为 31fe2e5f。本次未运行单测、pre-commit 或 PPO 训练，仍保持 Draft。

训练启动脚本

#!/usr/bin/env bash
# 通用 PPO (Math) 训练启动配置脚本
#
# 使用方法:
#   1. 填写下方 "必填" 部分的路径
#   2. source 本脚本或直接 bash 执行

# ============================================================
# 必填：请填写你自己的模型 / 数据集路径
# ============================================================

# 模型路径（本地模型目录或 HuggingFace 名称），例如 /path/to/Qwen2.5-Math-1.5B
export MODEL_PATH="${MODEL_PATH:?请设置 MODEL_PATH，例如: export MODEL_PATH=/path/to/your/model}"

# 训练数据集配置文件路径，例如 /path/to/configs/train_dapo.json
export DATA_PATH="${DATA_PATH:?请设置 DATA_PATH，例如: export DATA_PATH=/path/to/your/train_data.json}"

# 评测数据集配置文件路径，例如 /path/to/configs/eval_benchmarks.json
export EVAL_DATA_PATH="${EVAL_DATA_PATH:?请设置 EVAL_DATA_PATH，例如: export EVAL_DATA_PATH=/path/to/your/eval_benchmarks.json}"

# ============================================================
# 必填：评测服务地址（内网评测服务，host:port 逗号分隔）
# ============================================================
export COMPASS_EVAL_HOSTS="${COMPASS_EVAL_HOSTS:?请设置 COMPASS_EVAL_HOSTS（评测服务地址，host:port 逗号分隔）}"

# ============================================================
# 可选配置
# ============================================================

# XTuner 代码根目录，默认取当前环境变量，未设置则报错
export XTUNER_PATH="${XTUNER_PATH:?请设置 XTUNER_PATH（xtuner 代码根目录）}"
cd "$XTUNER_PATH" || exit 1

# 训练输出目录（可用环境变量覆盖）
export WORK_DIR="${WORK_DIR:-work_dirs/exp_final/ppo_colocated}"
rm -rf "$WORK_DIR"

# 长跑轮次：不设置 TOTAL_TRAIN_STEPS 时使用配置文件默认值（128）

# XTuner 代码根目录，默认取当前环境变量，未设置则报错
export XTUNER_PATH="${XTUNER_PATH:?请设置 XTUNER_PATH（xtuner 代码根目录）}"
cd "$XTUNER_PATH" || exit 1

# 训练输出目录（可用环境变量覆盖）
export WORK_DIR="${WORK_DIR:-work_dirs/exp_final/ppo_colocated}"
rm -rf "$WORK_DIR"

# 长跑轮次：不设置 TOTAL_TRAIN_STEPS 时使用配置文件默认值（128）
# export TOTAL_TRAIN_STEPS=

export COMPARE_PACK_MISMATCH=0
export PYTHONUNBUFFERED=1

# ============================================================
# 启动训练
# ============================================================
bash examples/v1/scripts/run_rl.sh recipe/ppo/config/rl_ppo_math.py "lmdeploy" "$MODEL_PATH" "$DATA_PATH" "$EVAL_DATA_PATH"

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 在索引式 packing 之上接入 PPO：actor Ray worker 内挂载独立 critic（共享进程、样本缓存与 pack 计划），controller 收集 token value、在原样本顺序上计算 GAE，再分发 returns / advantages，并新增 LPT packing 策略与 LR schedule 辅助。整体数据流设计清晰，但 role 挂载采用 stringly-typed 反射接口、critic 大量依赖 host 私有方法、critic 核心路径完全没有单测，另有一处与 PPO 无关的 distillation padding 回归。

ProduceBatchResult impact: not affected
RoutedExperts impact: load_train_batch 的 worker 侧缓存会在 critic 阶段失败时继续持有 seq_ctx.rollout_routed_experts object refs（transfer 之后的泄漏）
Ray concurrency impact: 新增 attach / call_attached / switch_role_modules / load_train_batch 均在默认 concurrency group 的同步 actor 上，串行行为与原先单个 fit 一致，无新增 starvation

Main Flowchart after this PR

flowchart TD
    A[_prepare_rollout_items] --> B[_build_pack_plan]
    B --> C[ray.put per-DP RolloutState]
    C --> D["load_train_batch<br/>缓存 samples/plan/batch_attr"]:::problem
    D --> E{has_critic?}
    E -- yes --> F["switch_role_modules('critic')"]:::changed
    F --> G["call_attached('critic','collect_values')"]:::problem
    G --> H["_compute_advantages / _compute_ppo_targets"]:::problem
    E -- no --> I[_prepare_advantages]
    H --> J["call_attached('critic','fit', token_returns)"]:::changed
    I --> K["switch_role_modules('actor')"]:::changed
    J --> K
    K --> L["worker.fit(advantages)<br/>finally 释放缓存"]:::changed
    L --> M[_build_trainer_log_info + _critic_log_scalars]

    classDef changed fill:#dbeafe,stroke:#2563eb
    classDef problem fill:#fee2e2,stroke:#dc2626

Loading

核心原理实现与单测

• GAE 数学语义（终局 reward 落在最后一个可控 action、轨迹边界不互相 bootstrap、observation 跳过、session 拼接、population 归一化不改写 critic return）通过 GAEAdvantageConfig(...).build().compute_gae 公开 API 覆盖，路径真实。
• LPT packing 的可行性判定、bin 增长、optimizer-step 分配与 shuffle 均有单测，并与 baseline reference plan 对齐；num_bins 始终为 DP 倍数，保证各 DP rank pack 数相同、collective 对称。
• TokenTensorBatch 的 cat/split 往返经由转换测试覆盖。
• 缺口：critic 的 collect_values / _materialize / _pack_one / fit / _critic_updates / save / resume，以及 controller 的 _compute_ppo_targets（value → 全局索引回填 → GAE → returns 分发）与 _critic_log_scalars 全部无测试，而这正是本 PR 的核心正确性主张。

抽象与信息隐藏评估

• Warning xtuner/v1/rl/trainer/worker.py:284、worker.py:497-540 role 挂载使用 critic_cfg: Any 与 call_attached(name, method, *args) 字符串反射，绕过 pydantic 校验与 strict mypy，且只有一个 Adapter，属于尚未成立的 Seam。
• Warning xtuner/v1/rl/trainer/critic_worker.py:107-245 CriticWorker 依赖 host 六个私有成员并重写了一套与 _pack_one_batch 平行的 padding 逻辑，host 的 Interface 几乎等于其 Implementation，padding/SP 规则失去 Locality。
• Warning xtuner/v1/rl/trainer/controller.py:237、controller.py:544-556 scalar advantage 规则在 _prepare_advantages 与 _compute_advantages 的非 GAE 分支中重复，后者在 fit 中不可达，_build_dispatches 也仅被测试维持存活。
• Warning xtuner/v1/rl/trainer/controller.py:131-139 _invoke / _resolve 的 .remote-或-直调双分派只为让测试传入普通对象，生产分派逻辑被 test double 反向塑形。

单测建议

• Warning critic 与 controller PPO 目标分发路径零覆盖；_compute_ppo_targets 是纯函数（接收 rollout_items / pack_plan / value_batches / data_replicate_size），无需 Ray 即可对 value→全局索引→GAE→returns 的映射补回归测试。

其他 Issues

• Critical xtuner/v1/rl/trainer/worker.py:1148-1153 删除 uses_sampled_token_targets 分支后，loss_mode="k1" 下任何需要 padding 的 pack 都会命中 assert top_k is not None 而失败。
• Warning xtuner/v1/rl/trainer/worker.py:897-913 训练批缓存仅在 actor fit 的 finally 释放，critic 阶段抛错时样本与 routed-experts refs 会一直驻留到下一次成功加载。
• Warning xtuner/v1/rl/trainer/critic_worker.py:216-218 collect_values 以 sequence_parallel=False 做全长 forward，sp_size>1 时同一 DP 组内各 SP rank 重复同样计算且只有 replica 0 的结果被读取。
• Warning xtuner/v1/datasets/rl_tokenize_fn/text_tokenize_fn.py:79-81 .get() 改为硬 KeyError，未映射的 data_source 由静默 None 变为中断整个 tokenization，属无测试的行为变更。
• Warning xtuner/v1/rl/weight_update/transport.py:112、transport.py:357 与 PPO 无关的权重同步热路径改动（移除 per-group empty_cache、放宽 barrier 条件）无 before/after benchmark，连同新增 judger/math_rule.py 与 model/base.py value-head 初始化改动，违反“一个 PR 一个逻辑变更”。

Verdict

REQUEST_CHANGES

[claude]

Claude: [性能] collect_values 用 sequence_parallel=False 做全长 forward，而 critic 训练在 _prepare_loss 中是 SP 切分的。sp_size>1 时同一 DP 组的每个 SP rank 都重复算同一份 value，开销约 sp_size 倍，且 controller 只读取 replica 0 的结果。

[claude]

Claude: [测试] 本 PR 的核心正确性主张——critic value 按 global_indices 回填到原样本顺序、再把 returns 分发回各 DP rank——完全没有测试。_compute_ppo_targets 是纯函数（只依赖入参），不需要 Ray 即可直接覆盖长度不匹配、未覆盖全部样本、以及 data_replicate_size>1 时只读 replica 0 这三条路径。CriticWorker.collect_values / _materialize / fit 同样零覆盖。

[YanhuiDua]

LGTM

[YanhuiDua]

TrainWorker, RoleWorker, ActorWorker, CriticWorker之间关系