The N 实现细节 of RLHF with PPO – Hugging Face 博客摘要
Hugging Face 博客文章复现了 OpenAI 的 2019 年 RLHF 代码库,匹配了其学习曲线,并枚举了影响可重复性的实现细节。
匹配学习曲线
该复现产生的学习曲线几乎与 OpenAI 原始代码库在情感和描述性任务上的学习曲线完全相同。作者在 AWS p3dn.24xlarge 实例上运行原始的 TensorFlow 1.x 仓库以获得基准指标,然后使用这些指标来验证他们的基于 PyTorch 的复现。
通用实现细节
奖励模型和策略价值头接收拼接的查询和响应作为输入。序列会用特殊填充 token 填充到固定长度,过长时会被截断,并且在 logit 计算期间会根据填充 token 调整位置索引。响应生成采样固定长度的输出,不会在 EOS token 处停止,并且学习率在奖励模型和策略训练中会被退火到零。每个 GPU 进程使用不同的随机种子以鼓励探索。
奖励模型实现细节
奖励模型仅输出拼接的查询‑响应对最后一个 token 的值。奖励头的权重从均值为 0、方差为 1/(sqrt(d_model)+1) 的正态分布初始化,偏置设为零。奖励归一化在训练前后通过计算固定参考策略的奖励经验均值和标准差,然后缩放到单位方差和零均值来完成。
策略训练实现细节
在计算对数概率之前,logits 会除以采样温度。价值头权重初始化为均值为 0、方差为 0。策略训练期间禁用 Dropout。拒绝采样在 token 16 和 24 之间强制一个周期,否则分配一个固定的低奖励。折扣因子 γ 设置为 1。训练使用批次、小批次和微批次进行梯度累积。每个 token 的 KL 惩罚会被加到奖励中,并且每个小批次的奖励和优势会进行白化处理(仅对优势进行均值偏移)。价值函数的裁剪方式与标准 PPO 相同,KL 惩罚系数 β 会根据观测到的 KL 散度进行自适应调整。
PyTorch Adam 优化器在 RLHF 中的数值问题
PyTorch 的 Adam 实现与 TensorFlow 的不同之处在于 epsilon 在分母中的位置,这导致训练早期的有效归一化项更小,从而产生更激进的梯度更新。实验表明,与 TensorFlow‑style Adam 相比,PyTorch Adam 会产生更高的对数概率方差、更大的比率极端值以及增加的近似 KL 和裁剪比例,尤其是在较大的模型如 gpt2‑xl 上。
局限性
该工作未尝试复现来自 CNN/DM 或 TL;DR 的摘要任务,原因是计算成本高且训练不稳定。原始代码需要 8× V100 32GB GPU,而复现过程中 GPU 利用率低(约 30%),导致运行成本高且脆弱。
结论
通过复现 OpenAI 的 RLHF 代码库并记录其实现细节,作者提供了一个理解 RLHF 工程实践的参考,并突出了一种微妙的优化器差异,这种差异可能会显著影响训练动态。