让异构显卡协同训练大模型强化学习,提升效率降低成本。
AReaL-Hex: Accommodating Asynchronous RL Training over Heterogeneous GPUs
- 分阶段调度异构显卡资源,动态分配生成与优化任务。
- 相比同质部署,吞吐量最高提升50%,成本降低46%。
- 适合追求性价比的RL训练团队,尤其在多型号显卡环境。
提升大模型强化学习训练的吞吐量和成本效率,是推动该技术普及的关键。一个有前景但具挑战性的方向是在异构GPU上部署计算流程。与传统的大型语言模型预训练不同,强化学习训练通常分为三个紧密耦合的阶段:采样生成、奖励计算和策略/价值更新,三者在计算强度、内存占用和通信模式上差异显著。近期研究表明,完全异步的强化学习训练可将各阶段拆分至独立硬件池中而不影响训练稳定性,为实际异构部署创造了机会。为此,我们提出AReaL-Hex,一种面向异构性的异步强化学习训练系统,能有效调度采样生成与策略模型训练在异构显卡上的执行,并保证数据陈旧性约束。具体地,采用两阶段调度器:(i) 基于混合整数线性规划(MILP)的受限搜索,根据资源预算选择各阶段并行策略与工作负载分配;(ii) 图划分步骤,将异构显卡与互连结构进行最优分配,以最大化端到端吞吐量。基于全异步架构,AReaL-Hex将高显存/带宽受限的生成任务映射到更低成本资源,同时平衡生产者-消费者关系,避免空闲与过时轨迹。在数学推理任务中,针对1.5B、7B和14B模型规模,相较于最先进的同质异步强化学习系统:(i) 在相同总预算下,训练吞吐量最高提升1.50倍;(ii) 在相同吞吐量目标下,训练成本最高降低1.46倍。
原文摘要 · Abstract (English)
Maximizing training throughput and cost-efficiency of RL for LLMs is essential to democratize this advanced technique. One promising but challenging approach is to deploy such a computational workflow over heterogeneous GPUs. Unlike conventional large-scale LLM pretraining, RL training generally decomposes into three coupled stages, i.e., rollout generation, reward computation, and policy/value updates, which exhibit markedly different compute intensities, memory footprints, and communication patterns. Recent research shows that fully asynchronous RL training can disaggregate these stages across disjoint hardware pools without sacrificing training stability, creating a great opportunity for real-world heterogeneous deployment. To this end, we present AReaL-Hex, a heterogeneity-aware asynchronous RL training system that effectively schedules how to execute rollout generation and policy model training over heterogeneous GPUs while enforcing data staleness bounds. Concretely, we use a two-phase scheduler: (i) a constrained search with MILP to select per-stage parallelization strategies and workload assignments given a resource budget, and (ii) a graph-partitioning step that allocates heterogeneous GPUs and interconnects to maximize end-to-end throughput. Built atop a fully asynchronous RL architecture, AReaL-Hex maps HBM-I/O-bound generation and compute-bound optimization to more cost-efficient resources and balances their producer-consumer interactions to avoid both idleness and stale rollout trajectories. On the mathematical reasoning task with various model scales (1.5B, 7B, and 14B), compared to homogeneous deployments of state-of-the-art asynchronous RL systems: (i) When maintaining the same total budgets, AReaL-Hex delivers up to 1.50x higher training throughput; (ii) When achieving the same training throughput, AReaL-Hex results in up to 1.46x reduction in training cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。