ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

verl 中的 Optimal Token Baseline(OTB):基于 Token 级路径方差的最优基线优势估计实践指南

verl 中的 Optimal Token Baseline(OTB):基于 Token 级路径方差的最优基线优势估计实践指南 verl 中的 Optimal Token BaselineOTB基于 Token 级路径方差的最优基线优势估计实践指南【免费下载链接】verlverl/HybridFlow: A Flexible and Efficient RL Post-Training Framework项目地址: https://gitcode.com/GitHub_Trending/ve/verlverlHybridFlow在examples/otb_trainer中提供了 Optimal Token BaselineOTB训练器的完整参考实现与 GRPO 那种“每个 prompt 组一个均值基线”的做法不同OTB 在每个时间步为同一 prompt 组内的所有采样轨迹计算一个由累计路径方差加权的最优基线从而在 token 粒度上降低优势估计的方差。读完本文你将掌握 OTB 的数学原理与源码实现、optimal_token_baseline与tir_optimal_token_baseline两种变体的区别、calculate_sum_pi_squaredTrue的开关作用以及如何基于run_qwen3_8b_fsdp.sh在 NVIDIA GPU 上用 vLLM 采样 FSDP 训练跑通 OTB。一、什么是 Optimal Token Baseline1.1 动机从“组均值基线”到“时间步最优基线”在以 GRPO 为代表的 group-based RL 训练中每个 prompt 会采样 N 条轨迹基线通常取该组奖励的均值即对整条轨迹或整个组使用单一基线。这种做法的局限在于一条轨迹内部不同位置的不确定性是不同的——有的时间步模型“很有把握”概率集中有的时间步模型“非常随机”概率接近均匀分布。把所有位置同等对待会带来不必要的方差。OTB 的核心思想是为每个时间步 t 单独计算一个最优基线并且用“累计路径方差”作为该时间步在组内加权聚合时的权重。verl 的 core_algos.py 中对该理论的描述为每个时间步的最优基线B_t* E[G_t × W_t] / E[W_t]其中W_t Σ_{j1}^t ||s_j||²是到时间步 t 为止的累计路径方差代理而每个时间步的方差代理为||s_j||² 1 - 2π_j Σπ²W_t刻画了轨迹“到目前为止已经实现的能量/不确定性”累计值越大说明该轨迹在高方差路径上走得越远其收益在基线估计中的权重就越高——直观上给“走钢丝走得越久”的轨迹更大的话语权能更准确地估计该时刻应当扣减的基线。1.2 与 TIR 变体的关系OTB 存在两个注册在 core_algos.py 中的变体枚举值字符串名适用场景OPTIMAL_TOKEN_BASELINEoptimal_token_baseline单轮 / 通用场景TIR_OPTIMAL_TOKEN_BASELINEtir_optimal_token_baselineTIRTool-Integrated Reasoning多轮工具调用场景两者的数学形式完全一致B_t* E[G_t × W_t] / E[W_t]差别在于处理变长轨迹与掩码的方式compute_optimal_token_baseline_advantage直接在[bs, response_length]的稠密矩阵上按response_mask分组计算compute_multi_turn_optimal_token_baseline_advantage先把每个轨迹的w_cumulative与token_returns按真实长度“压实”unpad到[bs*n, max_response_length]再计算基线最后再写回原始[bs, turn * response_length]布局因此天然适配 TIR 多轮multi-turn数据的拼接结构。如果你在跑 TIR 多轮训练用环境变量ADV_ESTIMATORtir_optimal_token_baseline即可切换到该变体见 run_qwen3_8b_fsdp.sh 第 11 行。二、OTB 在 verl 中的注册与调用链2.1 优势估计器注册机制verl 用一个可扩展的注册表管理所有优势估计器GAE、GRPO、RLOO、REMAX、OPO、GPG 等OTB 只是其中之一。定义位于 core_algos.pyclass AdvantageEstimator(str, Enum): ... OPTIMAL_TOKEN_BASELINE optimal_token_baseline TIR_OPTIMAL_TOKEN_BASELINE tir_optimal_token_baseline注册装饰器register_adv_est(...)core_algos.py把两个实现函数分别绑定到这两个字符串名上compute_optimal_token_baseline_advantagecore_algos.pycompute_multi_turn_optimal_token_baseline_advantagecore_algos.py因此配置algorithm.adv_estimatoroptimal_token_baseline就能在训练时被解析并调用对应实现。2.2 训练器中的调用点在 PPO 训练器 ray_trainer.py 中当检测到adv_estimator属于 OTB 两个变体之一时会断言 batch 中必须存在sum_pi_squaredassert sum_pi_squared in data.batch, ( Step-dependent optimal baseline requires sum_pi_squared from actor. Please set actor.calculate_sum_pi_squaredTrue in config. ) adv_kwargs[sum_pi_squared] data.batch[sum_pi_squared] adv_kwargs[old_log_probs] data.batch[old_log_probs] adv_kwargs[rollout_is_weights] data.batch.get(rollout_is_weights, None)这段代码同时说明了三个关键事实calculate_sum_pi_squaredTrue是硬性前置条件——不开这个开关训练会在优势计算处直接断言失败OTB 使用old_log_probs采样时的策略概率来构造路径方差代理w_t 1 - 2·exp(old_log_probs) sum_pi_squared当启用 Rollout Correctionrollout_is_weights存在时OTB 还会把重要性采样权重平方ρ̄²(t)乘到W_t上用于在截断 IS 下最小化 MSE。2.3sum_pi_squared从哪里来sum_pi_squared记为 Σπ²即模型在词表上的概率平方和由 actor 前向时逐 token 计算相关开关定义在 actor.pycalculate_sum_pi_squared: bool False在 FSDP 引擎 transformer_impl.py 中当calculate_sum_pi_squaredTrue时前向会调用verl_F.calculate_sum_pi_squared_from_logits(logits)并把结果写入model_output[sum_pi_squared]。其数值实现位于 torch_functional.pydef calculate_sum_pi_squared_from_logits(logits: torch.Tensor): Formula: Σπ² exp(logsumexp(2*logits) - 2*logsumexp(logits)) return torch.exp(torch.logsumexp(2.0 * logits, dim-1) - 2.0 * torch.logsumexp(logits, dim-1))该式利用 log-sum-exp 的数值技巧直接给出词表概率平方和的精确值全程不显式计算 softmax 概率数值稳定性更好。需要特别注意的是FSDP 实现中会检查calculate_sum_pi_squaredTrue与use_fused_kernelsTrue不能同时开启transformer_impl.py因为融合核路径不产出该统计量。三、OTB 优势计算的核心步骤源码级拆解compute_optimal_token_baseline_advantagecore_algos.py的完整流程可以拆成五步Step 1计算每个时间步的收益reward-to-goreturns (token_level_rewards * response_mask).flip(dims[-1]).cumsum(dim-1).flip(dims[-1])对 token 级奖励从序列末尾向前做 cumsum得到每个位置的累计回报G_t。Step 2逐时间步方差代理pi_t torch.exp(old_log_probs) w_per_timestep 1 - 2 * pi_t sum_pi_squared由1 - 2π Σπ²得到每个位置的“不确定性度量”。当策略在该位置完全确定π 为 one-hotΣπ²1时该项为 0当策略接近均匀分布时该项接近 1从而在数值上捕捉每个位置的熵/方差信息。Step 3累计路径方差w_cumulative (w_per_timestep * response_mask).cumsum(dim-1)W_t Σ_{j1}^t w_j即从轨迹起点累计到当前时间步的方差代理之和。Step 4按 prompt 分组计算每时间步最优基线numerator (returns_group * w_cumulative_group * mask_group).sum(dim0) denominator (w_cumulative_group * mask_group).sum(dim0) epsilon baseline_per_step numerator / denominator对组内所有轨迹按时间步做加权平均B_t* Σ[G_t × W_t] / Σ[W_t]epsilon1e-8防除零。注意两个边界行为组内只有 1 条轨迹时不计算基线优势直接等于收益handle_zero_tailTrue默认开启时会把组内最长轨迹超出第二长轨迹长度的那段“尾巴”的基线置零——因为那段只有一条轨迹参与没有可比的组内参照。Step 5计算优势advantages (returns - baselines) * response_mask每个 token 的优势A_t G_t - B_t*最终乘response_mask屏蔽 padding 位置。TIR 变体core_algos.py的数学流程一致但多了一步“按真实长度压实 → 在[bs*n, max_response_length]紧凑张量上计算 → 写回原始多轮布局”的转换Step 4这正是它能正确处理 TIR 多轮数据的关键。四、官方标准脚本run_qwen3_8b_fsdp.sh 逐段解读verl 为 OTB 提供了官方 canonical 脚本 run_qwen3_8b_fsdp.sh配置矩阵为文本任务 vLLM 采样 FSDP 训练 NVIDIA GPU。脚本顶部给出了可调环境变量及其默认值环境变量默认值作用MODEL_PATHQwen/Qwen3-8B基座模型HuggingFace Hub 路径或本地路径NNODES/NGPUS_PER_NODE1/8节点数与每节点 GPU 数ADV_ESTIMATORoptimal_token_baseline优势估计器TIR 场景设为tir_optimal_token_baselineTRAIN_BATCH_SIZE128训练 batch 大小PPO_MINI_BATCH_SIZE128PPO mini-batch 大小MAX_PROMPT_LENGTH1024最大 prompt 长度MAX_RESPONSE_LENGTH2048最大响应长度PPO_MAX_TOKEN_LEN_PER_GPU24576动态 bsz 下每 GPU 最大 token 数ACTOR_LR1e-6actor 学习率ENTROPY_COEFF0熵正则系数ROLLOUT_TP2vLLM 张量并行度ROLLOUT_GPU_MEM_UTIL0.75vLLM GPU 显存利用率ROLLOUT_N8每个 prompt 采样的轨迹数组大小TOTAL_EPOCHS/SAVE_FREQ/TEST_FREQ15/20/5训练总轮数、保存频率、评测频率PROJECT_NAME/EXPERIMENT_NAMEverl_otb_gsm8k_math/qwen3_8b_vllm_fsdpwandb/日志项目与实验名4.1 数据部分DATADATA( algorithm.adv_estimator${adv_estimator} algorithm.use_kl_in_rewardFalse data.train_files$train_files data.val_files$val_files data.train_batch_size${train_batch_size} data.max_prompt_length${max_prompt_length} data.max_response_length${max_response_length} data.filter_overlong_promptsTrue data.truncationerror )脚本默认混用 GSM8K 与 MATH 两个数据集train_files[$HOME/data/gsm8k/train.parquet, $HOME/data/math/train.parquet]测试集同理使用前请按 prepare_data.rst 准备好 parquet 格式的数据。值得注意的是algorithm.use_kl_in_rewardFalse——OTB 示例走的是“KL 从损失侧约束use_kl_lossFalse时则完全不做 KL 约束 组内优势估计”的路线与经典的 in-reward KL 惩罚不同。4.2 actor 部分OTB 的两个关键开关ACTOR( actor_rollout_ref.actor.optim.lr${actor_lr} actor_rollout_ref.actor.ppo_mini_batch_size${ppo_mini_batch_size} actor_rollout_ref.actor.use_dynamic_bszTrue actor_rollout_ref.actor.ppo_max_token_len_per_gpu${ppo_max_token_len_per_gpu} actor_rollout_ref.actor.use_kl_lossFalse actor_rollout_ref.actor.entropy_coeff${entropy_coeff} actor_rollout_ref.actor.calculate_sum_pi_squaredTrue actor_rollout_ref.actor.fsdp_config.param_offloadFalse actor_rollout_ref.actor.fsdp_config.optimizer_offloadFalse )actor_rollout_ref.actor.calculate_sum_pi_squaredTrueOTB 的必需开关。如前文所述缺失它会在 ray_trainer.py 处断言失败。同时注意不要与use_fused_kernelsTrue同时启用。use_dynamic_bszTrueppo_max_token_len_per_gpu24576开启动态 batch 大小以 token 数为单位控制显存。entropy_coeff0示例默认不施加熵正则。4.3 采样与参考策略部分ROLLOUT / REFROLLOUT( actor_rollout_ref.rollout.namevllm actor_rollout_ref.rollout.tensor_model_parallel_size${rollout_tp} actor_rollout_ref.rollout.gpu_memory_utilization${rollout_gpu_mem_util} actor_rollout_ref.rollout.n${rollout_n} actor_rollout_ref.rollout.log_prob_use_dynamic_bszTrue actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu${ppo_max_token_len_per_gpu} ) REF( actor_rollout_ref.ref.log_prob_use_dynamic_bszTrue actor_rollout_ref.ref.log_prob_max_token_len_per_gpu${ppo_max_token_len_per_gpu} actor_rollout_ref.ref.fsdp_config.param_offloadTrue )采样由 vLLM 承担tensor_model_parallel_size2表示 rollout 阶段做 2 路张量并行n8表示每个 prompt 采 8 条轨迹OTB 的组大小。参考策略ref只用于计算 KL尽管本示例use_kl_in_rewardFalse其 FSDP 参数卸载param_offloadTrue以节省显存。4.4 训练器与启动方式TRAINER / LAUNCHTRAINER( trainer.balance_batchTrue trainer.critic_warmup0 trainer.logger[console,wandb] trainer.project_name${project_name} trainer.experiment_name${experiment_name} trainer.n_gpus_per_node${NGPUS_PER_NODE} trainer.nnodes${NNODES} trainer.save_freq${save_freq} trainer.test_freq${test_freq} trainer.total_epochs${total_epochs} )critic_warmup0说明 OTB 不需要 value network与 GRPO 同族属于无 critic 的优势估计因此也无需 critic 预热。启动部分脚本 104-123 行直接调用python3 -m verl.trainer.main_ppo传入上述全部参数数组若环境变量VERL_USE_UV默认 1且设备为 GPU则改用uv run --frozen --all-packages --extra vllm --extra fsdp启动 driver 与 Ray worker保证 vLLM × FSDP 依赖组合与仓库锁定的uv.lock一致。运行时需在 verl 仓库根目录下执行bash examples/otb_trainer/run_qwen3_8b_fsdp.sh五、关键配置速查OTB 训练的最小配置面配置项推荐值说明algorithm.adv_estimatoroptimal_token_baselineTIR 多轮场景用tir_optimal_token_baselineactor_rollout_ref.actor.calculate_sum_pi_squaredTrue必需否则训练断言失败勿与use_fused_kernelsTrue同开actor_rollout_ref.rollout.n 2示例 8组大小组内只有 1 条轨迹时 OTB 退化为无基线algorithm.use_kl_in_rewardFalse示例示例采用组内优势路线actor_rollout_ref.actor.use_dynamic_bszppo_max_token_len_per_gpuTrue/24576动态 bsz 控制显存六、实践要点与注意事项两个开关必须同时就位algorithm.adv_estimatoroptimal_token_baseline与actor_rollout_ref.actor.calculate_sum_pi_squaredTrue。前者决定调用 OTB 实现后者保证 batch 中携带sum_pi_squared张量缺任一都会中断训练。理解handle_zero_tail的边界处理默认开启时组内最长轨迹的“无人可比”尾巴基线被置 0该段优势等于原始收益避免用单一轨迹的统计量污染基线。rollout IS 兼容OTB 通过rollout_is_weights以ρ̄²(t)缩放W_t可平滑接入 algorithm.py 中的 Rollout Correction 配置缓解 rollout 与训练策略不一致带来的偏差。数据准备与运行环境脚本默认数据位于$HOME/data/gsm8k与$HOME/data/math训练入口是verl.trainer.main_ppo运行前请确认依赖vLLM、FSDP extras已按 install.rst 安装。OTB 在 verl 中的实现路径清晰配置开关 → actor 前向产出sum_pi_squaredtorch_functional.py→ 训练器组装参数ray_trainer.py→ 核心算法按组逐时间步加权求基线core_algos.py。无论你是想复现 OTB 训练还是希望把它作为自定义优势估计器的起点verl 的register_adv_est注册表支持以字符串名注册新实现本文给出的脚本、源码链路与配置矩阵都足以支撑你直接上手。【免费下载链接】verlverl/HybridFlow: A Flexible and Efficient RL Post-Training Framework项目地址: https://gitcode.com/GitHub_Trending/ve/verl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表