LMSYS、分散型強化学習システムにおけるロールアウトデータの効率的な大量 I/O 手法「Mooncake」を発表
本文の状態
日本語全文を表示中
詳細モードで約22分の本文を読めます。
同じ出来事の情報源
この情報源を基点に整理
LMSYS Blog
LMSYS Blog は、大規模言語モデルの強化学習においてロールアウト生成とトレーニングを分離する分散アーキテクチャのデータ転送効率化手法「Miles」を発表した。
AI深層分析を開く2026年8月21日 10:15
AI深層分析
キーポイント
分散型RLシステムにおけるワークロードの分離
大規模な強化学習では、推論中心のロールアウト生成とテンソル計算中心のトレーニングを別々のワーカーグループで実行する分散アーキテクチャが採用される。
パイプラインレベルでの並行処理の実現
Miles は、トレーニングがロールアウト N を処理している間に、ロールアウトワーカーが N+1 の生成を並行して行うことでリソース利用率を向上させる。
構造化データ転送の課題と解決
ロールアウトデータは単一の連続テンソルではなく異種構造を持つため、推論側からトレーニング側への効率的な移動がシステム境界におけるボトルネックとなる。
非同期処理による柔軟性の向上
各ステージを独立してスケジューリングできることで、一方のステージが他方を待たずに異なる速度で進行し、柔軟な容量プロビジョニングが可能になる。
ロールアウトデータの複雑な構造と異種性
ロールアウトバッチは生成トークンやメタデータなど多様な型を持つ非連続オブジェクトであり、効率的な転送には物理表現と論理構造の両方を理解する必要がある。
重要な引用
The two stages have fundamentally different resource and execution characteristics.
This separation also enables pipeline-level concurrency.
Efficiently moving this structured rollout data from the inference side to the training side therefore becomes an important part of the end-to-end RL pipeline.
Rollout data can contain a very large number of small memory allocations.
編集コメントを表示
編集コメント
LMSYS が提示する「Miles」の概念は、大規模強化学習の実用化に向けたインフラ課題への鋭い洞察を示している。推論と学習の境界を跨ぐデータ移動の最適化が、次世代RLシステムの性能を左右する鍵となるだろう。
Source Article
元記事を日本語で読む
本文に関係しない購読案内、埋め込み通知、サイト内プロモーションは除いています。
分散型 RL システムにおけるロールアウトデータ
大規模言語モデルに対する強化学習は、ロールアウト生成とモデル学習という、非常に性質の異なる二つのワークロードを組み合わせるものです。
ロールアウト期間中、推論ワーカーは現在のポリシーを用いて一連のプロンプトを実行し、応答を生成します。その過程で、学習アルゴリズムに必要な情報が生成されます。具体的には、生成されたトークン、マスク、対数尤度確率、報酬値、シーケンス長、サンプル識別子、その他のメタデータなどが含まれます。これらすべての出力が一体となり、次の学習ステップで消費されるロールアウトデータを形成します。
小規模なスケールでは、生成と学習は同じ実行環境を共有できます。しかし大規模になると、現代の RL システムではますます分散アーキテクチャを採用する傾向にあります。これは、ロールアウト生成と学習を別々のワーカーグループとしてデプロイする方式で、異なるプロセス、GPU、あるいはマシンにまたがって展開されることも珍しくありません。
この二つの工程は、根本的に異なるリソース要件と実行特性を持っています。ロールアウトは推論集約型のワークロードであり、そのスループットはデコード効率、バッチ処理、およびリクエストスケジューリングに依存します。一方、学習は大規模で高度に同期されたテンソル計算を必要とします。これらを分離することで、両方のワークロードを同じ実行パターンに無理やり合わせるのではなく、それぞれを独立してスケーリングし、スケジュールすることが可能になります。
この分離により、パイプラインレベルの並行処理が可能になります。Miles はロールアウト N のトレーニング中にも、ロールアウトワーカーがすでにロールアウト N+1 を生成している状態を維持できます。非同期 RL システムでは、各オペレーションのたびに一方が他方を待たせるのではなく、ロールアウトとトレーニングのワーカーが異なる速度で進行することを許容します。
この利点は、リソース利用率の向上と、推論・トレーニング容量のプロビジョニング方法における柔軟性の拡大です。ただし、システムを分割することにより新たな境界が生じます。ロールアウトワーカーによって生成されたデータは、トレーニング側で消費される前に別のワーカー群へ移動させる必要があるという点です。
この引き渡しは、生成と次のポリシー更新の間に直接位置します。転送が遅れると、トレーナーがデータの到着を待たされ、ロールアウトワーカー上のメモリが不要なほど長く占有されたままになるリスクがあります。Miles のような分散 RL システムでは、推論側からトレーニング側へ構造化されたロールアウトデータを効率的に移動させることが、エンドツーエンドの RL パイプラインにおいて重要な要素となります。
RL ロールアウトデータの転送が難しい理由
RL のロールアウトデータは、分散トレーニングで一般的に移動される大規模な規則的なテンサーとは大きく異なります。
ロールアウトバッチは、単一の連続したテンソルではなく、通常は異種構造化オブジェクトです。使用するフレームワークやアルゴリズムによって、生成されたトークン、損失マスク、対数尤度、報酬、シーケンス長、サンプル識別子、ルーティング情報、メタデータ、その他の補助フィールドなどが含まれる場合があります。これらの値は、テンソル、NumPy 配列、Python スカラーリスト、サンプルごとの可変長配列、バイト列、あるいは任意の Python オブジェクトとして表現されます。
このデータを効率的に転送する上で特に困難となる性質がいくつかあります。
課題 1: 異種データ型と複雑な意味構造
ロールアウト内の各フィールドは、根本的に異なる表現と意味を持ちます。密なテンソルや数値配列は型付きバッファとして効率的に転送できますが、不規則なシーケンスには行境界の情報が必要であり、スカラーリストでは元の値を保持する必要があり、メタデータや Python オブジェクトはより汎用的なエンコーディングを必要とします。同時に、トレーニング側は RL フレームワークが期待する正確な構造(データ型、形状、行の順序、null 状態、メタデータなど)を再構築しなければなりません。一般的なシリアライザーはこの種のオブジェクトを機能的に扱えますが、その代償として追加の変換、コピー、再構築が必要になることが多々あります。したがって、ロールアウト転送を効率化するには、各フィールドの物理的な表現と論理的な構造の両方を理解する必要があります。
課題 2: 高度に断片化したメモリアウトプットレイアウト
ロールアウトデータには非常に多数の小さなメモリ割り当てが含まれることがあります。キャプチャされた Miles ワークロードでは、トークン、ロースマスク、ロールアウトログプロブなど主要なフィールドは list[np.ndarray] として表現され、各サンプルごとに NumPy アレイが個別に割り当てられています。バッチサイズが大きくなると、これらの断片を個々に転送する際にメモリ登録と Store オペレーションが繰り返されるほか、Python オブジェクト全体をシリアライズするには大規模なオブジェクトグラフの走査、コピー、再構築が必要になります。課題は、元の構造を保ちながら、断片的な論理データを効率的な一括転送に変換することです。
これらの課題を合わせると、ロールアウトデータの移動は単なる帯域幅の問題ではありません。システムは、その構造を保ちつつ、断片的で多様なデータを効率的に移動させる必要があります。
効果的なデータパスは、同時にいくつかの要件を満たす必要があります:
- 効率性: 過度なシリアライズ、コピー、オブジェクト再構築、および割り当てごとの転送オーバーヘッドを避けること。
- 正しさ: ループ全体を通じてフィールド型、形状、行境界、null 状態、メタデータ、Python レベルの値を保持すること。
- スケーラビリティ: サンプル数、オブジェクト断片数、総ペイロードサイズが増加しても性能が低下しないこと。
- 柔軟性: RL フレームワークにネイティブなロールアウト表現を平坦化したり書き換えさせたりせず、多様なフィールドをサポートすること。
予測可能なハンドオフ遅延:トレーニングがロールアウトデータの待機で停止しないよう、バッチを迅速に配信する。
Mooncake を活用した Miles のロールアウトデータ転送
Miles は大規模モデルのポストトレーニング向け高性能強化学習フレームワークです。SGLang を用いた高スループットなロールアウト生成と、Megatron-LM によるスケーラブルなトレーニングを組み合わせ、必要に応じて Hugging Face の実装を直接学習するワークロードには PyTorch FSDP2 バックエンドも提供します。Miles は非同期 RL を完全サポートしており、ロールアウトワーカーとトレーニングワーカーが独立して動作可能。これにより、ループ内の高速な重み更新、エージェント型ロールアウト、低精度学習、大規模生産環境向けのフォールトトレランスといった機能を実現しています。
この分離・非同期設計により、ロールアウトからトレーニングへのデータ経路は RL パイプラインの要となっています。ロールアウトバッチは構造化されつつも異種混合であり、サンプルごとの断片的なデータやフレームワーク固有のメタデータを含むことが多いため、ワークロード規模が拡大するほど効率的な転送と再構築が重要になります。
Mooncake は分散 AI ワークロード向けに高性能なデータプレーンを提供します。RL のロールアウトデータにおいては、このデータプレーンを拡張し、構造化オブジェクトの転送機能を備えることで、異種混合で断片的なロールアウトオブジェクトを、その元の構造と意味を保ったまま移動させることが可能になります。
Mooncake は現在、Miles のロールアウトデータ転送バックエンドとして統合されています。Miles から収集されたロールアウトデータを用いたこの統合により、既存の Ray パスと比較して転送レイテンシが大幅に短縮されました。具体的には、リモート GET が約 10〜14 倍高速化され、PUT も約 1.2〜1.6 倍向上しています。
その結果、強化学習(RL)のプログラミングモデルを変更することなく、ロールアウトからトレーニングへの引き継ぎが高速化されました。これにより、大規模で構造化されたロールアウトワークロードに対して、Miles はより効率的なデータパスを獲得しました。
RL パイプラインにおけるロールアウトデータの移動
Miles は同期および非同期の学習ループの両方をサポートしています。いずれのモードにおいても、ロールアウトワーカーとトレーナーが別プロセスまたは別マシンにデプロイされた場合、完了したロールアウトバッチは次のトレーニングステップで消費される前に、その境界を越える必要があります。
有用な設計原則として、コントロールプレーンとデータプレーンを分離することが挙げられます。フレームワークのスケジューラーはロールアウトバッチがどこへ向かうかを決定しますが、大量のペイロード自体を運ぶべきではありません。代わりに、テンソルペイロードや Store のチャンクレイアウトを除外した軽量な転送参照(transfer reference)を渡します。必要に応じて JSON 対応のメタデータは保持しつつもです。
実際のロールアウトペイロードは別の経路を通ります:
- フレームワークスケジューラーが転送参照を渡す
- Mooncake がペイロードを保存・転送する
- トレーニングワーカーがトレーニングステップを実行する前に、元のオブジェクトを再構築する
- すべてのリーダーが完了した後、フレームワークが一時的な Store オブジェクトを削除する
この分離により、スケジューリングの決定を軽量に保ちつつ、バッチロールアウトペイロードは専用のデータ経路を流れることができます。
Miles が実際に転送するもの
Miles から取得されたロールアウトデータによって、このデータ経路が具体的に定義されます。特定のフレームワーク構成においてはフィールドの契約は安定していますが、各フィールドが共通の便利なメモリ内表現を共有しているわけではありません。これらは大きく以下の 3 つのグループに分類されます:
| フィールドグループ | Miles の表現形式 | 転送の懸念点 |
|---|---|---|
| トークン、マスク、および対数尤度行 | サンプルごとの数値行 | 多数の個別のオブジェクト行;長さはサンプルによって異なる場合があります。 |
| ID、長さ、報酬、およびフラグ | Python スカラーリスト | バイト数は少ないが、トレーニングに必須。 |
| オプションおよびオブジェクトメタデータ | Python 値、テンソル辞書、またはネストされたオブジェクト | 構造体、null 値、および有効な機能のセマンティクスを保持。 |
キャプチャデータにおいて、最初のグループが最も多くのバイトを占めています。他のフィールドはサイズが小さいものの、サンプルの識別情報、長さ、報酬、特徴状態、および管理情報を担っているため、削除や正規化によって無視することはできません。これらのフィールドの正確なサイズはワークロードに依存しますが、ベンチマークセクションでは測定例が示されています。
Miles は完成した辞書形式の手渡しを採用
Miles では、完全なロールアウト辞書が準備された後に現在の手渡しが開始されます。プロデューサーは put(data, type="dict") を呼び出し、スケジューラーが返された参照を引き継ぎ、トレーナーは get を実行して元の辞書を再構築します。
図 1. Miles は完了したロールアウト辞書に対して同期 put と get を使用します。参照はスケジューラーを経由しますが、Mooncake はペイロードを Store データプレーン経由で転送します。
転送呼び出しは同期処理です。プロデューサーは put が完了するまで参照を返さず、トレーナーもオブジェクトを消費する前に get の完了を待ちます。
これは RL パイプライン自体が非同期であることに反するものではありません。Miles は、N 番目のロールアウトのトレーニング中に N+1 番目のロールアウトを生成できます。つまり、並行処理は転送操作よりも上位に位置します。個々のハンドオフは同期的ですが、異なるロールアウトとトレーニングの段階はパイプラインレベルで重複して実行可能です。
Mooncake が Miles のロールアウトデータをどのように保存し転送するか
Mooncake は、2 つの課題を異なるレイヤーで解決します。各フィールドに対して効率的な表現形式を選択できるよう、構造を十分に可視化して維持します。具体的には、テンソルや配列は型情報を保持し、不規則な行(ragged rows)はコンパクトな境界メタデータを含み、Python 値は GET が再構築するために必要な情報を持ち続けます。その後、断片化したメモリをバッチ I/O へ変換します。
コピープランでは、条件を満たす小さな行を再利用可能な登録済み BufferPool チャンクに直接パッキングします。一方、大きな連続的なテンソルや配列は、ネイティブの Store パスを利用できます。これにより、辞書全体を単一の非可視なブロブとしてシリアライズする極端なケースも、小さなアロケーションごとに Store 操作を発行する別の極端なケースも回避されます。
生成されたペイロードメンバーはバンドルとして公開されます。そのマニフェストには、各メンバーの保存場所と結合方法が記録され、ペイロードが完成した後にのみ参照可能になります。GET はこの記述を逆順に追って、元の Miles 辞書を取得・再構築します。Miles は引き続き小さなパブリックインターフェイス put(data, type="dict") と get(ref, type="dict") を使用しますが、スキーマ、フィールドレイアウト、転送プラン、および再構築ロジックはすべて Mooncake の内部に保持されます。
図 2:スキーマとリーフ展開により、各フィールドの型と構造が明確になります。フィールド固有のエンコーディングによって、型付きペイロードメンバーと再構築メタデータが生成され、最後に Bundle Store がそのマニフェストを公開します。対象となる断片化転送は BufferPool をバックアップとするステージング領域を使用し、GET 操作も同様の構造を逆順で実行します。
フィールドごとにレイアウトを選択する
Mooncake はまず、効率的にエンコード可能なリーフへ辞書を展開します。配列とテンソルは型を保持したまま、不規則な行(ragged rows)はコンパクトな境界メタデータを付与され、サポート対象の Python 値は再構築に必要な情報を維持されます。曖昧さやパフォーマンスが敏感に依存するフィールドについては、フレームワークが提供するスキーマによって保存形式を固定します。スキーマが指定されていない場合、Mooncake は観測された値から自動的に推論してスキーマを生成します。
この選択基準はマルチモーダルデータにも適用されます。処理済みのピクセルや関連するモデル入力は、テンソルパスのまま維持できます。PIL 画像、エンコード済み PNG または JPEG のバイト列、可変長のメディアリストについては、それぞれの境界と再構築メタデータを保持しつつ、すべての表現を同一のシリアライザーに通す必要がないよう、メディア対応型のレイアウトを採用します。
断片化されたメモリをバッチ I/O に変換
行ごとにデータを転送すると、数千回の登録と Store リクエストが発生します。フィールド全体を先に連結すればリクエスト数は減りますが、フルサイズの一時領域が必要になります。Mooncake は代わりにコピー計画を作成し、必要な時に再利用可能な登録済みの BufferPool チャンクにデータを埋めていきます。ネイティブパスでは、対象となる数値行を直接これらのチャンクへ転送するため、Python の行ループや一時連結が不要です。Store が対応している場合、大規模な連続配列とテンソルも直接またはネイティブパスを使用します。
完全なバンドルの公開と直接再構築
Mooncake は、すべてのペイロードとメタデータが準備完了した後にのみバンドルマニフェストを公開するため、読者が不完全な Miles ディクトを見ることはありません。GET リクエストでは、マニフェストが取得すべきメンバを特定し、構造化されたメタデータが各フィールドの再構築方法を記述します。対象となる読み取り操作は中間 bytes オブジェクトを経由せず、BufferPool 対応先の宛先に直接アクセスできます。また、型付きの不規則な行(ragged rows)も結果バッファを直接参照可能です。
トレーニングプロセスでは、使用済みのローカル BufferPool 対応の結果が解放されます。すべての読者が完了した後、Miles は短期間の Store オブジェクトを削除し、Mooncake がペイロードチャンクとマニフェストを回収します。
パフォーマンス結果
フレームワークの構成を固定した場合、モデルを変更してもフィールド契約自体が変化するわけではありません。転送サイズは、サンプル数やプロンプト・レスポンスの長さ、教師の対数確率やルーティング情報、マルチモーダル入力といった任意のフィールドの数に依存します。
ベンチマークペイロード
Miles は、数学的なプロンプトに対して Qwen3-0.6B を使用してソースデータを生成しました(rollout_id=0、8 サンプル)。このキャプチャ内のすべてのレスポンスは 256 トークンです。ベンチマークでは、3 つの主要な数値フィールドを型付き ndarray の行として正規化しましたが、その値や行長、データ型、サンプルごとの断片化状態は保持しました。
これらのサンプルを平均化した結果、論理的なペイロードは以下のようになります:
| キャプチャされたサンプルの一部 | 計算 | 論理バイト数 |
|---|---|---|
tokens | 平均プロンプト+レスポンストークン数 x 4 B (int32) | ~1,338 B |
loss_masks | 256 エントリ x 4 B (int32) | 1,024 B |
rollout_log_probs | 256 エントリ x 4 B (float32) | 1,024 B |
| スカラーおよびオブジェクトフィールド | ID、長さ、報酬、フラグ、バージョンメタデータ | ~36 B |
| 合計 | ~3,422 B |
最初の3つのフィールドは、この特定のサンプルレイアウトにおいて約3,386バイトを占め、全体の98.9%に相当します。より大きなベンチマークペイロードでは、これら8種類のキャプチャされたサンプルが繰り返し配置され、フィールドタイプと断片化された割り当てパターンは維持されたまま、論理的なサンプル数が増加します。この計算は特定のキャプチャを記述するものであり、Milesのサンプルサイズが固定されていることを定義するものではありません。
転送結果
本ベンチマークは、Milesが使用する完了したフラット辞書のハンドオフを測定し、Miles の Ray バックエンドと Mooncake による構造化オブジェクト転送を比較します。
PUT は、プロデューサーがペイロードを読み込んだ後の1回のタイム計測付きバックエンド呼び出しです。GET は、ウォームアップ後に3回のリモートコンシューマー試行の平均値です。参照シリアライゼーションやスケジューラのハンドオフは、この計測範囲に含まれません。これらの数値はペイロード転送と再構築に焦点を当てたものであり、エンドツーエンドのトレーニングスループット全体を対象としたものではありません。
テストしたすべてのペイロードサイズにおいて、Mooncake を使用すると Miles の GET が Ray バックエンドより約10〜14倍高速化されます。PUT についても、約1.2〜1.6倍の改善が見られました。
PUT の効果は比較的小さくなります。これは、タイムパスに Python オブジェクトの走査、不規則な行のパッキング、メタデータとマニフェストの構築、ペイロード転送が含まれるためです。一方、このワークロードでは GET の改善幅がより大きくなり、トレーニングステップを開始する前にトレーナーはこれを完了させる必要があります。
次のステップ
Miles 統合により基本となるデータパスが確立されました。次の優先事項は、より多様な実世界の強化学習(RL)ワークロードにおいて Miles コミュニティと連携してこのパスを検証することです。
- モデルとワークロードの検証範囲を拡大する。 次期検証では、マルチモーダルおよび VLA ワークロード、エージェント型 RL、ワールドモデルのトレーニング、動画生成や拡散モデル向けの RL などに対応します。これらのワークロードはメディア、軌跡、行動、報酬、中間状態などを異なる組み合わせで利用するため、各ワークロードについて実ロールアウトデータを用いた検証と、正確性・互換性・転送性能を確認するためのエンドツーエンドのトレーニングを実施します。
- ロールアウトデータを KV キャッシュワークロードから分離する。 Mooncake では、短寿命のロールアウトデータと KV キャッシュデータに対して、それぞれ独立した集計管理、クォータ設定、およびエビクションポリシーを適用する必要があります。これにより、ロールアウトトラフィックの一時的な急増が、レイテンシに敏感なキャッシュエントリの削除を引き起こさないようにします。
謝辞
本稿は複数のリポジトリにまたがる内容であり、そのレビューと検証も同様に複数機関を跨いで行われました。Mooncake の実装およびレビュー、Miles への統合と検証、設計フィードバック・レビュー・テストにおいて貢献いただいた方々に対し、心より感謝申し上げます。
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 プロジェクト:https://github.com/kvcache-ai/Mooncake
- Mooncake ドキュメント:https://kvcache-ai.github.io/Mooncake/
- Miles と Mooncake のロールアウト転送ユーザーガイド:radixark/miles#2535
原文を表示
Rollout Data in Disaggregated RL Systems
Reinforcement learning for large language models combines two very different workloads: rollout generation and model training.
During rollout, inference workers run the current policy on a set of prompts and generate responses. Along the way, they produce the information required by the learning algorithm, including generated tokens, masks, log probabilities, rewards, sequence lengths, sample identifiers, and other metadata. Together, these outputs form the rollout data that will be consumed by the next training step.
At small scale, generation and training can share the same execution environment. At larger scale, however, modern RL systems increasingly adopt a disaggregated architecture, where rollout generation and training are deployed as separate worker groups, often across different processes, GPUs, or machines.
The two stages have fundamentally different resource and execution characteristics. Rollout is an inference-heavy workload whose throughput depends on decoding efficiency, batching, and request scheduling, while training relies on large, highly synchronized tensor computations. Separating them allows each side to be scaled and scheduled independently instead of forcing both workloads into the same execution pattern.
This separation also enables pipeline-level concurrency. Miles can train on rollout *N* while rollout workers are already generating rollout *N+1*. Asynchronous RL systems can therefore allow rollout and training workers to progress at different rates rather than making one stage wait for the other after every operation.
The benefit is better resource utilization and greater flexibility in how inference and training capacity are provisioned. But disaggregation also introduces a new systems boundary: the data produced by rollout workers must now move to a different set of workers before training can consume it.
That handoff sits directly between generation and the next policy update. A slow transfer can leave trainers waiting for data and keep memory occupied on rollout workers longer than necessary. For a distributed RL system such as Miles, efficiently moving this structured rollout data from the inference side to the training side therefore becomes an important part of the end-to-end RL pipeline.
What Makes RL Rollout Data Transfer Challenging
RL rollout data differs significantly from the large, regular tensors commonly moved in distributed training.
A rollout batch is usually a heterogeneous structured object rather than a single contiguous tensor. Depending on the framework and algorithm, it may contain generated tokens, loss masks, log probabilities, rewards, sequence lengths, sample identifiers, routing information, metadata, and other auxiliary fields. These values can be represented as tensors, NumPy arrays, Python scalar lists, variable-length per-sample arrays, bytes, or arbitrary Python objects.
Several properties make this data particularly challenging to move efficiently.
Challenge 1: Heterogeneous Data Types and Complex Semantics
Different rollout fields have fundamentally different representations and semantics. Dense tensors and numeric arrays can be transferred efficiently as typed buffers, while ragged sequences need row-boundary information, scalar lists need their original values preserved, and metadata or Python objects may require more general encoding. At the same time, the trainer must reconstruct the exact structure expected by the RL framework, including dtype, shape, row order, null state, and metadata. A generic serializer can handle these objects functionally, but often at the cost of extra conversion, copying, and reconstruction. Efficient rollout transfer therefore needs to understand both the physical representation and the logical structure of each field.
Challenge 2: Highly Fragmented Memory Layout
Rollout data can contain a very large number of small memory allocations. In the captured Miles workload, major fields such as tokens, loss_masks, and rollout_log_probs are represented as list[np.ndarray], with one NumPy array allocated for each sample. As batch size grows, transferring these fragments individually introduces repeated memory registration and Store operations, while serializing the full Python object requires walking, copying, and rebuilding a large object graph. The challenge is to turn fragmented logical data into efficient bulk transfers without losing its original structure.
Together, these challenges make rollout data movement more than a bandwidth problem: the system must efficiently move fragmented, heterogeneous data while preserving its structure.
An effective data path needs to satisfy several requirements at the same time:
- Efficiency: avoid excessive serialization, copying, object reconstruction, and per-allocation transfer overhead.
- Correctness: preserve field types, shapes, row boundaries, null state, metadata, and Python-level values through the round trip.
- Scalability: continue to perform well as the number of samples, object fragments, and total payload size increase.
- Flexibility: support heterogeneous fields without forcing the RL framework to flatten or rewrite its native rollout representation.
- Predictable handoff latency: deliver the batch quickly enough that the trainer does not stall waiting for rollout data.
Powering Miles Rollout Data Transfer with Mooncake
Miles is a high-performance reinforcement learning framework for large-scale model post-training. It combines SGLang for high-throughput rollout generation with Megatron-LM for scalable training, and also provides a PyTorch FSDP2 backend for workloads that prefer to train Hugging Face model implementations directly. Miles supports fully asynchronous RL, where rollout and training workers are decoupled and can progress independently, together with features such as fast in-loop weight updates, agentic rollout, low-precision training, and fault tolerance for large-scale production RL workloads.
This disaggregated and asynchronous design makes the rollout-to-training data path a critical part of the RL pipeline. Rollout batches are structured and heterogeneous, often containing fragmented per-sample data and framework-specific metadata, making efficient transfer and reconstruction increasingly important as workload scale grows.
Mooncake provides a high-performance data plane for distributed AI workloads. For RL rollout data, it extends this data plane with structured-object transfer, allowing heterogeneous and fragmented rollout objects to be moved while preserving their original structure and semantics.
Mooncake has now been integrated into Miles as a rollout data-transfer backend. On rollout data captured from Miles, the integration delivers substantially lower transfer latency than the existing Ray path: remote GET is roughly 10–14× faster, while PUT improves by about 1.2–1.6×.
The result is a faster rollout-to-training handoff without changing the RL programming model, providing Miles with a more efficient data path for large, structured rollout workloads.
How Rollout Data Moves Through the RL Pipeline
Miles supports both synchronous and asynchronous training loops. In either mode, once rollout workers and trainers are deployed in separate processes or on separate machines, each completed rollout batch must cross that boundary before it can be consumed by the next training step.
A useful design principle is to separate the control plane from the data plane. The framework scheduler decides where a rollout batch should go, but it should not carry the bulk payload itself. Instead, it passes a lightweight transfer reference that excludes tensor payloads and Store chunk layouts, while still allowing JSON-safe metadata when needed. The actual rollout payload takes a separate path:
- the framework scheduler passes the transfer reference;
- Mooncake stores and transfers the payload;
- the training worker reconstructs the original object before running the training step;
- after all readers finish, the framework removes the short-lived Store object.
This separation keeps scheduling decisions lightweight while allowing the bulk rollout payload to move through a dedicated data path.
What Miles Actually Transfers
The rollout data captured from Miles makes this data path concrete. For a given framework configuration, the field contract is stable, but the fields do not share one convenient in-memory representation. They fall into three broad groups:
| Field Group | Miles Representation | Transfer Concern |
|---|---|---|
| Token, mask, and log-probability rows | Per-sample numeric rows | Many separate row objects; lengths may differ by sample. |
| IDs, lengths, rewards, and flags | Python scalar lists | Small in bytes but required by training. |
| Optional and object metadata | Python values, tensor dictionaries, or nested objects | Preserve structure, nulls, and enabled-feature semantics. |
The first group carries most of the bytes in our capture. The other fields are smaller, but they cannot be discarded or normalized away: they carry sample identity, lengths, rewards, feature state, and bookkeeping information. Their exact size depends on the workload; the benchmark section gives one measured example.
Miles Uses a Completed-Dict Handoff
In Miles, the current handoff begins after the complete rollout dictionary is ready. The producer calls put(data, type="dict"), the scheduler carries the returned reference, and the trainer calls get to reconstruct the original dictionary.
The individual transfer calls are synchronous. The producer returns the reference only after put completes, and the trainer waits for get to finish before consuming the object.
This does not prevent the RL pipeline itself from being asynchronous. Miles can generate rollout *N+1* while training on rollout *N*. In other words, the concurrency sits above the transfer operation: each individual handoff is synchronous, while different rollout and training stages can overlap at the pipeline level.
How Mooncake Preserves and Transfers Miles Rollout Data
Mooncake tackles the two challenges at different layers. It keeps the structure visible long enough to choose an efficient representation for each field: tensors and arrays stay typed, ragged rows carry compact boundary metadata, and Python values retain what GET needs to rebuild them. It then turns fragmented memory into bulk I/O. A copy plan packs eligible small rows directly into reusable, registered BufferPool chunks, while large contiguous tensors and arrays can use native Store paths. This avoids both extremes: serializing the entire dict as one opaque blob or issuing a Store operation for every small allocation.
The resulting payload members are published as a bundle. Its manifest records where those members live and how they fit together, and becomes visible only after the payload is ready. GET follows that description in reverse to fetch and reconstruct the original Miles dict. Miles still uses the small public interface, put(data, type="dict") and get(ref, type="dict"); the schema, field layouts, transfer plan, and reconstruction remain inside Mooncake.
Choose a Layout for Each Field
Mooncake first expands the dict into leaves that can be encoded efficiently. Arrays and tensors remain typed, ragged rows carry compact boundary metadata, and supported Python values retain the information needed to reconstruct them. A framework-provided schema fixes the stored representation for ambiguous or performance-sensitive fields; when no schema is supplied, Mooncake infers one from the observed values.
The same choice applies to multimodal data. Processed pixels and related model inputs can stay on the tensor path. PIL images, encoded PNG or JPEG bytes, and variable-length media lists use media-aware layouts that preserve their boundaries and reconstruction metadata without forcing every representation through the same serializer.
Turn Fragmented Memory into Bulk I/O
Sending every row separately creates thousands of registrations and Store requests. Concatenating a whole field first avoids that request count, but adds a full-size temporary. Mooncake instead builds a copy plan and fills reusable, registered BufferPool chunks as they are needed. Its native path copies eligible numeric rows directly into those chunks, avoiding Python row loops and temporary concatenations. Large contiguous arrays and tensors still use direct or native paths when the Store supports them.
Publish Complete Bundles and Rebuild Directly
Mooncake publishes the bundle manifest only after all payloads and metadata are ready, so a reader never sees a half-written Miles dict. On GET, the manifest identifies the members to fetch and the structured metadata describes how to rebuild each field. Eligible reads can target BufferPool-backed destinations without an intermediate bytes object, and typed ragged rows can view the result buffer directly.
The trainer releases its local BufferPool-backed result after use. Once all readers finish, Miles removes the short-lived Store object and Mooncake reclaims its payload chunks and manifest.
Performance Results
For a fixed framework configuration, switching models does not by itself change the field contract. Transfer size instead follows the number of samples, prompt and response lengths, and optional fields such as teacher log probabilities, routing information, or multimodal inputs.
Benchmark Payload
Miles generated the source data with Qwen3-0.6B on math prompts (rollout_id=0, 8 source samples). Every response in this capture has 256 tokens. For the benchmark, the three large numeric fields were normalized to typed ndarray rows while preserving their values, row lengths, dtypes, and per-sample fragmentation. Averaged across those samples, the logical payload breaks down as follows:
| Part of One Captured Sample | Calculation | Logical Bytes |
|---|---|---|
tokens | Average prompt-plus-response token count x 4 B (int32) | ~1,338 B |
loss_masks | 256 entries x 4 B (int32) | 1,024 B |
rollout_log_probs | 256 entries x 4 B (float32) | 1,024 B |
| Scalar and object fields | IDs, lengths, rewards, flags, and version metadata | ~36 B |
| Total | ~3,422 B |
The first three fields account for about 3,386 bytes, or 98.9% of this particular sample layout. Larger benchmark payloads repeat the eight captured samples, preserving their field types and fragmented allocation pattern while increasing the logical sample count. The calculation describes this capture; it does not define a fixed Miles sample size.
Transfer Results
The benchmark measures the completed flat-dict handoff used by Miles and compares the Miles Ray backend with Mooncake structured-object transfer.
PUT is one timed backend call after the producer loads the payload. GET is the mean of three remote-consumer trials after one warmup. Reference serialization and scheduler handoff are outside the timed region. These numbers cover payload transfer and reconstruction, not end-to-end training throughput.
Across the tested payload sizes, Mooncake makes Miles GET roughly 10–14x faster than the Ray backend. PUT improves by about 1.2–1.6x.
PUT shows a smaller gain. Its timed path includes Python-object traversal, ragged-row packing, metadata and manifest construction, and payload transfer. GET improves more for this workload, and the trainer must finish it before starting the training step.
What Comes Next
The Miles integration establishes the basic data path. The next priority is to validate it with the Miles community across a wider range of real-world RL workloads.
- Validate more models and workloads. We will expand validation to multimodal and VLA workloads, agentic RL, world-model training, and RL for video-generation or diffusion models. Because these workloads combine media, trajectories, actions, rewards, and intermediate state in different ways, each workload will be validated with real rollout data and end-to-end training for correctness, compatibility, and transfer performance.
- Isolate rollout data from KV cache workloads. Mooncake needs separate accounting, quotas, and eviction policy for short-lived rollout data and KV cache data, so a burst of rollout traffic cannot evict latency-sensitive cache entries.
Acknowledgements
This work crossed several repository boundaries, and so did its review and validation. We thank Xinpeng Zhao (@zxpdemonio), Teng Ma (@stmatengss), Xingyuan Wu (@yokinoshitayoki), @fzyzcjy, @guapisolo, Xuchun Shang (@XucSh), Bo Gao (@Bo-Vincent), He Zhou (@ehuohz), and Yufeng He (@he-yufeng) for their contributions across Mooncake implementation and review, Miles integration and validation, and design feedback, review, and testing.
- Mooncake project: https://github.com/kvcache-ai/Mooncake
- Mooncake documentation: https://kvcache-ai.github.io/Mooncake/
- Miles Mooncake rollout transfer user guide: radixark/miles#2535
関連記事
今日のまとめ
AIデイリーブリーフで今日の重要ニュースをまとめ読み