语言模型正从"回答问题"迈向"完成任务", 于RL后训练里, 模型不但生成文本, 还会调用搜索、代码执行等外部工具, 依据环境返回继续推理。如此交互令模型具备更强行动能力, 还使训练系统面临一种相对普通RLHF较不稳定的工作负载: 同一批请求能够产生长度相差数十倍的轨迹, 少量超长轨迹拖累整个进程；与此同时, 训练以及对GPU的需求会随着策略演化持续变化, 另起一句, 这使得培训体系面对的工作负载比普通RLHF更加不稳定, 同一批请求可能生成长度相差数十倍的轨迹, 少数极度的超长轨迹会拖慢所有进程, 同时, 训练和对GPU 的需要且会随着政策的演变持续改变。
针对上述那些情况引发的问题, 来自香港中文大学以及香港恒生大学的研究团队提出来了Libra, 它是一个面向RL Post-的资源管理系统。Libra不再把某事物视作固定不变的瓶颈, 而是把训练与另一事物当作一个耦合系统来统一进行优化, 并且借助异构推理集群、因果感知调度以及弹性资源切换, 让有限的GPU资源能够随着实时的工作负载去动态地流动。
有48张A800 GPU, 在其中, Libra于 -R1、DAPO-Math-17K以及R2E-Gym一共三类任务里, 都获取到了最高的吞吐, 其最高能去到基线的3.0倍。并且, 当有着相近的最终奖励时, 达到目标奖励所要花费的时间, 最多能够缩短到基线的1/2.5。现阶段, 论文跟代码都已经公开了。
图1, 展示的是Libra系统的整体概况, 该系统包含全局资源规划器、异构集群、C-MLFQ调度器以及弹性执行机制, 其图片来源于论文。
RL 带来的不只是 "更长的输出"
往往一个标准的 RL 后训练迭代是包含着轨迹生成的, 还有轨迹评估, 以及策略更新。鉴于评估阶段是相对轻量的, 所以系统效率主要取决于两个环节, 一是能多么快速地去产生轨迹, 二是能多么快速地吸收这些轨迹进而更新策略。
在传统的那种推理里面, 请求的长度一般说来和输入提示有着比较强的相关性, 然而, 在RL当中, 轨迹的长度会被运行的时候出现的事件给改变, 比如说, 搜索工具可能会返回一大段的内容, 代码执行有可能失败并且触发多轮的修复, 模型也有可能依据环境反馈去扩展后续的推理, 所以, 轨迹最终的长度在生成之前想要可靠地预测是很难的。
研究团队于R2E - Gym上观察到, 最长的那10%轨迹, 占据了超过50%的时间。更为关键的是, 此分布并非稳定, 随着策略于训练里逐步改变, 模型的工具使用方式以及推理长度也会出现漂移。
这种漂移会将与的结构性差异予以放大, 实验表明, 当序列长度从1K递增至32K token的时候, 延迟增长了95倍, 然而训练时间仅仅增长3.9倍, 缘由是其在解码环节依赖非同步方式递进推理还原连贯正确信息, 对序列长度以及KV cache敏感度更高, 训练却能够由摊薄长度变化产生的影响达成。
图2, 其中（a）部分, 伴随序列朝着更长的方向发展, 延迟的增加速度显著快于训练的进程；（b）部分指出, 在训练的整个过程里, 平均序列的长度还有时间呈现出持续漂移的状态。源自这图片这一来源的便是这个展示图, 而来源处是论文。
这表明, 存在这样一种情况, 即在训练起始阶段还算合理的固定GPU划分方式, 于经过数百步之后, 极有可能演变成严重的不均衡状态。要是出现了速度变慢的情形, 用于训练的GPU便会等待新的数据；倘若训练进程变慢, 所生成的轨迹又会在队列里积压积聚。从头到尾的迭代用时实际上是由这两者当中速度更为缓慢的一方来决定的:
= max (, )
因而, 问题并非仅仅依靠"持续优化"便可化解, 而是得从整体视角以动态方式去寻觅训练与之的平衡契合点。
全局资源规划：同时决定训练和 怎么用 GPU
全球资源规划器, 是天秤座的首要核心设计。于固定图形处理器预算范围内, 它协同进行搜索:
在训练侧, 是运用拓扑感知的决策树来枚举可行的并行策略然后基于显存、通信开销之类的约束事先进行剪枝。而在另一个部分, 则是将请求依据历史长度予以排序, 通过动态规划找异构TP实例和请求区间间的最优分配。且底层Cost于同一时间对两侧的执行时间进行建模, 使得规划器能够对比不同的全局配置。
规划并非仅于启动之际运行一回, Libra会按周期读取最新的轨迹统计情况, 进而重新求解资源配置, 只有在预计收益超出重配置成本的情况下, 才切实触发资源移动, 如此一来既能追踪drift, 又能规避频繁切换所引发的抖动。
Pool：不重建核心通信组，也能移动算力
推出新配置, 并不等同于能够以低成本来执行新配置。在传统分布式训练里, 增添或者去除训练, 常常意味着要重新构建通信组, 还要再度分发状态, 频繁实施操作的代价实在过高。
Libra 将资源划分为三个池：
其关键的原则是, 维持核心训练拓扑始终保持不变, 以完整的那种数据并行副本的形式加入进来, 且不会去改变核心的TP/PP结构, 系统进一步地将副本内部的NCCL通信, 跟副本之间的梯度交换进行了解耦, 成员的变化只是发生在独立的跨副本通信域。
当步入重新投入到再一次参与训练的进程之中时, 它会以一种并非同步的方式去获取最新的模型以及优化器的快照。在恢复的这段期间之内, 核心的训练这一行为依旧持续性地向前推进；处于参入进程里的个体, 经由侧通道传发出携带零梯度这个特征的占位, 以此让核心的 All-达到与"该个体未曾加入"这种在数学领域上面所持有的等价情景。当状态达成对齐之后, 它才会从接下来的那一步起始贡献出真切的梯度。
不预测最终长度，而在工具返回时做因果路由
第二项 Libra 的核心设计, 是 C-MLFQ, 也就是 Multi-Level Queue。
传统的长度预测办法, 是尝试于请求起始之际去推测最终长度, 然而, RL的关键变动常常是在中途出现的。Libra 的观察结果为: 工具返回的大小、成功或者失败状态, 并非是普通的相关特征, 而是后续轨迹扩展的直接因果信号。比如说, 大就会马上增添上下文, 工具若是失败了, 那么就有可能引发重试、诊断以及代码修改。
首先, 运用历史轨迹去构建一棵因果感知前缀树, 这棵树名为C-MLFQ。其次, 树节点是由ID以及此前所有工具返回状态的有序序列来确定的。再者, 要保存从当前节点到轨迹结束的剩余长度分布。最后, 运行时流程被划分成三步。
请求开始时先进入适合短序列的小 TP ；
每次工具返回后，根据工具类型、 大小和执行状态查询前缀树；
只有在剩余长度的均值跟 P90 朝着同一个方向的时候, 才会进行迁移, 不然的话, 就继续留在当前的位置；在轨迹结束之后才行离线更新树。
图示3, C-MLFQ于工具返回之处读取因果状况, 进而判定是否迁移请求, 待完成之后再度更新前缀树。此图片源自论文。
这种设计躲开了额外模型推理, 并且也用不着伴随策略变换反复去训练长度预测器。相较于传统MLFQ等到长度越界了之后再逐级进行迁移, C-MLFQ能够在工具返回完结后更早一些作出一次性决策。
对于 -R1 而言, C-MLFQ 的单次路由准确率达成了 91.1%, 显著高于借助长度预测的方法 （65.2%）以及传统 MLFQ（44.8%）；与此同时, 它的迁移 token 比例仅仅为 8.2%, 系统吞吐实现了 2700 token/s。
48 张 A800 上的端到端结果
团队于6个节点上做实验, 这6个节点共有48张A800 - SXM4 - 80GB GPU, 在节点内部使用/, 在结节之间使用支持RDMA的200 Gb/s RoCE网络, 实验采用GRPO, 该实验的最大模型长度对应40960 token, 并且每个采样16条轨迹。
工作负载覆盖三个差异明显的 RL 场景：
对比方法涵盖了verl-, verl--, verl-- , 还有基于初始选出最优静态配置的ATeaL--。
图4展示的是, 在 -R1方面, Libra所呈现的吞吐及奖励收敛结果, 同时包含在DAPO - Math - 17K 中的相应情况, 还有在R2E - Gym上的相关结果, 其中红线所代表的正是Libra, 该图片的来源为论文。
关于 -R1这点, Libra的平均吞吐大概是/s, 跟AReaL的情况相较起来, 提升幅度约为63%那样, 与verl部分比起来, 则有80%的提高现象 , 而跟verl一对比, 提升比例能达300%这样。在DAPO - Math - 17K那儿, 还有R2E - Gym之上, Libra也同样维持着最高的吞吐状态。
多种方法对同一模型展开训练, 且执行相同步数, 最终它们抵达的奖励较为接近, 差异主要体现在wall - clock time方面。Libra于 - R1上运用17.9小时完成训练, 在DAPO - Math - 17K上运用26.7小时完成训练, 在R2E - Gym上运用63.2小时完成训练, 其达到目标奖励的速度提升幅度最高为2.5倍。
消融实验对各模块作用做了进一步说明, 在R2E - Gym上, 基线吞吐是423 token/s, 加入同构资源规划后提升到510 token/s, 异构TP又带来41 token/s, C - MLFQ增加115 token/s, 最后弹性执行继续增加97 token/s，让完整Libra达到763 token/s, 总体提升约80.4%。
图5, 呈现的是, 从静态均分开始, 逐步添加全局规划、异构TP、弹性执行以及C- -MLFQ之后, 所产生的吞吐变化情况, 其图片来源, 则是论文。
系统实现与开源
Libra含有大约1.3万行, 还有C++/CUDA代码, 其核心RL训练循环是以verl为基础, 生成方面运用vLLM , 训练的一侧采用 -LM。开源的仓库给出了Slurm以及非Slurm的快速入门, 还有数据准备、配置参考、可观测性说明以及实验脚本。
天秤座的核心看法是: 在强化学习里, 并非始终固定不变的系统阻碍。伴随策略以及轨迹分布演变, 实际的阻碍会在训练与别的事物之间变动。唯有一并处理横跨阶段资源配置、阶段以内异构执行以及低成本资源转换, 系统才可以在长时间训练进程中持续靠近最优情形。
作者简介