Google Cloud Blog

Minimize idle accelerators: Native RL job interleaving with co-operative time-slicing in llm-d

8.5内容质量
Minimize idle accelerators: Native RL job interleaving with co-operative time-slicing in llm-d

TL;DR · AI 摘要

Google推出llm-d项目通过动态调度技术,将强化学习训练中加速器利用率从40%提升至70%,显著降低计算成本。

核心要点

  • co-operative time-slicing技术使加速器利用率提升至70%
  • llm-d项目通过动态调度减少RL任务空闲时间
  • 同步/异步工作负载均能受益于时间切片优化

结构提纲

按章节快速跳转。

  1. 介绍RL训练中加速器利用率低的行业痛点及llm-d解决方案

  2. 通过动态调度将采样和训练步骤作为可调度实体进行时间切片

  3. 初始测试显示加速器利用率从40%提升至70%且不影响模型收敛

  4. llm-d栈包含吞吐驱动推理引擎和高吞吐代理沙箱

  5. 同步设置减少空闲窗口,异步工作负载利用碎片化空闲时间

思维导图

用一张图看清主题之间的关系。

查看大纲文本(无障碍 / 无 JS 友好)
  • co-operative time-slicing
    • 技术原理
      • 动态调度RL步骤
      • 时间切片机制
    • 应用效果
      • 利用率提升70%
      • 降低TCO
    • 架构组件
      • llm-d-router
      • Agent Sandbox

金句 / Highlights

值得收藏与分享的关键句。

#强化学习#llm-d#Google Cloud#资源优化
打开原文

为 llm-d 中的强化学习引入协作时间片分配 | Google Cloud 博客

容器与 Kubernetes

最小化空闲加速器:llm-d 中通过协作时间片分配实现原生 RL 任务交错

2026 年 7 月 24 日

##### Poonam Lamba

高级产品经理

##### Aishu Kamal

软件工程师

##### 今天立即体验 Gemini Enterprise Business Edition

职场 AI 的入口

立即体验

大型语言模型(LLM)训练后阶段的强化学习(RL)数学计算以苛刻著称。当前沿 AI 实验室使用 Group Relative Policy Optimization(GRPO)等 RL 训练后算法推动推理和编码模型边界时,他们经常遭遇严重的架构和基础设施限制。尽管业界大部分关注点仍集中在获取原始加速器容量上,但要实现运行多个 RL 任务所需的高效率并推动模型达到更高智能水平,基础设施效率同样关键。在大规模场景中,分布式 RL 因同步采样和训练严格按顺序执行的阶段划分而面临严重资源瓶颈,导致训练器和采样器资源交替空闲。同时,异步架构虽然尝试重叠这些阶段,但训练器在等待特定轨迹批次完成前仍会频繁出现空闲间隙。

今天,我们通过 llm-d 项目引入协作时间片分配方案解决这种结构性浪费。通过将离散的 RL 步骤(如采样 rollout 和梯度训练)视为动态可调度实体,我们能够将独立的 RL 任务交错分配到共享物理硬件上。初步基准测试表明,这种平台级复用技术在不损害模型收敛性和准确性的情况下,将加速器整体负载率从约 40% 的基准提升至 70%。这通过消除随时间累积的浪费计算显著提升了性价比并降低了总拥有成本(TCO)。

对于同步设置,平台会交错调度采样器和训练器以最小化交替空闲窗口;而异步工作负载则利用时间片分配动态回收并利用 RL 训练器迭代之间的碎片化空闲间隙。

在本文中,我们将详细介绍时间片分配方案,涵盖技术流程、当前版本和未来路线图。

llm-d 用于 RL 基础设施效率(整体视角)

从一开始就预见了大规模 RL 训练后阶段的严重基础设施瓶颈,我们投入资源解决 RL 工作负载的基础设施效率问题。

我们已将 llm-d 构建为一个高度可组合的基础设施栈,专注于消除加速器空闲时间,适用于推理、智能体和 RL 工作负载。llm-d 的 RL 基础设施栈包含以下特性:

  • 以吞吐量驱动的推理(llm-d-router):一个经过生产验证的引擎,部署于各类 RL 工作负载,专注于最大化 rollout 生成吞吐量以持续饱和流水线。
  • 高吞吐智能体沙箱(recipe):经过规模和密度测试,可在 rollout 生成和评估期间提供安全的亚秒级工具使用和隔离代码执行。智能体沙箱作为奖励信号生成的高速进气歧管,确保沙箱永远不会成为导致时间片分配 NVIDIA GPU 瓶颈的延迟源头。
  • 核心管道原语:为应对权重传输中的可靠性和速度问题,我们正在构建权重传播接口(Weight Propagation Interface,WPI),同时专注于提升强化学习(RL)的整体可观测性和可靠性。

