人工智能领域，2024年2月7日 – 近日，随着 DeepSeek-R1 的发布，群组相对策略优化（GRPO）因其在大型语言模型强化学习中的有效性和易用性而备受关注。然而，不少开发者在使用 GRPO 过程中遇到了内存占用过高的问题，尤其是在 GPU 资源有限的情况下。针对这一痛点，AI 公司 Oxen.ai 的 CEO Greg Schoeninger 分享了他的实践经验，为开发者提供了宝贵的指导。
GRPO 是一种在线学习算法，它通过使用训练过程中由训练模型自身生成的数据来进行迭代改进。其核心目标是最大化生成补全的优势函数，同时确保模型保持在参考策略附近。然而，GRPO 的高内存需求也给许多开发者带来了挑战。
Schoeninger 在文章中提到，他在使用配备 16GB 显存的 Nvidia GeForce RTX 3080 移动版笔记本电脑进行训练时，遇到了显存不足（OOM）的错误。这促使他进行了一系列实验，以确定训练各种大小模型所需的显存（VRAM）要求。实验涵盖了参数数量从 5 亿到 140 亿不等的模型，并比较了权重的完全微调与参数高效微调（PEFT）。所有训练运行都在英伟达 H100 上完成，即使在这样的高端硬件上，OOM 仍然意味着超过 80GB 的 VRAM 需求。
实验结果表明，内存需求随着模型大小和训练方式的不同而显著变化。例如，全参数微调比 PEFT 需要更多的内存。那么，为什么 GRPO 对内存需求如此之高？
原因在于 GRPO 的工作原理。它涉及多个模型，并且在训练数据中每个查询会产生多个输出。策略模型、参考模型和奖励模型各自都是一个需要进行推理的 LLM。
为了缓解内存压力，Schoeninger 建议采用以下两种技术：
- 8-Bit 优化器： 使用像 AdamW 这样的 8-bit 优化器版本，它们能更高效地存储跟踪数据，同时仍保持良好的性能。
- 梯度检查点： 在训练过程中拍摄快照，而不是记录所有内容。虽然这会使训练速度减慢约 20-30%，但它显著减少了内存使用。
这些技术能够显著降低内存占用，使得即使 GPU 资源有限的开发者也能够训练更大的模型。
Schoeninger 还提到，像 trl 这样的库已经开始支持 GRPO，使得微调由 transformers 构成的 LLM 变得非常简单。
总结
DeepSeek 的 GRPO 算法在大型语言模型强化学习中展现出强大的潜力，但其高内存需求也给开发者带来了挑战。通过理解 GRPO 的工作原理，并采用 8-Bit 优化器和梯度检查点等技术，开发者可以在有限的 GPU 资源下成功训练大型模型。
参考文献
- Greg Schoeninger. GRPO VRAM Requirements for the GPU Poor. https://www.oxen.ai/blog/grpo-vram-requirements-for-the-gpu-poor
- 机器之心. DeepSeek用的GRPO占用大量内存？有人给出了些破解方法. 原文链接
Views: 0