GKE の NVIDIA NeMo RL を使用して強化学習をファインチューニングしてスケーリングする

このチュートリアルでは、Google Kubernetes Engine（GKE）で強化学習（RL）用の分散トレーニング環境をオーケストレートする方法について説明します。Ray と NVIDIA NeMo RL フレームワークを使用して、モデルをファインチューニングするための分散トレーニング環境を設定します。

このチュートリアルでは、Ray と NeMo RL を使用した GKE のグループ相対ポリシー最適化（GRPO）トレーニング パイプラインについて説明します。GRPO は、モデルの推論能力を向上させるように設計された強化学習アルゴリズムです。このメモリ効率の高いアルゴリズムは、Critic（値モデル）を排除し、相対グループベースの計算を使用することで、RL プロセスを簡素化します。

このチュートリアルを実行する前に、GKE で verl を使用して強化学習をファインチューニングしてスケーリングするチュートリアルを完了することをおすすめします。次のチュートリアルでは、ファインチューニングとスケーリング RL with verl チュートリアルと同じクラスタ設定と構成を使用します。

背景

以降のセクションでは、このチュートリアルで使用するコンセプトの概要を説明します。

強化学習（RL）

RL は、静的な模倣ではなく、経験、探索、フィードバックを通じてモデルを学習させます。事前トレーニングではモデルに何を言うべきかを学習させますが、人間からのフィードバックを用いた強化学習（RLHF）では、有用性、安全性、論理性を学習させます。RL は、ベースモデルと特殊なユースケースのファイン チューニングされたモデルの間の橋渡しとして機能します。

詳細については、強化学習とはをご覧ください。

グループ相対ポリシーの最適化（GRPO）

DeepSeek によって普及したアルゴリズムである GRPO は、Critic モデルを削除することで、LLM アライメント用の Proximal Policy Optimization（PPO）に代わるメモリ効率の高い方法を提供します。GRPO は、Critic ネットワークの代わりに、同じプロンプトに対するレスポンスのグループを生成し、そのグループの平均報酬をベースラインとして使用します。

詳細については、GRPO をご覧ください。

NVIDIA NeMo RL

NeMo RL は、スケーラブルな RL 用に設計された NVIDIA のオープンソースの事後トレーニング ライブラリです。より広範な NeMo フレームワーク エコシステムの一部である NeMo RL を使用すると、単一の GPU で小規模な実験を行うことも、数千の GPU にわたるマルチノード デプロイを行うこともできます。

詳細については、NVIDIA NeMo RL をご覧ください。

GSM8k データセット

このチュートリアルでは、8,500 個の高品質で言語的に多様な小学校の数学の文章問題を含む GSM8k データセットを使用します。

GSM8k と GRPO を使用することで、モデルは同じ問題に対して n 個の異なる回答を生成します。GRPO は、これらの回答をグループの平均と比較します。モデルは、グループの他の回答と比較して、一貫して正しく論理的に妥当なパスに対してより多くの報酬を受け取ります。時間の経過とともに、モデルは手順を明確に説明することが報酬を最大化する最も信頼性の高い方法であることを学習し、パフォーマンスの低い回答に対する報酬を効果的に減らします。

詳細については、GSM8k をご覧ください。

目標

このチュートリアルでは、次の手順に沿って NeMo RL を使用して GKE で RL を設定する方法について説明します。

  1. 環境を準備します。
  2. B200 または H200 GPU を使用して GKE クラスタを設定します。
  3. 分散 Ray クラスタを管理するように KubeRay を構成します。
  4. 高性能ストレージには Managed Lustre を使用します。
  5. NeMo RL を使用する GRPO トレーニング ジョブを実行します。

