图解OpenRLHF中基于Ray的分布式训练流程

文章预览

本文着重分析OpenRLHF中的 PPO-Ray训练架构设计 ，没有使用过Ray的朋友也可以通过本文快速上手，本文共分成四块： 1. 为什么用Ray 2. 使用图例抽象出整体训练流程 3. Ray核心知识速过 4. 使用图例，进一步抽象出核心代码细节，包括： 训练入口 部署PPO-Actor/Ref/Critic/RM实例 部署vllm_engines实例 PPO-Actor与vllm_engines之间的通讯 PPO-Actor/Critic训练 一、为什么要使用Ray 对于通常的rlhf框架， 在训练时会在单卡上同时部署actor/ref/reward/critic四类模型 ，这种单一的部署方式可能存在如下问题： 难以突破单卡显存的限制。 无法实现更多的并行计算 。例如在收集exp阶段，拿到(prompt, responses)结果的四类模型其实可以做并行推理；在训练阶段，拿到exp的actor和critic也可以做并行训练。但受到单卡显存等因素影响，通常的rlhf框架中使用更多的是串行。 无法独立优化训练和推 ………………………………