强化学习循环的效率问题

分布式强化学习的后训练过程表现为一个碎片化的持续循环,交替进行生成(采样轨迹)和优化(梯度更新)。由于传统云基础设施是为持续稳定的负载设计的,标准的Kubernetes集群无法适应这种交替的节奏。

在大规模部署时,这种结构性的节奏会引入两大系统性效率问题:

  • 空闲的加速器:由于这些阶段是按顺序执行的,GPU集群在生命周期的40%到60%时间内完全空闲(0%利用率)。训练器在等待采样轨迹完成时处于空闲状态;采样器在梯度更新和权重分发期间也处于空闲状态。这每年可能造成数百万美元的资本浪费。
  • 固定上下文:强化学习训练和采样器在整个运行时间(包括空闲阶段)都会保持对加速器的占用,因为NVIDIA CUDA上下文和所有设备内存必须保持驻留。标准调度器将这些Pod视为静态、孤立的分配,而非与实时强化学习循环的交替阶段状态对齐,导致即使在非活跃阶段,宝贵的硬件资源仍被锁定。

重要的是,这不仅是同步强化学习的问题。异步变体虽然会重叠生成和训练过程,但无法完全消除空闲时间。生成仍然是强化学习循环的固有瓶颈,这意味着训练器加速器在等待轨迹数据累积时仍会处于饥饿状态。异步任务越接近策略梯度,这些空闲窗口就会越大——生成和训练之间的滞后期受到有界陈旧性限制,当没有新鲜轨迹数据时,流水线会立即停滞。

合作式时间切片(强化学习任务交错)如何提供帮助

为消除强化学习任务中的空闲加速器,llm-d项目下的合作式时间切片技术使基础设施能够动态地将独立的强化学习任务交错分配到共享的硬件块上,而非强制硬件等待上游阶段。这有助于在不改变底层模型收敛性或准确性的情况下,提升加速器的整体利用率。

当同步强化学习中的任务A在阶段边界处空闲(或在异步强化学习中因等待新鲜轨迹数据而停滞)时,基础设施会将物理加速器进行时间切片,切换到任务B的活跃采样或训练阶段。在底层,这种切换是通过检查点/恢复机制实现的:任务A的整个设备状态会从加速器内存检查点保存到主机DRAM中,而任务B之前保存的状态则会恢复到其位置。由于任何时候只有一个任务的状态占用加速器,切换过程可以安全进行,不会引发框架级别的干扰或内存溢出(OOM)故障。

时间切片:高层架构

时间切片系统架构分为三个层级:工作负载级(应用逻辑)、集群级(协调)和节点级(硬件管理)。

工作负载作用域层(应用运行时) 此处运行用户的代码——训练循环、推理服务器和强化学习框架。新增的是时间片客户端库,它在时间片协调器上暴露了两个 gRPC API:acquire() 用于请求独占加速器访问权限,yield() 用于释放权限。用户使用这些调用将任何涉及加速器的阶段包裹起来,以向协调器发出阶段边界信号。其余部分——机器学习框架(PyTorch FSDP、vLLM 等)、CUDA 上下文、加速器内存分配——均无需修改。

集群作用域层(控制与调度平面) 该层决定哪些任务获得加速器访问权限以及何时获得。共享相同物理加速器的任务(例如,两个强化学习任务在相同 GPU 节点组上交替运行)会被分组。对于每个组,时间片协调器维护一个锁队列——一个等待独占访问该组加速器的任务有序列表。只有队列头部的任务持有锁并运行在硬件上;其他所有任务会等待,阻塞在 acquire() 调用上。当运行中的任务调用 yield() 时,协调器将锁传递给队列中的下一个任务,并在组内所有节点上触发协调的上下文切换。未来,工作负载部署优化器将能够分析工作负载阶段模式,自动将具有互补空闲阶段的任务配对,无需用户显式指定任务分组。