始める前に

  • Google Cloud アカウントにログインします。 Google Cloudを初めて使用する場合は、 アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

  • Google Cloud CLI をインストールします。

  • 外部 ID プロバイダ（IdP）を使用している場合は、まず連携 ID を使用して gcloud CLI にログインする必要があります。

  • gcloud CLI を初期化するには、次のコマンドを実行します。

    gcloud init

  • Google Cloud プロジェクトを作成または選択します。

    プロジェクトの選択または作成に必要なロール

    • プロジェクトを選択する: プロジェクトの選択に特定の IAM ロールは必要ありません。ロールが付与されているプロジェクトであれば、どのプロジェクトでも選択できます。
    • プロジェクトを作成する: プロジェクトを作成するには、resourcemanager.projects.create 権限を含むプロジェクト作成者ロール（roles/resourcemanager.projectCreator）が必要です。詳しくは、ロールを付与する方法をご覧ください。

    • Google Cloud プロジェクトを作成します。

      gcloud projects create PROJECT_ID

      PROJECT_ID は、作成する Google Cloud プロジェクトの名前に置き換えます。

    • 作成した Google Cloud プロジェクトを選択します。

      gcloud config set project PROJECT_ID

      PROJECT_ID は、 Google Cloud プロジェクトの名前に置き換えます。

  • Google Cloud プロジェクトに対して課金が有効になっていることを確認します。

  • 必要な API を有効にします。

    API を有効にするために必要なロール

    API を有効にするには、serviceusage.services.enable 権限を含む Service Usage 管理者 IAM ロール（roles/serviceusage.serviceUsageAdmin）が必要です。ロールを付与する方法を確認する。

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com

  • Google Cloud CLI をインストールします。

  • 外部 ID プロバイダ（IdP）を使用している場合は、まず連携 ID を使用して gcloud CLI にログインする必要があります。

  • gcloud CLI を初期化するには、次のコマンドを実行します。

    gcloud init

  • Google Cloud プロジェクトを作成または選択します。

    プロジェクトの選択または作成に必要なロール

    • プロジェクトを選択する: プロジェクトの選択に特定の IAM ロールは必要ありません。ロールが付与されているプロジェクトであれば、どのプロジェクトでも選択できます。
    • プロジェクトを作成する: プロジェクトを作成するには、resourcemanager.projects.create 権限を含むプロジェクト作成者ロール（roles/resourcemanager.projectCreator）が必要です。詳しくは、ロールを付与する方法をご覧ください。

    • Google Cloud プロジェクトを作成します。

      gcloud projects create PROJECT_ID

      PROJECT_ID は、作成する Google Cloud プロジェクトの名前に置き換えます。

    • 作成した Google Cloud プロジェクトを選択します。

      gcloud config set project PROJECT_ID

      PROJECT_ID は、 Google Cloud プロジェクトの名前に置き換えます。

  • Google Cloud プロジェクトに対して課金が有効になっていることを確認します。

  • 必要な API を有効にします。

    API を有効にするために必要なロール

    API を有効にするには、serviceusage.services.enable 権限を含む Service Usage 管理者 IAM ロール（roles/serviceusage.serviceUsageAdmin）が必要です。ロールを付与する方法を確認する。

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
    •

  • ユーザー アカウントにロールを付与します。次の IAM ロールごとに次のコマンドを 1 回実行します。 roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    次のように置き換えます。

    • PROJECT_ID: プロジェクト ID。
    • USER_IDENTIFIER: ユーザー アカウントの識別子。例: myemail@example.com
    • ROLE: ユーザー アカウントに付与する IAM ロール。

環境を準備する

このチュートリアルでは、Cloud Shell を使用します。

  1. Google Cloud コンソールに移動します。

  2. Google Cloud コンソール ウィンドウの上部にある [Cloud Shell をアクティブにする] ボタンをクリックします。

  3. 次の環境変数を設定します。

    export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export CONTROL_PLANE_REGION=CONTROL_PLANE_REGION
export NODE_ZONE=NODE_ZONE
export CLUSTER_NAME=CLUSTER_NAME
export GPU_TYPE=GPU_TYPE
export MACHINE_TYPE=MACHINE_TYPE
export KSA_NAME=generic-ksa
export NAMESPACE=default
export RESERVATION=RESERVATION
export LUSTRE_NAME=LUSTRE_NAME
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

    次の値を置き換えます。

    • CLUSTER_NAME: GKE クラスタの名前。
    • CONTROL_PLANE_REGION: GKE クラスタ コントロール プレーンの Compute Engine リージョン。
    • NODE_ZONE: ノードのゾーン。NVIDIA B200 GPU または H200 GPU を使用できるゾーンを選択します。
    • GPU_TYPE: Compute Engine の容量予約で予約したアクセラレータ。次のいずれかの値を指定する必要があります。
      • nvidia-b200: NVIDIA B200（180 GB）
      • nvidia-h200-141gb: NVIDIA H200（141 GB）
    • MACHINE_TYPE: 使用するマシンのタイプ:
      • NVIDIA B200（180 GB）GPU の場合は、a4-highgpu-8g 以降を使用します。
      • NVIDIA H200（141 GB）GPU の場合は、a3-ultragpu-8g 以降を使用します。
    • RESERVATION: GPU 予約の名前。
    • LUSTRE_NAME: Lustre インスタンスの名前。
    • YOUR_HF_TOKEN: Hugging Face トークン。
    • YOUR_WANDB_API_KEY: Wandb API キー。

  4. ネットワークに対して、次の環境変数を作成します。

    export NETWORK="NETWORK-NAME"
