RL训练中rollout推理与模型训练负载差异大,Mooncake通过分离调度和批量预取实现高效资源利用。
大规模语言模型的强化学习融合了两种截然不同的工作负载:rollout 生成和模型训练。
在 rollout 阶段,推理 worker 在一组提示词上运行当前策略并生成响应。在此过程中,它们产生学习算法所需的信息,包括生成的 token、mask、log probability、reward、序列长度、样本标识符以及其他元数据。这些输出共同构成了 rollout 数据,将被下一个训练步骤所消费。
在规模较小时,生成和训练可以共享同一执行环境。然而在更大规模下,现代 RL 系统越来越多地采用分散式架构(disaggregated architecture),其中 rollout 生成和训练被部署为独立的 worker 组,通常跨不同进程、GPU 或机器进行。
这两个阶段具有根本不同的资源和执行特征。Rollout 是推理密集型工作负载,其吞吐量取决于解码效率、批处理和请求调度;而训练依赖于大规模、高度同步的张量计算。将二者分离使得各自可以独立扩展和调度,而无需将两种工作负载强制纳入同一执行模式。
这种分离还支持流水线级并发。Miles 可以在 rollout worker 已经在生成 rollout N+1 的同时,对 rollout N 进行训练。异步 RL 系统因此允许 rollout 和训练 worker 以不同速率推进,而不必在每个操作后让一个阶段等待另一个阶段。其好处是更好的资源利用率,以及在配置推理和训练容量方面更大的灵活性。但分散化也引入了一个新的系统边界:rollout worker 产生的数据必须先移动到另一组 worker,然后训练才能消费它。
这个交接点恰好位于生成和下一次策略更新之间。缓慢的传输会让训练器等待数据,并使 rollout worker 上的内存占用时间超过必要长度。对于像 Miles 这样的分布式 RL 系统,因此将这种结构化的 rollout 数据高效地从推理侧移动到训练侧,成为端到端 RL 流水线的关键一环。
RL rollout 数据与分布式训练中常见的大型规则张量有很大不同。
一个 rollout 批次通常是一个异构的结构化对象,而不是单个连续张量。根据框架和算法的不同,它可能包含生成的 token、loss mask、log probability、reward、序列长度、样本标识符、路由信息、元数据以及其他辅助字段。这些值可以表示为张量、NumPy 数组、Python 标量列表、可变长度的每样本数组、字节或任意 Python 对象。
以下几个特性使得这些数据特别难以高效移动。
不同的 rollout 字段具有根本不同的表示方式和语义。稠密张量和数值数组可以作为类型化缓冲区高效传输,而参差不齐的序列(ragged sequences)需要行边界信息,标量列表需要保留其原始值,元数据或 Python 对象可能需要更通用的编码方式。与此同时,训练器必须重建 RL 框架所期望的精确结构,包括 dtype、shape、行顺序、空状态和元数据。通用序列化器可以在功能上处理这些对象,但通常以额外的转换、复制和重建为代价。因此,高效的 rollout 传输需要同时理解每个字段的物理表示和逻辑结构。
Rollout 数据可能包含大量小型内存分配。在捕获的 Miles 工作负载中,tokens、loss_masks 和 rollout_log_probs 等主要字段表示为 list[np.ndarray],即每个样本分配一个 NumPy 数组。随着批次规模增长,逐个传输这些碎片会引入重复的内存注册和 Store 操作,而序列化完整的 Python 对象则需要遍历、复制和重建大型对象图。挑战在于将碎片化的逻辑数据转化为高效的批量传输,同时不丢失其原始结构。
这两个挑战使得 rollout 数据移动不仅仅是一个带宽问题:系统必须在保留其结构的同时高效移动碎片化、异构的数据。
一条有效的数据路径需要同时满足以下几个要求:
效率:避免过多的序列化、复制、对象重建和每次分配的传输开销。
正确性:在往返过程中保留字段类型、shape、行边界、空状态、元数据和 Python 级别的值。
可扩展性:随着样本数量、对象碎片数和总载荷规模的增长,继续保持良好性能。
灵活性:支持异构字段,而不强迫 RL 框架扁平化或重写其原生 rollout 表示。
可预测的交接延迟:足够快速地交付批次,使训练器不会因等待 rollout 数据而停滞。
Miles 是一个高性能的大规模模型后训练强化学习框架。它结合了用于高吞吐量 rollout 生成的 SGLang 和用于可扩展训练的 Megatron-LM,同时为偏好直接训练 Hugging Face 模型实现的 workloads 提供 PyTorch FSDP2 后端。Miles 支持完全异步的 RL,其中 rollout 和训练 worker 解耦且可以独立推进,同时还具备诸如快速环内权重更新、agentic rollout、低精度训练和大规模生产 RL 工作负载的容错等功能。
这种分散式和异步的设计使得 rollout 到训练的数据路径成为 RL 流水线的关键部分。Rollout 批次是结构化且异构的,通常包含碎片化的每样本数据和框架特定的元数据,随着工作负载规模的增长,高效的传输和重建变得越来越重要。
Mooncake 为分布式 AI 工作负载提供了高性能的数据平面。对于 RL rollout 数据,它扩展了这种数据平面,支持结构化对象传输,允许在保留原始结构和语义的同时移动异构和碎片化的 rollout 对象。
Mooncake 现已集成到 Miles 作为 rollout 数据传输后端。在从 Miles 捕获的 rollout 数据上,该集成比现有的 Ray 路径实现了显著更低的传输延迟:远程 GET 快约 10–14 倍,而 PUT 改善约 1.2–1.6 倍。
结果是实现了更快的 rollout 到训练交接,而无需改变 RL 编程模型,为 Miles 提供了更高效的大规模结构化 rollout 工作负载数据路径。
Miles 支持同步和异步两种训练循环。在任一模式下,一旦 rollout worker 和训练器部署在独立进程或独立机器上,每个已完成的 rollout 批次必须跨越该边界,才能被下一个训练步骤所消费。
一个有用的设计原则是将控制平面与数据平面分离。框架调度器决定 rollout 批次应该去哪里,但不应携带主体 payload 本身。相反,它传递一个轻量级的传输引用,不包含张量 payload 和 Store chunk 布局,同时在需要时仍允许 JSON 安全的元数据。实际的 rollout payload 走单独的路径:
这种分离使调度决策保持轻量,同时允许主体 rollout payload 通过专用数据路径移动。
从 Miles 捕获的 rollout 数据使这条数据路径具体化。对于给定的框架配置,字段契约是稳定的,但字段并不共享一种便捷的内存表示。它们大致分为三组:
第一组携带了我们抓包中的大部分字节。其他字段较小,但不能被丢弃或归一化:它们携带样本标识、长度、奖励、特征状态和记账信息。其确切大小取决于工作负载;基准测试部分给出了一个实测示例。
Miles 使用完成后的 Dict 交接
在 Miles 中,当完整的 rollout 字典准备就绪后,当前交接开始。生产者调用 put(data, type="dict"),调度器携带返回的引用,然后训练器调用 get 来重建原始字典。
图 1. Miles 对完整的 rollout 字典使用同步的 put 和 get。引用通过调度器传输,而 Mooncake 通过 Store 数据平面移动有效载荷。
各个传输调用是同步的。生产者在 put 完成之前不返回引用,训练器在消费对象之前等待 get 完成。
这并不阻止 RL 流水线本身是异步的。Miles 可以在训练 rollout N 的同时生成 rollout N+1。换句话说,并发位于传输操作之上:每次单独的交接是同步的,而不同的 rollout 和训练阶段可以在流水线层面重叠。
Mooncake 如何保留和传输 Miles Rollout 数据
Mooncake 在不同层次上处理这两个挑战。它让结构保持可见足够长时间,以便为每个字段选择高效的表示:tensor 和数组保持类型化,不规则行携带紧凑的边界元数据,Python 值保留 GET 重建它们所需的信息。然后它将碎片化内存转换为批量 I/O。复制计划将符合条件的较小行直接打包到可复用的已注册 BufferPool 块中,而大型连续 tensor 和数组可以使用原生 Store 路径。这避免了两个极端:既不将整个字典序列化为一个不透明 blob,也不为每个小分配发起 Store 操作。
结果的有效载荷成员作为 bundle 发布。其清单记录了这些成员的位置及其如何组装,并在有效载荷就绪后才可见。GET 按相反的描述来获取和重建原始 Miles 字典。Miles 仍然使用小型公共接口 put(data, type="dict") 和 get(ref, type="dict");schema、字段布局、传输计划和重建逻辑都保留在 Mooncake 内部。
图 2. Schema 和叶子展开暴露每个字段的类型和结构。字段特定编码生成类型化的有效载荷成员和重建元数据;Bundle Store 最后发布其清单。符合条件的碎片化传输使用 BufferPool 支持的暂存区,GET 按相同结构反向操作。
为每个字段选择布局
Mooncake 首先将字典展开为可以高效编码的叶子。数组和 tensor 保持类型化,不规则行携带紧凑的边界元数据,支持的 Python 值保留重建所需的信息。框架提供的 schema 为模糊或性能敏感的字段固定存储表示;当没有提供 schema 时,Mooncake 从观察到的值推断。
同样的选择也适用于多模态数据。处理后的像素和相关模型输入可以保留在 tensor 路径上。PIL 图像、编码的 PNG 或 JPEG 字节以及可变长度媒体列表使用媒体感知布局,保留其边界和重建元数据,而不会强制每种表示都通过相同的序列化器。
将碎片化内存转换为批量 I/O
单独发送每一行会产生数千次注册和 Store 请求。首先拼接整个字段可以避免请求计数,但会增加一个完整大小的临时对象。Mooncake 则构建一个复制计划,并根据需要填充可复用的已注册 BufferPool 块。其原生路径将符合条件的数值行直接复制到这些块中,避免了 Python 行循环和临时拼接。当 Store 支持时,大型连续数组和 tensor 仍然使用 direct 或原生路径。
发布完整 Bundle 并直接重建
Mooncake 仅在所有有效载荷和元数据就绪后才发布 bundle 清单,因此读者永远不会看到半写的 Miles 字典。在 GET 时,清单标识要获取的成员,结构化元数据描述如何重建每个字段。符合条件的读取可以定向到 BufferPool 支持的目标,而无需中间 bytes 对象,类型化不规则行可以直接查看结果缓冲区。
训练器使用完本地 BufferPool 支持的结果后释放。一旦所有读者完成,Miles 移除短寿命的 Store 对象,Mooncake 回收其有效载荷块和清单。
图 3. 两个主要的 Miles rollout 数据挑战直接映射到 Mooncake 的结构化编码和批量传输优化。
对于固定的框架配置,切换模型本身不会改变字段契约。传输大小反而跟随样本数量、提示和响应长度以及可选字段(如教师对数概率、路由信息或多模态输入)变化。
Miles 使用 Qwen3-0.6B 在数学提示上生成源数据(rollout_id=0,8 个源样本)。此抓包中每个响应都有 256 个 token。对于基准测试,三个大型数值字段被归一化为类型化的 ndarray 行,同时保留其值、行长度、dtype 和每样本碎片化。在这些样本上平均,逻辑有效载荷分解如下:
前三个字段约占 3,386 字节,即此特定样本布局的 98.9%。更大的基准测试有效载荷重复八个抓包样本,保留其字段类型和碎片化分配模式,同时增加逻辑样本计数。该计算描述了此次抓包;它不定义固定的 Miles 样本大小。
图 4. Qwen3-0.6B 基准测试抓包中一个样本的实测组成。
基准测试测量了 Miles 使用的完整 flat-dict 交接,并将 Miles Ray 后端与 Mooncake 结构化对象传输进行比较。
PUT 是生产者加载有效载荷后的一次计时后端调用。GET 是一次预热后三次远程消费者试验的平均值。引用序列化和调度器交接在计时区域之外。这些数字涵盖有效载荷传输和重建,不包括端到端训练吞吐量。
在测试的有效载荷大小范围内,Mooncake 使 Miles GET 比 Ray 后端快约 10–14 倍。PUT 提升约 1.2–1.6 倍。
图 5. 对于这种碎片化的 Miles rollout 布局,Mooncake 提供约 10–14 倍更快的 GET。纵轴使用对数刻度。
PUT 收益较小。其计时路径包括 Python 对象遍历、不规则行打包、元数据和清单构建以及有效载荷传输。GET 对此工作负载改进更大,且训练器必须在开始训练步骤之前完成 GET。
Miles 集成为数据路径奠定了基础。下一个优先级是在更广泛的真实世界 RL 工作负载范围内与 Miles 社区一起验证。
验证更多模型和工作负载。我们将把验证扩展到多模态和 VLA 工作负载、agentic RL、世界模型训练以及视频生成或扩散模型的 RL。由于这些工作负载以不同方式组合媒体、轨迹、动作、奖励和中间状态,每种工作负载都将使用真实 rollout 数据和端到端训练进行验证,以确保正确性、兼容性和传输性能。
将 rollout 数据与 KV 缓存工作负载隔离。Mooncake 需要对短寿命 rollout 数据和 KV 缓存数据分别进行计量、配额和驱逐策略,这样 rollout 流量突发就不会驱逐对延迟敏感的缓存条目。
这项工作跨越了多个代码库边界,其审查和验证也是如此。我们感谢 Xinpeng Zhao(@zxpdemonio)、Teng Ma(@stmatengss)、Xingyuan Wu(@yokinoshitayoki)、@fzyzcjy、@guapisolo、Xuchun Shang(@XucSh)、Bo Gao(@Bo-Vincent)、He Zhou(@ehuohz)和 Yufeng He(@he-yufeng)对 Mooncake 实现和审查、Miles 集成和验证以及设计反馈、审查和测试的贡献。
Mooncake 项目:https://github.com/kvcache-ai/Mooncake
Mooncake 文档:https://kvcache-ai.github.io/Mooncake/
Miles Mooncake rollout 传输用户指南:radixark/miles#2535