LlamaRL实现千卡级大模型强化学习高效训练,速度最高提升10.7倍。
LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Training
- 单控制器异步分布式架构,基于原生PyTorch实现无缝扩展。
- 405B模型下相比DeepSpeed-Chat快10.7倍,规模越大优势越明显。
- 适合超大规模语言模型的后训练场景,尤其适用于千卡集群。
强化学习(RL)已成为提升大语言模型(LLM)能力最有效的后训练方法。然而,由于对延迟和内存的高要求,构建一个能可靠管理数百亿至数千亿参数策略模型的高效RL框架极具挑战。本文提出LlamaRL,一个完全分布式的异步强化学习框架,可支持8B、70B和405B参数量的LLM在从少量到数千个GPU的集群上高效训练。该框架采用全原生PyTorch实现的轻量化单控制器架构,具备模块化、易用性和可扩展性。我们提供了理论分析,形式化证明其异步设计能带来严格的速度提升。在Llama 3后训练中,通过采用共置模型卸载、异步离策略训练和分布式直接内存访问权重同步等最佳实践,LlamaRL在405B参数模型上相较DeepSpeed-Chat类系统实现最高达10.7倍的速度提升。效率优势随模型规模增大而持续增长,表明其适用于未来更大规模的强化学习训练。
原文摘要 · Abstract (English)
Reinforcement Learning (RL) has become the most effective post-training approach for improving the capabilities of Large Language Models (LLMs). In practice, because of the high demands on latency and memory, it is particularly challenging to develop an efficient RL framework that reliably manages policy models with hundreds to thousands of billions of parameters. In this paper, we present LlamaRL, a fully distributed, asynchronous RL framework optimized for efficient training of large-scale LLMs with various model sizes (8B, 70B, and 405B parameters) on GPU clusters ranging from a handful to thousands of devices. LlamaRL introduces a streamlined, single-controller architecture built entirely on native PyTorch, enabling modularity, ease of use, and seamless scalability to thousands of GPUs. We also provide a theoretical analysis of LlamaRL's efficiency, including a formal proof that its asynchronous design leads to strict RL speed-up. Empirically during the Llama 3 post-training, by leveraging best practices such as colocated model offloading, asynchronous off-policy training, and distributed direct memory access for weight synchronization, LlamaRL achieves significant efficiency gains -- up to 10.7x speed-up compared to DeepSpeed-Chat-like systems on a 405B-parameter policy model. Furthermore, the efficiency advantage continues to grow with increasing model scale, demonstrating the framework's suitability for future large-scale RL training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。