export GVNIC_NETWORK_PREFIX="GVNIC-NAME"
export RDMA_NETWORK_PREFIX="RDMA-NAME"

    次の値を置き換えます。

    • NETWORK-NAME: GKE のネットワーク名。
    • GVNIC-NAME: gVNIC ネットワーク名の接頭辞。任意の接頭辞を使用できます。
    • RDMA-NAME: リモート ダイレクト メモリ アクセス（RDMA）ネットワークの接頭辞。任意の接頭辞を使用できます。

インフラストラクチャを設定する

このセクションでは、VPC ネットワークと GKE クラスタを作成します。

VPC ネットワークの作成

  1. gVNIC インターフェース用の VPC ネットワークを作成します。

    gcloud compute networks create ${NETWORK} --subnet-mode=auto
gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
    --subnet-mode=custom
gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --region=${CONTROL_PLANE_REGION} \
    --range=192.168.0.0/24
gcloud compute firewall-rules create ${GVNIC_NETWORK_PREFIX}-internal \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --action=ALLOW \
    --rules=tcp:0-65535,udp:0-65535,icmp \
    --source-ranges=192.168.0.0/16

  2. 8 個の GPU 用に 8 個のサブネットを含む RDMA 用の VPC ネットワークとサブネットを作成します。

    gcloud compute networks create ${RDMA_NETWORK_PREFIX}-net \
    --network-profile=${NODE_ZONE}-vpc-roce \
    --subnet-mode=custom

for N in $(seq 0 7); do
  gcloud compute networks subnets create ${RDMA_NETWORK_PREFIX}-sub-$N \
    --network=${RDMA_NETWORK_PREFIX}-net \
    --region=${CONTROL_PLANE_REGION} \
    --range=192.168.$((N+1)).0/24 &
done
wait

GKE クラスタを作成する

GKE Standard クラスタで NeMo RL を設定できます。

  1. Standard クラスタを作成します。

    gcloud container clusters create ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --enable-dataplane-v2 \
    --enable-ip-alias \
    --enable-multi-networking \
    --addons=RayOperator,LustreCsiDriver \
    --enable-legacy-lustre-port \
    --machine-type=n2-highmem-80 \
    --num-nodes=1 \
    --min-nodes=1 \
    --max-nodes=5 \
    --enable-autoscaling \
    --network=${NETWORK}

  2. クラスタの認証情報を取得します。

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION}

  3. GPU ノードプールを作成します。

    gcloud container node-pools create gpu-pool \
    --cluster=${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION} \
    --node-locations=${NODE_ZONE} \
    --machine-type=${MACHINE_TYPE} \
    --accelerator=type=${GPU_TYPE},count=8,gpu-driver-version=DEFAULT \
    --reservation-affinity=specific \
    --reservation=${RESERVATION} \
    --enable-autoscaling \
    --num-nodes=0 \
    --total-max-nodes=2 \
    --additional-node-network=network=${GVNIC_NETWORK_PREFIX}-net,subnetwork=${GVNIC_NETWORK_PREFIX}-sub \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-0 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-1 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-2 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-3 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-4 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-5 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-6 \
    --additional-node-network=network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-7

  4. NCCL RDMA インストーラをインストールします。

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer.yaml

ネットワーク マッピングを構成する

  1. 次のマニフェストを network-mapping.yaml として保存します。

    # Copyright 2026 Google LLC. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: gvnic-1
spec:
  vpc: ${GVNIC_NETWORK_PREFIX}-net
  vpcSubnet: ${GVNIC_NETWORK_PREFIX}-sub
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: gvnic-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: gvnic-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-0
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-0
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-0
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-0
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-1
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-1
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-2
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-2
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-2
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-2
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-3
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-3
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-3
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-3
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-4
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-4
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-4
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-4
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-5
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-5
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-5
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-5
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-6
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-6
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-6
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-6
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-7
spec:
  vpc: ${RDMA_NETWORK_PREFIX}-net
  vpcSubnet: ${RDMA_NETWORK_PREFIX}-sub-7
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-7
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-7

  2. 次のようにマニフェストを適用します。

    envsubst < network-mapping.yaml > network-mapping-updated.yaml
