GKE で LLM の大規模な強化学習（RL）を実行する

※この投稿は米国時間 2025 年 11 月 11 日に、Google Cloud blog に投稿されたものの抄訳です。

大規模言語モデル（LLM）の進化に伴い、強化学習（RL）は、強力なモデルを人間の好みや複雑なタスクの目標に合わせるための重要な手法になりつつあります。

しかし、LLM の RL を実装してスケーリングする必要がある企業は、インフラストラクチャの課題に直面しています。主な課題としては、複数の大規模モデル（アクター、クリティック、報酬、参照モデルなど）を同時にホストすることによるメモリ競合、高レイテンシの推論生成における反復的な切り替え、そして高スループットのトレーニング フェーズが挙げられます。

本ブログでは、カスタム TPU ハードウェアから GKE オーケストレーション層まで、Google Cloud のフルスタック統合アプローチを詳しく解説し、大規模 RL におけるハイブリッドかつ重要な要求への対応方法を紹介します。

クイックガイド: LLM における強化学習（RL）

RL は、トレーニングと推論の両方の要素を組み合わせた継続的なフィードバック ループです。LLM の RL ループの概要は次のとおりです。

1. LLM が、指定されたプロンプトに対する回答を生成します。

2. 「報酬モデル」（多くの場合、人間の好みに基づいてトレーニングされる）は、出力に定量的なスコア、つまり報酬を割り当てます。

3. RL アルゴリズム（例: DPO、GRPO）は、この報酬シグナルを使用して LLM のパラメータを更新し、そのポリシーを調整して、その後のインタラクションでより高い報酬が得られる出力を生成します。

この生成、評価、最適化により、事前定義された目標に基づいて LLM のパフォーマンスが継続的に向上します。

RL ワークロードはハイブリッドで循環的です。RL の主な目標は、エラーの最小化（トレーニング）や高速な予測（推論）ではなく、反復的なインタラクションを通じて報酬を最大化することです。RL ワークロードの主な制約は計算能力だけでなく、システム全体の効率性にもあります。具体的には、エンドツーエンドのステップ時間を効率化するために、サンプラー全体の合計レイテンシを最小限に抑え、重みのコピー速度を最大化することが求められます。

Google Cloud の RL に対するフルスタック アプローチ

システム全体の課題を解決するには、統合的なアプローチが必要です。高速なハードウェアや優れたオーケストレーターだけでは不十分で、スタックのすべてのレイヤが連携する必要があります。RL の特定のニーズを解決するために構築された Google のフルスタック アプローチは次のとおりです。

1. 柔軟で高性能なコンピューティング（TPU と GPU）: お客様を 1 つのパスに固定するのではなく、2 つの高性能オプションを提供します。Google のTPU スタックは、垂直統合された JAX ネイティブのソリューションです。行列演算に優れたカスタム ハードウェアが、ポストトレーニング ライブラリ（MaxText と Tunix）と共同設計されています。並行して、Google は NVIDIA GPU エコシステムを全面的にサポートし、最適化された NeMo RL レシピについて NVIDIA と提携しているため、お客様は既存の専門知識を GKE で直接活用できます。

2. 包括的なフルスタックの最適化: ベアメタルから上位レイヤまで最適化を統合します。これには、Google のカスタム TPU アクセラレータ、高スループット ストレージ（Managed Lustre、Google Cloud Storage）、そして何よりも重要な GKE が提供するオーケストレーションとスケジューリングが含まれます。スタック全体を最適化することで、ハイブリッド RL ワークロードのボトルネックとなるシステム全体のレイテンシに対処できます。

3. オープンソースのリーダーシップ: RL インフラストラクチャは複雑で、幅広いツール上に構築されています。Google のリーダーシップは、Kubernetes のオープンソース化から始まり、Ray などのオーケストレーターとの積極的なパートナーシップにまで及んでいます。Google は、vLLM などの主要なプロジェクトに貢献し、費用対効果の高いサービングのための llm-d などのオープンソース ソリューションを開発し、独自の高性能 MaxText および Tunix ライブラリをオープンソース化しています。これにより、単一のベンダーのツールだけでなく、作業に最適なツールを統合できます。

4. 実績のあるメガスケールのオーケストレーション: トレーニング後の RL には、事前トレーニングに匹敵するコンピューティング リソースが必要になる場合があります。これには、大規模な分散ジョブを単一のユニットとして管理できるオーケストレーション レイヤが必要です。GKE AI メガクラスタは現在最大 65,000 ノードをサポートしており、Google は単一クラスタの制限を超えて RL ワークロードをスケーリングするために、MultiKueue などのマルチクラスタ ソリューションに多額の投資を行っています。

GKE で RL ワークロードを実行する

既存の GKE インフラストラクチャは、要求の厳しい RL ワークロードに最適であり、インフラストラクチャ レベルでさまざまな効率性を提供します。

以下の画像は、RL を大規模に実装するためのアーキテクチャと主な推奨事項の概要を示しています。

基盤となるインフラストラクチャ レイヤは、サポートされているコンピューティング タイプ（CPU、GPU、TPU）などの基礎的なハードウェアを提供します。Run:ai モデル ストリーマーを使用すると、3 つのコンピューティング タイプすべてでモデル ストリーミングを高速化できます。高性能ストレージ（Managed Lustre、Cloud Storage）を RL のストレージ ニーズに使用できます。

中間レイヤは、GKE を利用したマネージド K8s レイヤです。このレイヤは、リソースのオーケストレーション、リソースの入手可能性（Spot または Dynamic Workload Scheduler を使用）、自動スケーリング、プレースメント、ジョブのキューイングやスケジューリングなどにメガスケールで対応します。

最後に、オープン フレームワーク レイヤが GKE 上で実行され、アプリケーションと実行環境が提供されます。これには、安全な分離されたタスク実行のための KubeRay、Slurm、gVisor サンドボックスなどのオープンソース ツールのマネージド サポートが含まれます。

RL ワークフローの構築

RL ワークロードを作成する前に、まず明確なユースケースを特定する必要があります。目標を定義したら、アルゴリズム（DPO、GRPO など）、モデルサーバー（vLLM、SGLang など）、ターゲット GPU/TPU ハードウェア、その他の重要な構成を選択して、コア コンポーネントを設計します。

次に、Workload Identity、GCS Fuse、DGCM 指標で構成された GKE クラスタをプロビジョニングできます。堅牢なバッチ処理を行うには、Kueue と JobSet の API をインストールします。この GKE スタックの上にオーケストレーターとして Ray をデプロイすることをおすすめします。そこから、Nemo RL コンテナを起動し、GRPO ジョブ用に構成して、実行のモニタリングを開始できます。詳細な実装手順とソースコードについては、こちらのリポジトリをご覧ください。

RL のスタートガイド

1. GPU で RL を実行する: GRPO アルゴリズムで MaxText と Pathways を使用する場合は TPU で RL レシピを試し、GPU を使用する場合は NemoRL レシピをお試しください。

2. オープンソース エコシステムとの連携: Google の AI におけるリーダーシップは、Kubernetes、llm-d、Ray、MaxText、Tunix などのオープン スタンダードの上に構築されています。Google と連携して、AI の未来を共に築きましょう。ぜひ llm‑d で開発にご協力ください。llm-d コミュニティに参加し、GitHub のリポジトリをチェックして、オープンソース LLM サービスの今後の発展に貢献しましょう。

-シニア プロダクト マネージャー Poonam Lamba 

-ソフトウェア エンジニア、Bogdan Berce