节点作用域层(硬件与数据平面隔离) 该层在每个加速器节点上执行检查点/恢复交换。快照代理(一个特权 DaemonSet)接收协调器的指令,并将其转换为硬件级操作——暂停加速器进程、将设备状态序列化到主机 DRAM,并在任务重新获得访问权限时恢复状态。代理基于可插拔后端接口构建,其中 cuda-checkpoint 是首个实现(后续将有更多实现)。未来的后端将引入更快的快照机制和更选择性的方法,例如卸载特定内存地址(如 LoRA 适配器)而非完整设备状态。代理本身设计为独立于 Kubernetes 运行,以支持裸金属和 Slurm 环境。

流程:所有部分如何协同工作 当工作负载完成当前加速器阶段时,其时间片客户端库会向时间片协调器调用 yield() 以释放访问权限。协调器通过向组内每个节点的快照代理发送指令来启动上下文切换。代理会冻结正在让出的工作负载进程,并将其设备状态从加速器内存转移到主机 DRAM。

加速器释放后,协调器将组锁授予队列中等待的下一个工作负载。它指示这些节点的快照代理将该工作负载先前保存的主机 DRAM 状态恢复到加速器内存,然后解除该工作负载的 pending acquire() 调用阻塞。工作负载会从暂停的位置无缝恢复执行——无需容器重启、无需框架重新初始化、无需从存储重新加载模型。

让出的工作负载状态保留在主机 DRAM 中。当协调器再次授予其锁时,快照代理会执行反向的相同交换操作。

开发者体验(客户端侧) 研究人员希望专注于核心建模逻辑,而非与低层级的CUDA上下文切换或自定义调度循环搏斗。如果你使用Ray或类似平台来编排强化学习任务,使用时间切片对客户端侧的影响将微乎其微。事实上,如果在平台层面将训练和采样任务分别排队,客户端侧甚至可能完全不受影响。

加载中...

python
from timeslice import TimeSliceOrchestratorClient
orchestrator = TimeSliceOrchestratorClient(target="orchestrator:50051")

@orchestrator.on_accelerators(group_id="trainer-group")
def train_phase(model, trajectories):
    return model.update(trajectories)

@orchestrator.on_accelerators(group_id="sampler-group")
def generate_phase(model, prompts):
    return model.generate(prompts)

# 标准顺序循环 —— 在底层与其他任务交错执行
for epoch in range(EPOCHS):
    trajectories = generate_phase(policy, dataset)
    rewards = compute_rewards(trajectories)
    train_phase(policy, rewards)

当前发布版本与未来规划

今天,我们正式发布完整的时间切片技术栈:快照代理(Snapshot Agent)、加速器编排器(Accelerator Orchestrator)以及Python客户端库,每项都附带用户指南,帮助你将时间切片集成到强化学习任务中。

关键路线图亮点包括:

  • 延迟与状态优化:通过更快的检查点/恢复后端扩展快照代理,以最小化上下文切换开销,并引入应用感知后端,实现选择性内存区域快照(例如交换LoRA适配器而非完整模型权重)。
  • 自动化调度与接入:推出自动化调度器,用于分析运行进程、识别可切片结构,并动态处理任务分配。
  • 跨硬件兼容性:将数据平面支持范围从GPU扩展至TPU和自定义加速器架构。

入门指南

构建稳健且高度优化的强化学习基础设施,需要与大规模运行这些任务的工程师和研究人员紧密协作。

如果你目前正面临训练后流水线中GPU利用率低、同步阻塞或复杂调度逻辑等问题,时间切片技术可以提供帮助。要立即开始,可参考以下资源,并别忘了给我们反馈!

  • 通过[用户指南](...)立即在强化学习运行中启用时间切片。
  • 在强化学习生成阶段,尝试使用[llm-d-router(Kubernetes原生)](...)或[RL Scheduler(Python库)](...)用户指南,以提升采样吞吐量。
  • 探索[Weight Propagation Interface仓库](...)。
  • 在llm-d Slack的#sig-rl频道参与讨论。
  • 通过分享参考实现、基准测试和边缘案例为我们的技术路线提供帮助。

感谢Dolev Ish Am和Bogdan Berce对本文的贡献。

发布于

  • 容器与Kubernetes
  • AI基础设施
  • llm-d