kubectl apply -f network-mapping-updated.yaml

ストレージを準備する

このセクションでは、RL ワークロードに必要な高性能ストレージをプロビジョニングする Managed Lustre インスタンスを作成します。

  1. プライベート サービス アクセス用の IP アドレス範囲を割り振ります。

    gcloud compute addresses create ${LUSTRE_NAME}-range \
    --global --purpose=VPC_PEERING \
    --prefix-length=20 --network=${NETWORK}

  2. ピアリングを接続します。

    gcloud services vpc-peerings connect \
    --service=servicenetworking.googleapis.com \
    --ranges=${LUSTRE_NAME}-range \
    --network=${NETWORK}

  3. Managed Lustre インスタンスを作成します。

    gcloud lustre instances create ${LUSTRE_NAME} \
    --per-unit-storage-throughput=500 \
    --capacity-gib=18000 \
    --filesystem=lustrefs \
    --location=${NODE_ZONE} \
    --network=projects/${PROJECT_ID}/global/networks/${NETWORK} \
    --gke-support-enabled

  4. Managed Lustre CSI ドライバを使用して既存の Managed Lustre インスタンスにアクセスする

    1. Managed Lustre インスタンスの IP アドレスを抽出します。

      export LUSTRE_IP=$(gcloud lustre instances describe ${LUSTRE_NAME} \
    --location=$NODE_ZONE --format="value(mountPoint)" | awk -F'@' '{print $1}')

    2. lustre-pv.yaml マニフェストを調べます。

      # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: PersistentVolume
metadata:
  name: lustre-pv
spec:
  storageClassName: lustre-rwx-500mbps-per-tib
  capacity:
    storage: 18000Gi
  accessModes:
    - ReadWriteMany
  persistentVolumeReclaimPolicy: Retain
  volumeMode: Filesystem
  claimRef:
    namespace: default
    name: lustre-pvc
  csi:
    driver: lustre.csi.storage.gke.io
    volumeHandle: "${PROJECT_ID}/${NODE_ZONE}/${LUSTRE_NAME}"
    volumeAttributes:
      ip: ${LUSTRE_IP}
      filesystem: lustrefs

    3. 次のようにマニフェストを適用します。

      envsubst < lustre-pv.yaml > lustre-pv-updated.yaml
kubectl apply -f lustre-pv-updated.yaml

    4. lustre-pvc.yaml マニフェストを調べます。

      # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: lustre-pvc
spec:
  accessModes:
    - ReadWriteMany
  storageClassName: lustre-rwx-500mbps-per-tib
  volumeName: lustre-pv
  resources:
    requests:
      storage: 18000Gi

    5. 次のようにマニフェストを適用します。

      kubectl apply -f lustre-pvc.yaml

RayCluster をデプロイする

このセクションでは、サンプル リポジトリのクローンを作成し、マニフェストを準備して、Ray クラスタをデプロイします。

  1. サンプル リポジトリのクローンを作成します。

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples

  2. 作業ディレクトリに移動します。

    cd ai-ml/nemo-rl-on-gke/nemoRL

  3. マニフェスト values.yaml を調べます。

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

image:
  repository: "nvcr.io/nvidia/nemo-rl"
  tag: "v0.5.0" 
  pullPolicy: Always

nameOverride: "kuberay"
fullnameOverride: ""

common:
  containerEnv: {}

