MooncakeがMilesを強化：断片化したRolloutデータから効率的なバッチI/Oへ

分離型RLシステムにおけるRolloutデータ

大規模言語モデルの強化学習は、Rollout生成とモデル学習という2つの全く異なるワークロードを融合している。Rolloutフェーズでは、推論ワーカーノードが現在のポリシーに基づいてプロンプトを実行しレスポンスを生成するとともに、学習アルゴリズムが必要とする情報（生成トークン、マスク、対数確率、報酬、シーケンス長、サンプル識別子およびその他のメタデータ）を生成する。これらの出力が合わさって、次のステップの学習で消費されるRolloutデータを構成する。

小規模では、生成と学習が実行環境を共有できるが、大規模では、現代のRLシステムは分離アーキテクチャを採用しており、Rollout生成と学習を独立したワーカーグループとしてデプロイし、プロセス・GPU・マシンをまたいで動作させることが多い。

RL Rolloutデータ転送が困難な理由

RL Rolloutデータは、分散学習でよく見られる大規模な単純テンソルとは大きく異なる。1つのRolloutバッチは通常、単一の連続テンソルではなく異種構造化オブジェクトであり、生成トークン、損失マスク、対数確率、報酬、シーケンス長、サンプル識別子、ルーティング情報、メタデータなどの補助フィールドを含む場合がある。これらの値は、テンソル、NumPy配列、Pythonスカラーのリスト、可変長配列、バイト列、または任意のPythonオブジェクトとして表現される。

主な課題は次の通りである：1. 異種データ型と複雑なセマンティクス；2. 高度に断片化されたメモリレイアウト。有効なデータパスは、効率性・正確性・スケーラビリティ・柔軟性・予測可能なハンドオフレイテンシを同時に満たす必要がある。

MooncakeによるMilesへのRolloutデータ転送サポート

Milesは大規模モデルのポストトレーニングを対象とした高性能RLフレームワークであり、高スループットなRollout生成のためのSGLangとスケーラブルな学習のためのMegatron-LMを組み合わせている。Mooncakeは分散AIワークロード向けの高性能データプレーンを提供しており、現在MilesのRolloutデータ転送バックエンドとして統合されている。Milesがキャプチャしたロールアウトデータ上で、リモートGETレイテンシは既存のRayパスと比較して約10〜14倍高速化し、PUTは約1.2〜1.6倍向上している。

Milesが実際に転送する内容

Milesは完成済み辞書によるハンドオフを採用しており、プロデューサーがput(data, type="dict")を呼び出し、学習ワーカーがgetを呼び出して元の辞書を再構築する。

コントロールプレーンとデータプレーンを分離することで、スケジューリングの決定を軽量に保ちながら、バッチRolloutのペイロードを専用データパスを通じて効率的に転送することを実現している。