首页 / 资讯中心 / 文章详情

超越前缀局部性:混合Rollout调度器的设计与工程实践

超越前缀局部性:混合Rollout调度器的设计与工程实践 ★ FEATURED ARTICLE
在强化学习RL训练循环中rollout 生成是最容易被低估、也最容易卡住训练进度的一环。以 LLM 对齐任务为例训练进程需要不断从当前策略采样一批 prompt让生成引擎批量产出回答再把这些回答整理成训练样本。当这批 prompt 来自不同任务、不同数据源甚至不同策略版本时调度问题就变成了“混合 rollout 调度”。前缀局部性prefix locality是这类场景里常用的优化前提如果多个请求共享同一个前缀那么 KV Cache 可以复用从而显著减少重复 prefill 计算。但在真实 RL 训练中前缀共享并不是唯一需要关注的维度。调度器设计者需要同时处理队列、优先级、资源预算、训练 deadline 和策略版本隔离才能真正让 rollout 环节稳定地喂饱训练循环。这篇文章从一个调度器设计者的视角讨论为什么要超越前缀局部性以及如何把混合 rollout 调度做到既保证吞吐又兼顾训练稳定。1. RL 训练循环中rollout 为什么需要专门的调度层1.1 训练循环里的“策略采样”并不是简单的 generate在 RLHF 或 RLVR 这类训练流程中训练进度通常由“策略采样 → 奖励计算 → 策略更新”三个步骤反复推进。策略采样阶段不是简单调用一次推理接口而是需要选择一批 prompt并携带当前策略版本的采样参数。把 prompt 发送给一个或多个生成引擎。等待生成引擎返回完整回答即 rollout。将 rollout 交给 reward model 或规则打分。把打分的样本交给训练进程做策略更新。这个环节里采样数量、采样时机和采样完成时间会直接影响训练稳定。传统在线推理场景中调度器只需要关注“用户请求能不能尽快返回”。而在 RL 训练场景中调度器还必须在指定训练 step 内把足够多的 rollout 样本准备好。否则训练进程只能空等GPU 浪费在 idle 上。因此RL 训练循环里需要独立出一个 rollout 调度层。它介于“采样请求生成方”和“生成引擎执行方”之间负责把训练侧的时间约束翻译成生成侧可执行的批次任务。1.2 前缀局部性在 rollout 场景里到底意味着什么前缀局部性这个概念来自自回归生成模型的运行方式。Transformer 生成每个 token 时都要读取已有的 KV Cache。如果两个请求使用同一个前缀例如相同的 system prompt、相同的指令模板、相同的情景对话开头那么可以复用该前缀对应的 KV Cache避免重复 prefill。在 rollout 场景中前缀局部性的典型收益来自两处训练集里大量 prompt 共享同一个固定的 instruction 前缀例如“请根据以下信息完成问答”多轮对话式任务里不同采样分支共享同一段历史对话。如果调度器能把共享前缀的请求合并到同一批次生成引擎可以只计算一次 prefix然后并行生成不同后缀从而显著降低显存占用和 prefill 延迟。不过前缀局部性只是“可利用的空间局部性”它描述的是请求特征而不是任务目标。在混合 rollout 调度中前缀是否相同只是其中一个维度。请求的优先级、deadline、策略版本、生成长度、资源占用可能比前缀复用更关键。1.3 为什么只按前缀局部性调度会翻车如果调度器把所有决策都建立在“前缀是否相同”上会碰到几类非常现实的问题。第一不同任务的前缀结构差异很大。有些任务共享长前缀但生成的 max_tokens 很长有些任务前缀不同但 deadline 马上就要到了。如果调度器一直等待更多共享前缀请求来凑 batch就会让 deadline 临近的请求被饿死。第二不同策略版本之间不能共享 KV Cache。策略模型经过一次更新后权重已经改变旧版本的 KV Cache 不能直接用于新版本。此时即使前缀完全一样把它们合并到一个批次里也是没有意义的。第三一个训练 step 内可能需要同时包含探索样本和利用样本。探索样本来自旧策略或随机采样利用样本来自当前策略。二者优先级不同超时容忍度也不同。前缀局部性无法表达这种优先级差异。所以混合 rollout 调度必须把 prefix locality 当作“可选的加分项”而不是“调度的唯一依据”。2. 混合 rollout 调度的问题定义与核心约束2.1 什么算“混合 rollout”这里说的“混合”不是单纯指多个 prompt 放在一起而是指一个调度周期内同时存在多种不同维度的请求。常见混合点包括任务来源混合训练数据包含数学推理、代码生成、问答、安全对齐等任务每个任务的 prompt 模板和长度分布不同。策略版本混合同一时刻可能存在当前策略、旧策略、探索策略的 rollout 请求它们对应不同模型权重。采样参数混合温度、top_p、max_tokens 不同导致生成时间差异巨大。优先级混合on-policy 样本决定当前训练 step 是否继续必须高优off-policy 或辅助任务样本可以延后。可以对比一下单一调度和混合调度的差异调度维度单一 rollout 调度混合 rollout 调度请求来源单一 prompt 池多个 prompt 池、多个策略版本前缀结构高度一致长短差异大共享程度不稳定优先级基本一致高优、普通、低优并存deadline宽松训练 step 强约束资源预算可预估难以预估依赖历史统计KV Cache 复用明显需要按策略版本判断是否有效2.2 调度器需要优化的多个目标混合 rollout 调度器不是只做“先来先服务”它需要同时面对多个目标吞吐量每分钟生成的 token 总数。等待时间请求从进入队列到开始生成的等待时间。Deadline miss 率在训练 step 截止前没有完成生成的请求比例。饥饿避免低优先级任务不能被高优任务无限压制。资源利用率显存和算力是否被充分使用。这些目标之间存在明显冲突。例如为了提升批次大小而等待更多请求会增大等待时间把所有资源都给高优任务则低优任务会长期饥饿。因此实际调度器很少追求全局最优解而是通过启发式规则在多个目标之间取平衡。2.3 混合调度中的三条硬约束在算法目标之外调度器还受到三条硬约束约束资源约束生成引擎的显存、并发度、最大 batch token 数都有限。调度器不能把无限请求塞给一个引擎。时间约束训练 step 有明确 deadline超时会导致训练暂停。调度器必须支持 deadline 检查和超时截断。一致性约束on-policy rollout 必须在策略更新前完成如果旧策略已经被淘汰那么相关请求应该被取消而不是继续占用资源。这三条约束决定了调度器的核心数据结构设计优先级队列、deadline 队列、策略版本隔离组、超时取消回调缺一不可。3. 超越前缀局部性调度策略的四个关键设计3.1 把请求拆成“可组合的调度单元”要超越前缀局部性首先要重新设计请求模型。不要把请求看成一个不可分割的整体而要把每个 rollout 请求看成一组可组合字段request_id全局唯一标识。prompt实际输入文本。prefix_key用于识别可复用前缀的键通常由模板 ID 前若干 token 构成。policy_version策略版本号决定使用哪个模型权重。priority请求优先级例如LOW、NORMAL、HIGH、URGENT。max_tokens允许生成的最大 token 数。deadline完成截止时间戳。estimated_tokens预估生成的 token 数可由历史统计获得。调度器在做决策时可以根据这些字段的任意组合进行拆分或合并。例如按policy_version先做隔离再按prefix_key做合并然后按deadline做排序。这样前缀局部性就从一个“全局调度目标”降级为“局部合并条件”。3.2 用“时间约束优先前缀复用加分”的两阶段选择一个简单且实用的调度策略是两阶段选择第一阶段从所有未调度的请求中根据优先级和 deadline 选出候选集合。规则可以是先取URGENT和HIGH队列中的请求再补入 deadline 在 10 秒内的普通请求。这一步保证时间敏感请求不被漏掉。第二阶段在候选集合内进行批次组装。此时才考虑前缀复用。对候选请求按prefix_key分组在每一组内按max_tokens装箱形成 batch。如果某些请求的前缀不同但依然能正好凑满一个 batch 的 token 预算也不要强行排到下一批。用公式表达就是候选条件priority P或deadline - now T。合并条件same policy_version且same prefix_key且sum(max_tokens) budget。兜底条件same policy_version且sum(max_tokens) budget且deadline最短优先。这样设计的好处是前缀复用只是“优先尝试”的路径不是“必须等待”的条件。3.3 为不同策略版本单独建立队列避免串扰策略版本是混合 rollout 调度中最容易忽略的维度。KV Cache 只有在同一模型权重下才能复用。如果两个 rollout 请求来自不同策略版本把它们放进同一个 batch生成引擎需要在 batch 内切换模型状态或者在共享前缀时错误复用 KV Cache都会造成严重错误。因此调度器内部应该按policy_version建立逻辑队列。每个队列独立维护自己的优先级、deadline 和 batch 预算。跨策略版本的调度只发生在资源分配层例如两个队列竞争同一个 GPU 资源时根据总优先级和 deadline 决定把算力分给谁。这种隔离设计也方便扩展。如果未来生成引擎支持在同一进程内加载多个策略模型调度器只需在 batch 创建时把policy_version作为第一拆分键后续合并逻辑无需改动。3.4 异步完成与超时回收机制生成式 rollout 往往不是立刻返回的。一个长回复可能需要几十秒。调度器不可能一直阻塞等待每个请求完成。必须要设计异步机制每个请求进入生成引擎后调度器只记录“已派发”状态。生成引擎通过回调或轮询方式返回结果。调度器定期检查未完成请求的 deadline。如果某个请求已经接近 deadline调度器可以发送 cancel 信号或者把剩余生成任务截断到当前已生成部分。超时回收是保证训练 step 不被打爆的关键。RL 训练通常对样本数量有最低要求。缺少样本时宁可截断部分 rollout 也要让训练继续而不是无限等待一两个长尾请求。4. 最小实现一个可运行的混合 Rollout 调度器4.1 数据结构设计下面用一个最小 Python 示例说明调度器核心逻辑。完整生产实现需要对接真实生成引擎和训练框架但队列、优先级、deadline、前缀分组这些思想是通用的。先定义请求数据结构from dataclasses import dataclass from enum import IntEnum import time import heapq class Priority(IntEnum): LOW 0 NORMAL 1 HIGH 2 URGENT 3 dataclass class RolloutRequest: request_id: str prompt: str policy_version: int prefix_key: str priority: Priority max_tokens: int deadline: float property def score(self) - float: # 分数越低越紧急 return self.deadline - int(self.priority) * 1000score是一个简易的紧急度分数把 deadline 往前移优先级越高越靠前。实际系统可以用更复杂的公式。4.2 核心调度逻辑调度器包含三个主要方法enqueue把请求放入按策略版本拆分的优先级队列。_pick_candidates从队列里挑出满足时间约束的候选。schedule执行两阶段选择生成 batch。class MixedRolloutScheduler: def __init__(self, max_batch_tokens: int 4096): self.max_batch_tokens max_batch_tokens self.queues {} def enqueue(self, request: RolloutRequest): if request.policy_version not in self.queues: self.queues[request.policy_version] [] queue self.queues[request.policy_version] heapq.heappush(queue, (request.score, request.request_id, request)) def _pick_candidates(self, now: float, policy_version: int): queue self.queues.get(policy_version, []) candidates [] urgent [] # 先从队列中取出高优先级或临近 deadline 的请求 for score, rid, req in queue: if req.priority Priority.HIGH or req.deadline - now 5.0: urgent.append((score, rid, req)) else: candidates.append((score, rid, req)) # 重新把未选中的放回队列这里仅作演示 urgent.sort(keylambda x: x[0]) return [r for _, _, r in urgent] def _build_batch(self, candidates): # 先按前缀分组再在组内按 deadline 排序 groups {} for req in candidates: groups.setdefault((req.policy_version, req.prefix_key), []).append(req) batch [] current_tokens 0 for key in sorted(groups.keys()): group_requests sorted(groups[key], keylambda r: r.deadline) for req in group_requests: if current_tokens req.max_tokens self.max_batch_tokens: batch.append(req) current_tokens req.max_tokens return batch def schedule(self, now: float None): now now if now is not None else time.time() all_batches [] for policy_version in list(self.queues.keys()): candidates self._pick_candidates(now, policy_version) if candidates: batch self._build_batch(candidates) all_batches.append({policy_version: policy_version, requests: batch}) return all_batches这个实现省略了从队列中移除已派发请求的细节。实际工程中需要注意请求一旦被选中进入 batch就应该从队列中弹出避免重复调度。这里为了演示清晰没有在_pick_candidates中删除原队列元素。4.3 运行一个小样例构造一组混合请求模拟不同策略版本、不同优先级和不同前缀now time.time() scheduler MixedRolloutScheduler(max_batch_tokens4096) scheduler.enqueue( RolloutRequest( request_ida1, prompttask: math\nquestion: 11?, policy_version10, prefix_keytask: math, priorityPriority.HIGH, max_tokens128, deadlinenow 10, ) ) scheduler.enqueue( RolloutRequest( request_ida2, prompttask: math\nquestion: 2*3?, policy_version10, prefix_keytask: math, priorityPriority.HIGH, max_tokens128, deadlinenow 10, ) ) scheduler.enqueue( RolloutRequest( request_idb1, prompttask: code\nquestion: write a python function, policy_version10, prefix_keytask: code, priorityPriority.NORMAL, max_tokens512, deadlinenow 30, ) ) scheduler.enqueue( RolloutRequest( request_idc1, promptold policy sample, policy_version8, prefix_keytask: generic, priorityPriority.LOW, max_tokens256, deadlinenow 60, ) ) batches scheduler.schedule(nownow) for batch in batches: print(batch[policy_version], [r.request_id for r in batch[requests]])预期输出中policy_version10的队列会优先合并a1和a2因为它们共享prefix_key且优先级高b1可能因为当前 batch token 预算不够或优先级低而没有进入当前 batchc1属于policy_version8会单独形成一个 batch。这个结果说明调度器在决定 batch 时既没有忽略前缀共享也没有单靠前缀共享决定一切。5. 关键参数、调优方向与环境差异5.1 调度器参数速查表实际落地时需要针对业务和资源情况调整参数。下面是几个关键参数及其影响。参数含义常见值调大影响调小影响max_batch_tokens单个 batch 的最大 token 数4096~8192批量吞吐提升显存压力增大出 batch 延迟变高等待时间变短资源利用率降低schedule_interval_ms调度器每隔多少毫秒触发一次50~200决策更及时CPU 开销增大调度延迟变高容易 miss deadlinedeadline_buffer距离 deadline 多少秒时提前调度3~10 秒降低超时率提前占用资源资源被过早占用浪费算力prefix_key 粒度如何生成前缀键模板 ID / 前 N token合并率高但可能错误复用不相关内容安全但前缀命中率降低priority 等级数优先级划分数量3~5 档表达更细粒度但队列管理复杂调度简单可能无法区分紧急程度policy_version 数量同时存在的策略版本数1~3样本多样性增加显存和队列开销增大资源占用少但探索不足5.2 学习环境与生产环境的部署差异上面的最小示例适合在单机、模拟请求环境下学习调度思想。进入生产后需要面对更多额外问题。学习环境只有少量请求进程内队列就够用调度直接打印 batch 即可。开发环境需要接入真实生成引擎验证请求格式、返回回调、deadline 取消是否正常。测试环境可以模拟大规模请求观察队列长度、批次 token 分布、超时率等指标。生产环境调度器本身要无状态化使用 Redis 或消息队列存储待调度请求需要支持多副本所有决策要有日志便于回放问题。生产环境中还应在调度器外部增加配置中心方便动态调整优先级权重、deadline buffer、batch 预算等参数而不是每次改代码后重启服务。6. 验证、监控和常见问题排查6.1 用三类指标检验调度器部署调度器后不能只说“看起来能运行”需要用指标验证它是否真的满足目标。指标类别具体指标正常范围/期望效果指标deadline miss 率尽量低于 5%效果指标高优任务平均等待时间低于普通任务等待时间效果指标饥饿请求数持续等待超过 T 的请求数应趋近 0资源指标batch 平均 token 数接近max_batch_tokens的 70%~90%资源指标prefix 命中率高于 30% 可认为有收益训练稳定指标训练 step 阻塞时间低于阈值训练稳定指标样本到达速率与训练消费速率匹配6.2 常见问题排查路径混合 rollout 调度最容易出现的问题集中在队列堆积、deadline 超时和资源利用不均。排查时按“进入队列 → 被调度 → 被生成引擎执行 → 返回结果”这条链路逐段检查。问题现象常见原因检查方式处理建议高优队列堆积请求优先级设置过高但生成引擎并发不足查看队列长度和生成引擎 QPS增大引擎并发或增加副本调整调度权重deadline miss 率升高调度周期太长或 deadline_buffer 太小查看调度日志中的“纳入调度时间”缩短调度周期提前调度时间prefix 命中率很低prompt 模板不一致或 prefix_key 粒度太细统计 prefix_key 分布统一 prompt 模板调整前缀键生成逻辑训练侧一直等 rollout生成任务处理慢或中途失败查看失败率和生成引擎负载增加重试和失败补偿逻辑显存 OOMbatch token 预算设置过大查看引擎显存占用降低max_batch_tokens缩小 batch过期请求占用资源没有超时取消机制检查队列中是否残留过期请求增加 deadline 检查和 cancel 回调6.3 一个可复用的排错清单每次 rollout 调度出问题可按下面顺序快速定位请求是否确实进入队列打印enqueue日志确认policy_version和deadline正确。请求是否被调度器选中检查schedule日志确认优先级和 deadline 条件是否触达。请求是否被发送到生成引擎检查发送接口返回码。生成引擎是否正常返回检查引擎日志和 token 数统计。如果有超时检查超时时间设置是否过短或生成引擎是否出现长尾请求。如果训练侧仍未拿到样本检查回调、消息队列和数据写入链路。把这条清单固化成巡检脚本比临时查日志更高效。7. 最佳实践与扩展方向7.1 四条可立即执行的建议第一让 deadline 和优先级成为调度器的一等公民。不要在算法实现到一半时才发现时间约束没有数据结构表达。第二把前缀复用当作“额外收益”而不是“调度前提”。前缀命中率高值得开心但绝不能因为等待相同前缀而饿死 deadline 请求。第三对策略版本显式建模。KV Cache 复用只能在同一个策略版本内发生跨版本调度必须隔离。第四调度器尽量保持无状态。请求可以放进外部队列调度结果写日志。这样即使调度器崩溃也能基于日志恢复现场。7.2 从迷你调度器到生产级调度器先用本章的最小示例跑通逻辑然后逐步向生产演进接入真实生成引擎把batch转换为引擎可接受的请求格式。将进程内队列替换为 Redis Stream 或 Kafka 等队列支持多实例消费。增加请求状态表记录pending、scheduled、running、completed、timeout状态。为每一个调度决策写入结构化日志包含请求 ID、批次 ID、策略版本、触发原因。增加监控大盘对队列长度、deadline miss 率、前缀命中率持续观测。7.3 扩展方向超越前缀局部性之后调度器还可以继续吸收更多维度基于历史的 token 预估利用历史请求的平均生成长度预测新请求的耗时让 deadline 判断更准。缓存预热把训练集中高频出现的 prompt 前缀提前放入 KV Cache而不是等请求到达后再重复计算。多租户调度当多个训练任务共享同一批生成引擎时调度器需要支持租户配额和共享池。连续批处理结合调度器负责“何时选谁”生成引擎负责“选好后如何抢占空闲 slot”。二者配合才能达到高吞吐。对新手而言最值得投入的时间是在一个最小调度器上反复修改优先级规则、deadline 规则和 batch 组装逻辑。只要把队列、deadline、策略版本这三件事理解透彻再叠加 prefix locality 就会非常自然。前缀局部性不是不好而是在混合 rollout 场景里它必须跟其他约束共享决策权。
阅读完成 · 觉得有帮助?
咨询建站