用协作式时间切片填平 RL 流水线的 GPU 空窗

2026-07-24 33 预计阅读时间: 1 分钟
来源: cloud.google.com AI 摘要 Original link

Disclaimer: This article is an AI-assisted summary. Read it together with the original source when precision matters. The summary may omit context, version differences, or edge cases and is not official documentation.

预计阅读时间:9 分钟

大模型强化学习后训练并不总是缺 GPU,很多时候真正的问题是 GPU 已经分配,却在等待下一阶段。同步 RL 在 rollout 采样与梯度训练之间严格交替;异步 RL 虽然允许两者重叠,训练器仍可能因新鲜轨迹不足而停顿。llm-d 引入的协作式时间切片把这些阶段变成可调度单元,让多个独立 RL 作业共享同一组物理加速器。

初步基准显示,这种平台级复用可以把加速器总体工作周期从约 40% 提高到 70%,同时不改变模型收敛和精度。它优化的不是单个训练步骤,而是整套硬件在一段时间内真正执行有效计算的比例。

RL 循环为什么会锁住昂贵算力

典型 RL 后训练循环包含生成轨迹、计算奖励、更新策略以及传播新权重等步骤。同步架构中的资源状态通常是:

当前阶段 Sampler GPU Trainer GPU
rollout 生成 忙碌 等待轨迹
梯度更新 等待新权重 忙碌
权重传播 部分等待 部分等待

空闲资源并不能直接交给其他任务。训练框架的 CUDA 上下文、模型参数、优化器状态和缓存仍驻留在显存中,Kubernetes 看到的也是一个持续占用 GPU 的 Pod,而不是一个暂时不需要计算能力的 RL 阶段。

异步架构只能缓解这个问题。生成通常仍是流水线瓶颈,而接近 on-policy 的训练又必须限制数据陈旧度。训练器追上 rollout 生产速度后,只能等待足够新鲜的轨迹,形成零散但频繁的 GPU 空窗。

调度对象从 Pod 下沉到 RL 阶段

协作式时间切片要求应用在接触加速器的阶段边界主动调用两个操作:

  • acquire():申请一组加速器的独占使用权。
  • yield():释放使用权,允许另一个作业接管硬件。

如果 Job A 进入奖励计算、数据准备或等待轨迹的阶段,它可以主动让出 GPU。编排器随后保存 Job A 的设备状态,并恢复 Job B 已保存的状态。任意时刻只有一个作业的状态占据设备显存,因此两个框架不会同时争抢显存,也不会因为普通的进程级并发而触发 OOM。

这不是容器重启。切换过程以 checkpoint/restore 为核心:设备状态从显存转移到主机 DRAM,稍后再恢复到显存。被换出的进程保持暂停和热状态,不需要重新初始化 PyTorch、vLLM 或重新从存储加载模型。

代价也很明确:主机必须准备足够的内存,切换会消耗 PCIe 或互连带宽,并产生暂停延迟。如果阶段本身很短,保存完整设备状态的成本可能超过回收空窗得到的收益。

三层架构如何完成一次切换

llm-d 将职责拆成三个范围:

  1. 工作负载层:训练循环、推理服务和 RL 框架通过 Python 客户端标记阶段边界,原有 CUDA 上下文和框架代码基本保持不变。
  2. 集群协调层:时间片编排器为共享同一批节点的作业维护锁队列。队首作业持有独占锁,其余作业阻塞在 acquire()
  3. 节点硬件层:以特权 DaemonSet 运行的 Snapshot Agent 暂停进程、序列化设备状态并执行恢复。首个后端基于 cuda-checkpoint,接口则允许未来接入更快或更细粒度的实现。

一次完整切换按照以下顺序发生:当前作业调用 yield();编排器通知每个节点上的 Agent 冻结进程并把设备状态移入主机内存;锁转交给下一个作业;Agent 恢复该作业的状态;等待中的 acquire() 返回,计算从暂停点继续。

多节点训练必须把这件事当作集体操作。只恢复部分 rank 会破坏 NCCL 等通信组的执行状态,因此编排器需要在组内所有节点上协调暂停与恢复,而不是各节点自行抢占。

在训练循环中标出可切换边界

下面的代码沿用 llm-d 客户端暴露的装饰器形式。运行前需要按项目用户指南安装并配置 timeslice 客户端,将 orchestrator:50051 改成实际编排器地址,并确保两个 group_id 对应正确的加速器组。

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)

装饰器将加速器访问限制在明确的函数范围内。业务代码仍然保持顺序执行,但底层可以在 generate_phasetrain_phase 之间插入其他 RL 作业。使用 Ray 等平台分别提交采样和训练任务时,也可以把时间片控制放到平台层,减少研究代码中的改动。

接入时不要把每个小算子都包装成一个时间片。更合理的边界是 rollout 批次、完整的若干次梯度更新,或者一次明确的数据等待周期。阶段需要足够长,才能摊薄保存和恢复设备状态的成本。

上线前应该测什么

采用时间切片不能只观察 nvidia-smi 的瞬时利用率。建议至少记录以下指标:

  • 聚合加速器 duty cycle,以及 sampler、trainer 各自的等待时间。
  • checkpoint 与 restore 的 P50、P95 延迟和传输字节数。
  • 主机 DRAM 峰值、内存带宽和多作业并发时的容量余量。
  • 每小时生成 token、每小时完成的训练 step,以及端到端实验耗时。
  • 作业切换前后的 loss、reward、数据陈旧度和最终评测结果。
  • Agent 或编排器故障时,锁回收、状态恢复和作业重试是否可靠。

最适合优先试点的是阶段较长、空闲窗口明显,并且存在互补采样与训练节奏的两到三个作业。短任务、高频切换或主机内存紧张的集群,可能暂时更适合通过批次调整和流水线并行提升利用率。

llm-d 后续计划包括更快的 checkpoint/restore 后端、只交换 LoRA 适配器等特定内存区域、自动识别互补作业,以及扩展到 TPU 和定制加速器。现阶段的关键判断仍然朴素:先量化空闲窗口,再测量一次完整切换的成本;只有可回收时间显著大于切换开销,时间切片才会真正降低训练总成本。


相关推荐