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）に代わるメモリ効率の高い方法を提供します。Critic ネットワークの代わりに、GRPO は同じプロンプトに対する一連のレスポンスを生成し、そのグループの平均報酬をベースラインとして使用します。

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

NVIDIA NeMo RL

NeMo RL は、スケーラブルな RL 用に設計された NVIDIA のオープンソースの事後トレーニング ライブラリです。NeMo RL は、より広範な NeMo フレームワーク エコシステムの一部であり、単一の 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 トレーニング ジョブを実行します。

始める前に

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

  2. フェデレーション ID（連携 ID）を使用するように gcloud CLI を構成します。

    詳細については、連携 ID を使用して gcloud CLI にログインするをご覧ください。

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

    gcloud init

  4. 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 プロジェクトの名前に置き換えます。

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

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

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

    API を有効にするには、serviceusage.services.enable 権限が必要です。プロジェクトを作成した場合は、オーナーロール（roles/owner）を介してこの権限がすでに付与されている可能性があります。それ以外の場合は、Service Usage 管理者ロール（roles/serviceusage.serviceUsageAdmin）を介してこの権限を取得できます。ロールを付与する方法をご覧ください。

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

  7. ユーザー アカウントにロールを付与します。次の 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

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

  8. Hugging Face アカウントを作成します（まだ作成していない場合）。
  9. read access を含む Hugging Face トークンがあることを確認します。
  10. Weights & Biases（Wandb）アカウントを作成します（まだアカウントを保有していない場合）。
  11. Wandb API キーを作成します。
  12. Google Cloud プロジェクトに B200 GPU と H200 GPU に十分な割り当てがあることを確認します。詳細については、GPU 割り当てを計画すると GPU 割り当てをご覧ください。

環境を準備する

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

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

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

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

    export CONTROL_PLANE_REGION="YOUR_REGION"
export NODE_ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GPU_TYPE="YOUR_GPU_TYPE"
export MACHINE_TYPE="YOUR_MACHINE_TYPE"
export KSA_NAME="generic-ksa"
export NAMESPACE="default"
export RESERVATION="RESERVATION_NAME"
export LUSTRE_NAME="CHOSEN_LUSTRE_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export WANDB_API_KEY="YOUR_WANDB_API_KEY"

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" --format="value(projectNumber)")

    次の値を置き換えます。

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

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

    export NETWORK="YOUR_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 | kubectl apply -f -

ストレージを準備する

このセクションでは、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 | kubectl apply -f -

    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

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

    cd kubernetes-engine-samples/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                2/2   Running 0      33h
ray-cluster-kuberay-worker-grp-0-worker-gkbxw 2/2   Running 0      33h
ray-cluster-kuberay-worker-grp-0-worker-kdg62 2/2   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. 別のターミナルで、kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL/gemma3-27b-it に移動して 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

  3. ray-head コンテナのディレクトリ構造を一覧表示します。

    tree /data/nemo_rl_gemma3_27b_3_17/

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

    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

クリーンアップ

このチュートリアルで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、個々のリソースを削除するか、リソースを含むプロジェクトを削除します。

リソースの削除

  1. Slurm クラスタを削除します。

    helm delete ray-cluster

  2. GKE クラスタを削除します。

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION} \
    --quiet

  3. Lustre ファイル システムを削除します。

    gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet

  4. VPC ピアリングを削除します。

    gcloud services vpc-peerings delete \
    --service=servicenetworking.googleapis.com \
    --network=${NETWORK}

  5. Lustre プライベート IP アドレス範囲を削除します。

    gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet

  6. RDMA サブネットと gVNIC サブネットを削除します。

    gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
    --region=${CONTROL_PLANE_REGION} --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

  7. ファイアウォール ルールとネットワークを削除します。

    echo "[$(date)] ========== Deleting firewall rules and networks =========="

NETWORKS=(
    "${RDMA_NETWORK_PREFIX}-net"
    "${GVNIC_NETWORK_PREFIX}-net"
    "${NETWORK}"
)

for NW in "${NETWORKS[@]}"; do

  echo "========== Deleting firewall rules for ${NW} =========="
  while true; do
      rules=$(gcloud compute firewall-rules list \
        --filter="network:${NW}" \
        --format="value(name)" \
        --project="${PROJECT_ID}")

        if [[ -z "${rules}" ]]; then
          echo "No firewall rules remain for ${NW}"
          break
        fi

        for rule in ${rules}; do
          echo "Deleting firewall rule ${rule}..."
          gcloud compute firewall-rules delete "${rule}" --project="${PROJECT_ID}" --quiet || true
        done

        sleep 3
      done

      echo "[$(date)] ========== Deleting network ${NW} =========="
      gcloud compute networks delete ${NW} --quiet || true
done

プロジェクトの削除

Google Cloud プロジェクトを削除する:

gcloud projects delete PROJECT_ID

次のステップ