configMap:
  fluentbit:
    data:
      fluent-bit.conf: |
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/worker-*
            Tag               ray-worker
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/raylet*
            Tag               raylet
        [INPUT]
            Name              tail
            Path              /tmp/ray/session_latest/logs/*
            Exclude_Path      /tmp/ray/session_latest/logs/debug_state.txt,/tmp/ray/session_latest/logs/raylet*,/tmp/ray/session_latest/logs/worker-*
            Tag               ray-misc
        [OUTPUT]
            Name              stackdriver
            Match             *
            resource          gce_instance
            labels_key        labels

# --- Head Node Configuration ---
head:
  enableInTreeAutoscaling: false
  serviceAccountName: ""
  rayStartParams:
    dashboard-host: '0.0.0.0'
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        networking.gke.io/default-interface: 'eth0'
  containerEnv:
  - name: RAY_GROUP
    value: "head"
  nodeSelector:
    cloud.google.com/gke-nodepool: default-pool
  resources:
    limits:
      cpu: "64"
      memory: "500G"
      nvidia.com/gpu: 0
    requests:
      cpu: "64"
      memory: "500G"
      nvidia.com/gpu: 0
  tolerations:
    # - operator: "Exists"
    #   key: "components.gke.io/gke-managed-components"
    # - key: "nvidia.com/gpu"
    #   operator: "Exists"
    #   effect: "NoSchedule"
  volumeMounts:
    - mountPath: /data
      name: lustre-data

  volumes:
    - name: log-volume
      emptyDir: {}
    - name: fluentbit-config-volume
      configMap:
        name: "ray-cluster-kuberay-fluentbit-config"
    - name: lustre-data
      persistentVolumeClaim:
        claimName: lustre-pvc
  sidecarContainers:
    - name: fluent-bit
      image: fluent/fluent-bit:latest
      env:
      - name: RAY_GROUP
        value: "head"
      volumeMounts:
        - name: fluentbit-config-volume
          mountPath: /fluent-bit/etc/
        - mountPath: /tmp/ray
          name: log-volume

  # --- HEAD POD STARTUP SCRIPT ---
  command:
    - "bash"
    - "-c"
    - |
      set -ex
      echo "--- Head Pod Setup ---"
      apt-get update
      apt-get install -y sudo netcat-openbsd pciutils
      cd /opt/nemo-rl
      /usr/bin/python -m pip install uv
      /usr/bin/python -m uv venv
      echo "Head pod setup complete. Starting Ray..."

      exec ${KUBERAY_GEN_RAY_START_CMD}

  args: []
  headService: {}
  # nodeSelector:
  #   cloud.google.com/gke-accelerator: nvidia-b200 #cloud.google.com/gke-nodepool: cpu-node-pool-llama #cpu-node-pool

# --- Default Worker (Disabled) ---
worker:
  disabled: true

# --- A4 GPU Worker Groups ---
additionalWorkerGroups:
  worker-grp-0:
    disabled: false
    replicas: 4
    annotations:
      networking.gke.io/default-interface: 'eth0'
      networking.gke.io/interfaces: |
        [
          {"interfaceName":"eth0","network":"default"},
          {"interfaceName":"eth1","network":"gvnic-1"},
          {"interfaceName":"eth2","network":"rdma-0"},
          {"interfaceName":"eth3","network":"rdma-1"},
          {"interfaceName":"eth4","network":"rdma-2"},
          {"interfaceName":"eth5","network":"rdma-3"},
          {"interfaceName":"eth6","network":"rdma-4"},
          {"interfaceName":"eth7","network":"rdma-5"},
          {"interfaceName":"eth8","network":"rdma-6"},
          {"interfaceName":"eth9","network":"rdma-7"}
        ]
    containerEnv:
      - name: RAY_GROUP
        valueFrom:
          fieldRef:
            fieldPath: metadata.labels['ray.io/group']
      - name: NCCL_NET  
        value: "gIB"
      - name: NCCL_IB_GID_INDEX
        value: "3"   
      - name: GLOO_SOCKET_IFNAME
        value: "eth0"
      - name: NCCL_CROSS_NIC
        value: "0"
      - name: NCCL_SOCKET_IFNAME
        value: "eth0"
      - name: TP_SOCKET_IFNAME # Specific to DTensor/PyTorch Distributed
        value: "eth0"
      - name: NCCL_TUNER_CONFIG_PATH
        value: "/usr/local/gib/configs/tuner_config_a4.txtpb"
      - name: NCCL_NET_GDR_LEVEL
        value: "PIX"
      - name: LD_LIBRARY_PATH
        value: /usr/local/nvidia/lib64
    resources:
      limits:
        nvidia.com/gpu: 8
        cpu: "206"
        memory: "2400Gi"
      requests:
        nvidia.com/gpu: 8
        cpu: "206"
        memory: "2400Gi"

    nodeSelector:
      cloud.google.com/gke-accelerator: nvidia-b200
    tolerations:
      - operator: "Exists"
        key: "nvidia.com/gpu"
      - operator: "Exists"
        key: "cloud.google.com/impending-node-termination"
      - operator: "Exists"
        key: "user-workload"
    securityContext:
      privileged: true
    volumes:
      - name: log-volume
        emptyDir: {}
      - name: shared-memory
        emptyDir:
          medium: "Memory"
          sizeLimit: 240Gi
      - name: ray-tmp
        emptyDir:
          medium: "Memory"
      - name: fluentbit-config-volume
        configMap:
          name: "ray-cluster-kuberay-fluentbit-config"
      - name: nvidia-install-dir-host
        hostPath:
          path: /home/kubernetes/bin/nvidia
      - name: gib-nccl-plugin-volume
        hostPath: 
          path: /home/kubernetes/bin/gib
      - name: lustre-data
        persistentVolumeClaim:
          claimName: lustre-pvc
    volumeMounts:
      - mountPath: /tmp/ray
        name: log-volume
      - name: shared-memory
        mountPath: /dev/shm
      - name: nvidia-install-dir-host
        mountPath: /usr/local/nvidia
      - name: gib-nccl-plugin-volume
        mountPath: /usr/local/gib
      - mountPath: /data
        name: lustre-data   
    # --- WORKER POD STARTUP SCRIPT ---
    command:
      - "bash"
      - "-c"
      - |
        set -ex

        echo "--- Worker Pod Setup ---"
        apt-get update
        apt-get install -y sudo netcat-openbsd pciutils
        cd /opt/nemo-rl
        /usr/bin/python -m pip install uv
        /usr/bin/python -m uv venv

        ldconfig /usr/local/nvidia/lib64/
        ldconfig -p | grep libcuda | sed 's/^/  /'
        export LD_LIBRARY_PATH="/usr/local/gib/lib64:$LD_LIBRARY_PATH"
        source /usr/local/gib/scripts/set_nccl_env.sh

        echo "Worker pod setup complete. Starting Ray..."

        exec ${KUBERAY_GEN_RAY_START_CMD}


    sidecarContainers:
      - name: fluent-bit
        env:
          - name: RAY_GROUP
            valueFrom:
              fieldRef:
                fieldPath: metadata.labels['ray.io/group']
        image: fluent/fluent-bit:latest
        volumeMounts:
          - name: fluentbit-config-volume
            mountPath: /fluent-bit/etc/
          - mountPath: /tmp/ray
            name: log-volume

# --- Service Config ---
service:
  type: ClusterIP

    このチュートリアルで使用するアクセラレータに基づいて、NCCL_TUNER_CONFIG_PATH を次のいずれかの値に置き換えます。

    • NVIDIA B200（180 GB）: /usr/local/gib/configs/tuner_config_a4.txtpb
    • NVIDIA H200（141 GB）: /usr/local/gib/configs/tuner_config_a3u.txtpb

    このマニフェストでは、ヘッドノードがジョブを管理し、Ray ダッシュボードをホストします。ワーカーノードはトレーニング ジョブを実行します。

  4. Ray クラスタをデプロイします。

    export REPLICA_COUNT=2
helm install ray-cluster . \
  --set additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT

    このチュートリアルでは、2 つのワーカーノードを使用します。ワーカーノードの数を変更する場合は、REPLICA_COUNT の値を変更します。

  5. ワーカーノードとヘッドノードが実行されていることを確認します。

    kubectl get pods

    出力は次のようになります。

    NAME                                          READY STATUS RESTARTS AGE
ray-cluster-kuberay-head-sw7dp                3/3   Running 0      33h
ray-cluster-kuberay-worker-grp-0-worker-gkbxw 3/3   Running 0      33h
ray-cluster-kuberay-worker-grp-0-worker-kdg62 3/3   Running 0      33h

  6. Ray クラスタが実行中であることを確認します。

    kubectl ray get cluster

    出力は次のようになります。

    NAME                 NAMESPACE DESIRED WORKERS AVAILABLE WORKERS CPUS GPUS TPUS MEMORY CONDITION STATUS AGE
ray-cluster-kuberay  default   2       2           618     17   0    1573741824k RayClusterProvisioned ready 33h

GRPO ジョブを起動する

Ray クラスタの準備ができたら、GKE で実行中の Ray クラスタに Ray ジョブを送信できます。NeMo RL は、RL トレーニング ジョブの実行中にモデルを自動的にダウンロードします。

Ray ジョブを送信するには、インタラクティブ セッションを開始してジョブを実行します。

  1. Ray クラスタへのローカル接続を確立するには、次のコマンドを実行します。

      kubectl ray session ray-cluster-kuberay

    このコマンドにより、ローカルマシンと GKE クラスタ内の Ray ヘッドノード間のポート転送が開始されます。このセッションがアクティブになっている間は、ターミナルが占有されます。続行するには、別のターミナル インスタンスを開きます。

  2. gemma3-27b-gsm8k.sh ファイルを編集します。

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!/bin/bash
WANDB_API_KEY='YOUR_WANDB_API_KEY' # Update this with your WANDB API key
HF_TOKEN='YOUR_HF_TOKEN' # Update this with your HF token
WORLD_SIZE=16

# --- Step 1: Find the Ray Head Pod ---
echo "Finding Ray head pod..."
export HEAD_POD_NAME=$(kubectl get pods --selector=ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
if [ -z "$HEAD_POD_NAME" ]; then
    echo "Error: No running Ray head pod found. Please check your cluster."
    exit 1
fi
echo "Found head pod: $HEAD_POD_NAME"
echo ""

# --- Step 2: Define the Job Script to Run ---
# This is the script that will be executed *inside* the head pod.
# It assumes the 'uv venv' setup from the values.yaml is already done.
JOB_SCRIPT=$(cat <<EOF
set -ex

echo "--- Running on Ray Head Pod ($HOSTNAME) ---"
cd /opt/nemo-rl

git pull && git checkout main

sed -i 's/subset: Optional\[str\] = None/subset: Optional[str] = "main"/' /opt/nemo-rl/nemo_rl/data/datasets/response_datasets/response_dataset.py
sed -i 's/raw_dataset = load_dataset(data_path)/raw_dataset = load_dataset(data_path, "main")/' /opt/nemo-rl/nemo_rl/data/datasets/utils.py

echo "Setting environment variables..."
export WANDB_API_KEY=$WANDB_API_KEY
export HF_TOKEN=$HF_TOKEN
export HF_HOME=/opt/nemo-rl/

###-----Example to launch Gemma3-27B on 2 nodes (16 GPUs)----------
uv run python examples/run_grpo_math.py \
  --config examples/configs/recipes/llm/grpo-gemma3-27b-it-8n4g-fsdp2tp4-actckpt-long.yaml \
  cluster.num_nodes=2 \
  cluster.gpus_per_node=8 \
  grpo.max_num_steps=10 \
  checkpointing.checkpoint_dir=/data/nemo_rl_gemma3_27b_3_17 \
  data.dataset_name=ResponseDataset \
  +data.train_data_path=openai/gsm8k \
  +data.val_data_path=openai/gsm8k \
  +data.val_split=test \
  +data.train_split=train \
  +data.subset="main" \
  +data.input_key="question" \
  +data.output_key="answer" \
  logger.tensorboard_enabled=False \
  logger.wandb_enabled=True \
  logger.wandb.name='nemo_rl_gemma3_27b_3_17' \
  grpo.num_prompts_per_step=16 \
  grpo.num_generations_per_prompt=32 \
  policy.generation.colocated.enabled=False \
  policy.generation.colocated.resources.num_nodes=1 \
  policy.generation.colocated.resources.gpus_per_node=8 \
  policy.generation.vllm_cfg.tensor_parallel_size=8 \
  policy.generation.vllm_cfg.gpu_memory_utilization=0.9 \
  policy.dtensor_cfg.tensor_parallel_size=8

echo "--- Job Finished ---"
EOF
)

# --- Step 3: Execute the Job ---
echo "Submitting job to $HEAD_POD_NAME..."
echo "$JOB_SCRIPT" | tr -d '\r' | kubectl exec -i $HEAD_POD_NAME -c ray-head -- /bin/bash

echo ""
echo "Job submission complete."

    gemma3-27b-gsm8k.sh ファイルで次の値を置き換えます。

    • YOUR_WANDB_API_KEY: WandB API キー。
    • YOUR_HF_TOKEN: Hugging Face トークン。

    このファイルには、GSM8k データセットで gemma3-27b-it モデルを使用してジョブを実行する構成が記述されています。GRPO トレーニング パイプラインを完了するために、このスクリプトでは次のパラメータが定義されています。

    • num_prompts_per_step: 16 と num_generations_per_prompt: 32: Gemma3-27b-it モデルは、プロンプトごとに大量のレスポンスを生成します。この構成では、モデルは合計 512 個のレスポンス（16 × 32 = 512）を生成します。
    • policy.generation.colocated.enabled=False: このパラメータは、コロケーション生成機能を無効にします。つまり、モデルはトレーニング プロセスと同じノードでレスポンスを生成しません。標準の RL では、同じ GPU がトレーニングと生成の両方を処理します。この NeMo RL 設定では、特定のノード（policy.generation.colocated.resources パラメータで管理）を vLLM 推論専用にし、クラスタの残りの部分を負荷の高いトレーニング計算に集中させます。これらのワークロードを分離することで、メモリ使用量の多いトレーニング バッファと計算負荷の高い推論ワークロード間のリソース競合を防ぐことができます。

  3. Job を送信するには、次のコマンドを実行します。

    bash gemma3-27b-it/gemma3-27b-gsm8k.sh

    ジョブの実行中、出力にはトレーニング結果、タイミング、パフォーマンス指標が表示されます。

GRPO ジョブの健全性をモニタリングする

Ray がジョブを完了すると、NeMo RL は構成されたパスにチェックポイントを保存します。

  1. GRPO ジョブの出力を確認するには、ray-head コンテナへの SSH セッションを作成します。

    kubectl exec -it $(kubectl get pods -l ray.io/node-type=head -o name) -c ray-head -- bash

  2. ray-head コンテナのターミナル内で apt tree ユーティリティをインストールします。

    apt update && apt install -y tree

    出力は次のようになります。

     root@ray-cluster-kuberay-worker-grp-0-worker-gkbxw:/opt/nemo-rl# tree /data/nemo_rl_gemma3_27b_3_17/
 /data/nemo_rl_gemma3_27b_3_17/
 `-- step_10
     |-- config.yaml
     |-- policy
     |   |-- optimizer
     |   |   |-- __0_0.distcp
     |   |   |-- __10_0.distcp
     |   |   |-- __11_0.distcp
     |   |   |-- __12_0.distcp
     |   |   |-- __13_0.distcp
     |   |   |-- __14_0.distcp
     |   |   |-- __15_0.distcp
     |   |   |-- __1_0.distcp
     |   |   |-- __2_0.distcp
     |   |   |-- __3_0.distcp
     |   |   |-- __4_0.distcp
     |   |   |-- __5_0.distcp
     |   |   |-- __6_0.distcp
     |   |   |-- __7_0.distcp
     |   |   |-- __8_0.distcp
     |   |   `-- __9_0.distcp
     |   |-- tokenizer
     |   |   |-- chat_template.jinja
     |   |   |-- special_tokens_map.json
     |   |   |-- tokenizer.json
     |   |   `-- tokenizer_config.json
     |   `-- weights
     |       |-- __0_0.distcp
     |       |-- __10_0.distcp
     |       |-- __11_0.distcp
     |       |-- __12_0.distcp
     |       |-- __13_0.distcp
     |       |-- __14_0.distcp
     |       |-- __15_0.distcp
     |       |-- __1_0.distcp
     |       |-- __2_0.distcp
     |       |-- __3_0.distcp
     |       |-- __4_0.distcp
     |       |-- __5_0.distcp
     |       |-- __6_0.distcp
     |       |-- __7_0.distcp
     |       |-- __8_0.distcp
     |       `-- __9_0.distcp
     |-- train_dataloader.pt
     `-- training_info.json

 6 directories, 39 files
 ```

クリーンアップ

課金されないようにするには、リソースを削除します。

helm delete ray-cluster
gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION} \
    --quiet
gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet
gcloud services vpc-peerings delete \
    --service=servicenetworking.googleapis.com \
    --network=${NETWORK}
gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet
gcloud compute firewall-rules delete ${GVNIC_NETWORK_PREFIX}-internal --quiet
gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
    --region=${CONTROL_PLANE_REGION} --quiet
gcloud compute networks delete ${GVNIC_NETWORK_PREFIX}-net --quiet
for N in $(seq 0 7); do
  gcloud compute networks subnets delete ${RDMA_NETWORK_PREFIX}-sub-$N \
    --region=${CONTROL_PLANE_REGION} --quiet &
done
wait
gcloud compute networks delete ${RDMA_NETWORK_PREFIX}-net --quiet
gcloud compute networks delete ${NETWORK} --quiet

次のステップ