Mooncake赋能Miles:从碎片化Rollout数据到高效批量I/O

Mooncake赋能Miles:从碎片化Rollout数据到高效批量I/O

分离式RL系统中的Rollout数据

大型语言模型的强化学习融合了两种截然不同的工作负载:Rollout生成模型训练。Rollout阶段,推理工作节点基于当前策略运行提示并生成响应,同时产生学习算法所需信息,包括生成token、掩码、对数概率、奖励、序列长度、样本标识符及其他元数据。这些输出共同构成供下一步训练消费的Rollout数据

小规模时,生成与训练可共享执行环境;大规模下,现代RL系统采用分离架构,Rollout生成与训练部署为独立工作组,常跨进程、GPU或机器运行。

Synchronous rollout transfer for Miles

RL Rollout数据传输为何具有挑战性

RL Rollout数据与分布式训练中常见的大规则张量显著不同。一个Rollout批次通常是异构结构化对象,而非单一连续张量,可能包含生成token、损失掩码、对数概率、奖励、序列长度、样本标识符、路由信息、元数据等辅助字段。这些值可表示为张量、NumPy数组、Python标量列表、变长数组、字节或任意Python对象。

Mooncake structured-object transfer architecture

主要挑战包括:1. 异构数据类型与复杂语义;2. 高度碎片化的内存布局。有效数据路径需同时满足效率、正确性、可扩展性、灵活性与可预测的交接延迟。

Rollout data challenges and Mooncake optimizations

Mooncake为Miles提供Rollout数据传输支持

Miles是面向大规模模型后训练的高性能RL框架,结合SGLang实现高吞吐Rollout生成与Megatron-LM实现可扩展训练。Mooncake为分布式AI工作负载提供高性能数据平面,现已集成至Miles作为Rollout数据传输后端。在Miles捕获的Rollout数据上,远程GET延迟较现有Ray路径快约10–14倍,PUT提升约1.2–1.6倍

Miles rollout object anatomy

Miles实际传输的内容

Miles采用已完成字典交接,生产者调用put(data, type="dict"),训练工作节点调用get重建原始字典。

Miles GET latency benchmark

控制平面与数据平面分离,确保调度决策轻量,同时批量Rollout负载通过专用数据路径高效移动。