RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
Reinforcement learning (RL) has emerged as the de-facto paradigm for improving the reasoning capabilities of large language models (LLMs). We have developed RLAX, a scalable RL framework on TPUs. RLAX employs a parameter-server architecture. A master trainer periodically pushes updated model weights to the parameter server while a fleet of inference workers pull the latest weights and generates new rollouts. We introduce a suite of system techniques to enable scalable and preemptible RL for a diverse set of state-of-art RL algorithms. To accelerate convergence and improve model quality, we have devised new dataset curation and alignment techniques. Large-scale evaluations show that RLAX improves QwQ-32B's pass@8 accuracy by 12.8% in just 12 hours 48 minutes on 1024 v5p TPUs, while remaining robust to preemptions during training.
Code (0)
등록된 구현이 없습니다.
Tasks
Reinforcement LearningSimilar Papers 제목 키워드 기반
MO-Playground: Massively Parallelized Multi-Objective Reinforcement Learning for Robotics
Multi-objective reinforcement learning (MORL) is a powerful tool to learn Pareto-optimal policy families across conflicting objectives. However, unlike traditional RL algorithms, existing MORL algorithms do not effective…
Reinforcement LearningDistributed Multi-Agent Reinforcement Learning Based on Graph-Induced Local Value Functions
Achieving distributed reinforcement learning (RL) for large-scale cooperative multi-agent systems (MASs) is challenging because: (i) each agent has access to only limited information; (ii) issues on convergence or comput…
Multi-agent Reinforcement Learningreinforcement-learningReinforcement LearningReinforcement Learning (RL)Optimization for Reinforcement Learning: From Single Agent to Cooperative Agents
This article reviews recent advances in multi-agent reinforcement learning algorithms for large-scale control systems and communication networks, which learn to communicate and cooperate. We provide an overview of this e…
Distributed OptimizationMulti-agent Reinforcement Learningreinforcement-learningReinforcement Learning+1Dynamic network congestion pricing based on deep reinforcement learning
Traffic congestion is a serious problem in urban areas. Dynamic congestion pricing is one of the useful schemes to eliminate traffic congestion in strategic scale. However, in the reality, an optimal dynamic congestion p…
Deep Reinforcement Learningreinforcement-learningReinforcement LearningReinforcement Learning (RL)LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Trainin
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 parti…
GPUReinforcement Learning (RL)