meta 近日推出了一款名为 LlamaRL 的强化学习框架,该框架采用全异步分布式架构设计,在处理 4050 亿参数模型时,成功将强化学习步骤的耗时从原来的 635.8 秒减少到 59.5 秒,整体效率提升超过 10 倍。
强化学习是一种通过反馈机制不断调整输出的技术,使模型能够更好地满足用户需求。随着对模型精确度和规则适配能力要求的提升,强化学习在大语言模型训练后期的作用愈加重要,已成为众多先进模型系统不可或缺的一部分。
然而,将强化学习应用于大规模语言模型一直面临资源消耗过大的难题。训练过程通常需要大量计算资源,并涉及多个模块的协同工作,例如策略模型、评分器等。当模型参数达到数百亿级别时,内存占用高、数据传输延迟以及 GPU 资源利用率低等问题成为制约因素。
LlamaRL 框架基于 PyTorch 构建,采用全异步分布式架构,有效简化了各组件之间的同步协调,同时支持模块化定制。该框架通过独立的执行器并行运行生成、训练和评分任务,显著降低了等待时间,提升了整体训练效率。
此外,LlamaRL 还利用分布式直接内存访问(DDMA)和 NVIDIA NVlink 技术,实现了高效的数据传输。在 4050 亿参数模型中,仅需 2 秒即可完成模型权重的同步操作。
实测数据显示,在 80 亿、700 亿和 4050 亿参数级别的模型上,LlamaRL 的训练时间分别缩短至 8.90 秒、20.67 秒和 59.5 秒,性能提升最高达 10.7 倍。同时,在 MATH 和 GSM8K 等标准测试中,模型表现稳定,甚至有所增强。
LlamaRL 成功缓解了大模型训练中常见的内存瓶颈与 GPU 利用率不足的问题,为未来更大规模模型的训练提供了更具扩展性的解决方案。