Google CloudがOpenRLなどのAIネイティブスタックを発表、RLの事後学習や分散推論を支援

OpenRLなど6つの構成要素

Google CloudがKubernetesとGKEにまたがって導入・拡充する統合スタックは、6つの構成要素に基づきます。うち5つはCNCF Sandboxに位置付けられ、役割は以下の通りです。^[1]

llm-d 1.0（CNCF Sandbox）は、オープンな分散推論の制御プレーンです。Agent Substrate（CNCF Sandbox）は、オープンなオーケストレーションランタイムに当たります。

Ray on GKE with KubeRay（CNCF Sandbox）は、分散Rayの計算とデータ処理に向けたフレームワークです。Kueueによるスケジューリングも提供します。

OpenRLの事後学習でできること

OpenRLは、Googleが作成してCNCF Sandboxに提供した、Kubernetesネイティブの事後学習サービスです。ML研究者とプラットフォームチームの間の摩擦をなくすのが狙いです。^[1]

ML研究者は、標準的な報酬関数を使い、PythonとPyTorchで素早く反復したいと考えます。一方でプラットフォームチームには、マルチテナントの統制やクォータの適用、予測可能なスケジューリングが必要です。

図1　OpenRLのアーキテクチャの特徴

研究者は、forward_backwardやoptim_step、sample、save_stateの4つを中心に学習コードをPythonで書きます。OpenRLがKubernetesのスケジューリングの流れへ変換するため、開発者はKubernetesの内部を学んだり、インフラのコードを管理したりする必要がありません。

OpenRLの設定不要のDRA

設定不要の動的リソース割り当て（DRA）では、OpenRLがモデルのコンテキスト長やバッチの次元、LoRAか完全なファインチューニングかを自動でプロファイルします。そのうえで、KubernetesネイティブのDRAクレームを使い、異なる種類のアクセラレータ群へワークロードを動的に詰め込む仕組みです。

図2　OpenRLの設定不要のDRA

OpenRLの学習と推論の連携

OpenRLは、協調的な時分割によって、学習ステップとロールアウト推論を同じ物理アクセラレータ上で交互に実行します。計算中心の学習と応答の遅れに敏感なロールアウト生成で、専用の静的クラスタを分ける必要がありません。

図3　OpenRLが避ける構成との違い

最適化ステップごとに数ギガバイトのモデル重み全体を送ると、クラスタのネットワークが大きく飽和します。OpenRLはパラメーターの差分だけをロールアウト用ワーカーへ送り、ロールアウトのコールドスタートを解消します。

GKEのRay Operatorの動作

Ray on GKE with KubeRay（CNCF Sandbox）とGKEのRay Operatorの関係は、今回の発表に含まれていません。GKEのRay Operatorは、KubeRayに基づき、Rayクラスタやアプリケーションのデプロイと管理に使う機能です。^[1][2]

図4　GKEのRay Operatorの位置付け

GKEクラスタでRay Operatorを有効にすると、GKEがKubeRayオペレーターを自動でインストールしてホストします。手順は「GKE で Ray オペレーターを有効にする」のページにあります。

KubeRayが提供するのは、Kubernetes上のRay Deploymentを管理するカスタムリソースです。例として、RayClusterやRayJob、RayServiceがあります。

図5　KubeRayのカスタムリソースの例

RayClusterカスタムリソースで指定するRayクラスタは、通常1つのヘッドPodと複数のワーカーPodの構成です。RayServiceでは、KubeRayがRayClusterを作成し、Ray Serveの構成どおりにアプリケーションをデプロイします。

RayJobでは、KubeRayがジョブ用のRayClusterと、そのヘッドPodへジョブを送信するKubernetes Jobを作成します。ジョブの正常完了後にRayClusterを自動でクリーンアップするKubeRayの構成は、リソースを効率的に管理するための手段です。

GKEのRay Operator利用時の責任分担

CNCF Sandboxに位置付けられているのは、AIネイティブスタックの6つの要素のうち5つです。該当するのは、OpenRLやllm-d 1.0、Ray on GKE with KubeRay、Agent Substrate、gVisorです。^[1]

図6　GKEのRay Operator利用時の責任分担の一部

Googleが担う責任の範囲

Ray Operatorを使ってGKEでRayワークロードを実行する場合、KubeRayオペレーターの信頼性と稼働時間はGoogleが維持します。バージョンアップグレードの管理と、RayCluster、RayJob、RayServiceカスタムリソースを管理するKubeRay固有の機能も、Googleの責任に含まれます。^[2]

図7　Ray Operator利用時のGoogleの責任

利用者が担う責任の範囲

Ray Operatorを使ってGKEでRayワークロードを実行する場合、利用者側の責任に含まれる項目は以下の通りです。

Ray Operatorを使ってGKEでRayワークロードを実行する場合は、利用者側が担う責任を導入前に確かめます。KubeRayオペレーターの更新はGoogle、RayのヘッドPodとワーカーPodの更新は利用者が担う分担です。

OpenRLなど構成要素の課題別の選び方

Google Cloudは、従来のコンテナ基盤で強化学習の事後学習や状態を持つ分散推論、自律エージェントの実行を拡張する際の課題に、このスタックで対処しています。編集部は、課題ごとに確かめる候補となる構成要素を、以下のように整理しました。^[1]

強化学習と推論で確かめる点

強化学習で学習とロールアウト生成を交互に行うと、GPUsやTPUsが学習の大半の時間で待機します。この状態に当てはまる場合は、OpenRL（CNCF Sandbox）の協調的な時分割が確かめる候補です。

図8　強化学習の待機と確かめる候補

Kubernetes標準のラウンドロビンや最小接続数の振り分けでは、マルチターンのエージェント呼び出しが任意のPodに散らばり、KVキャッシュの局所性が失われます。この場合は、llm-d 1.0（CNCF Sandbox）のプレフィックスキャッシュを考慮したルーティングを確かめます。

図9　KVキャッシュの課題と確かめる候補

クラスタ運用と安全性で確かめる点

独自のRayやSlurmのような分離した計算クラスタをKubernetesと並行して運用すると、容量が分断され、保守の手間が増えます。Kubernetesと並行して独自のRayクラスタを運用している場合は、Ray on GKE with KubeRay（CNCF Sandbox）が確かめる候補です。

図10　並行クラスタの課題とCNCF Sandboxの候補

共有環境でモデルが生成した信頼できないシェルコマンドやスクリプトを実行すると、クラスタがサンドボックスからの脱出や権限昇格の危険にさらされます。マルチテナントのクラスタでこうしたコードを実行する場合は、gVisor（CNCF Sandbox）のシステムコール分離を確かめます。

図11　共有環境の実行の課題と確かめる候補