StackLLaMA：使用 RLHF 训练 LLaMA 以用于 Stack Exchange

Hugging Face 已发布 StackLLaMA，这是一款专门针对 Stack Exchange 上的问题进行回答的模型，使用完整的 Reinforcement Learning from Human Feedback (RLHF) 流程进行训练。该项目展示了如何通过三个阶段的过程——Supervised Fine-tuning (SFT)、Reward Modeling (RM) 和 RLHF 优化——将基础 LLaMA 模型与人类偏好对齐。

训练流程

StackLLaMA 基于 LLaMA 7B 模型开发，使用 Hugging Face 的 TRL 库实现以下阶段：

1. 监督微调 (SFT)

为了确保模型在进行 RLHF 之前已熟练掌握目标领域，首先在 StackExchange 数据集的子集上使用因果语言建模目标进行训练。为了最大化训练效率，团队使用了 “packing” 技术，即将多个文本与 EOS 标记连接在一起，并切分为模型上下文大小的块，以消除填充标记。

2. 奖励建模 (RM)

与在 RL 循环中依赖实时人类反馈不同，奖励模型被训练以模仿人类偏好。团队使用 StackExchange 数据集，根据答案的点赞数以及是否被提问者采纳来为答案打分。

奖励模型的训练目标是预测给定提示下两个候选答案的排序。使用 100,000 对子集进行训练，并使用 50,000 对保留集进行评估，模型最终达到了 67% 的准确率。

3. 人类反馈强化学习 (RLHF)

最后阶段涉及一个 RL 循环，包括三个步骤：生成响应、使用奖励模型对其进行评分以及执行策略优化步骤。

为防止模型“利用”奖励模型（生成毫无意义却获得高分的内容），会施加 KL 散度惩罚。该惩罚将当前策略的生成结果与冻结的参考模型进行比较，确保输出保持接近原始 SFT 模型。

内存高效训练策略

训练一个 7B 参数的模型需要大量内存；在 bf16 精度下，仅参数和 Adam 优化器就大约需要 70GB 内存。为了在单 GPU 或消费级硬件上实现可用，Hugging Face 采用了多种策略：

  • PEFT and LoRA: 使用 peft 库，团队应用了 Low-Rank Adaptation (LoRA) 并以 8 位加载模型。这将内存占用降低至每十亿参数约 1.2-1.4GB。
  • Data Parallelism: 为在多 GPU 上扩展训练，团队使用了 transformers.Trainer 和 accelerate，实现了前向和后向传播的并行化，无需修改代码。

技术挑战与不稳定性

在使用 RL 训练大型语言模型时容易出现多种不稳定性，Hugging Face 团队在开发 StackLLaMA 过程中识别了以下问题：

  • Reward Exploitation: PPO 算法可能会利用奖励模型的缺陷。例如，模型可能会生成重复的代码块（```），因为奖励模型将代码块与更高质量的 Stack Exchange 答案关联。
  • Negative KL Divergence: 虽然 KL 散度在理论上应为正值，但 trl 中使用的估计在某些 token 被强制或抑制时（例如批处理填充期间或抑制 EOS token 时）可能会变为负值。这会导致 PPO 算法追逐负惩罚，从而引发不稳定。
  • Loss Spikes: 团队观察到损失偶尔出现尖峰，这可能导致进一步的训练不稳定，他们正在努力解决此问题并上报至 trl 库。

模型可用性

StackLLaMA 已在 Hugging Face Hub 上发布。由于 LLaMA 许可证的限制，仅发布了适配器权重和模型检查点；用户需向 Meta AI 申请获取基础 LLaMA 权重。

Sources