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

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

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

このチュートリアルを実行する前に、GKE で verl を使用して強化学習をファインチューニングしてスケーリングするチュートリアルを完了することをおすすめします。次のチュートリアルでは、ファインチューニングとスケーリング RL と 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 トレーニング ジョブを実行します。

始める前に

  • 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_LOCATION=CONTROL_PLANE_LOCATION
export NODE_LOCATION=NODE_LOCATION
export CLUSTER_NAME=CLUSTER_NAME
export GPU_TYPE=GPU_TYPE
export MACHINE_TYPE=MACHINE_TYPE
export GKE_VERSION=GKE_VERSION
export KSA_NAME=generic-ksa
export NAMESPACE=default
export GS_BUCKET=BUCKET_NAME-${PROJECT_ID}
export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

    次の値を置き換えます。

    • CLUSTER_NAME: GKE クラスタの名前。
    • CONTROL_PLANE_LOCATION: GKE クラスタ コントロール プレーンの Compute Engine リージョン。
    • NODE_LOCATION: ノードのロケーション。NVIDIA B200 または 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 以降を使用します。

    • GKE_VERSION: 使用する GKE のバージョン。

      • NVIDIA B200（180 GB）GPU の場合は、1.32.2-gke.1422000 以降を使用します。
      • NVIDIA H200（141 GB）GPU の場合は、1.31.4-gke.1183000 以降を使用します。

    • BUCKET_NAME: Cloud Storage バケットのベース名。

    • YOUR_HUGGING_FACE_TOKEN: Hugging Face トークン。

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

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

    次の値を置き換えます。

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

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

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

VPC ネットワークの作成

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

    gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
    --project=${PROJECT_ID} \
    --subnet-mode=custom
gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --location=${CONTROL_PLANE_LOCATION} \
    --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=${CONTROL_PLANE_LOCATION}-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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --range=192.168.$((N+1)).0/24 &
done
wait

GKE クラスタを作成する

GKE Autopilot クラスタまたは GKE Standard クラスタで NeMo RL を設定できます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードについてをご覧ください。

Autopilot

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

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-multi-networking  \
    --enable-ray-operator

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

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

  3. Autopilot 用の NCCL RDMA インストーラをインストールします。

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

標準

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

    gcloud container clusters create ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-dataplane-v2 \
    --enable-ip-alias \
    --enable-multi-networking \
    --addons=RayOperator \
    --num-nodes=1

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

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

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

    gcloud container node-pools create gpu-pool \
    --cluster=${CLUSTER_NAME} \
    --node-locations=${NODE_LOCATION} \
    --machine-type=${MACHINE_TYPE} \
    --accelerator=type=${GPU_TYPE},count=8 \
    --spot \
    --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. 次のようにマニフェストを適用します。

    kubectl apply -f network-mapping.yaml

ストレージを準備する

このセクションでは、Cloud Storage バケットと Managed Lustre インスタンスを作成します。これにより、RL ワークロードに必要な高性能ストレージがプロビジョニングされます。

  1. Cloud Storage バケットを作成します。

    gcloud storage buckets create gs://${GS_BUCKET} \
    --location=${CONTROL_PLANE_LOCATION} \
    --enable-hierarchical-namespace \
    --uniform-bucket-level-access

  2. Kubernetes サービス アカウント（KSA）を作成し、バケットにバインドします。

    kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}

gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
    --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
    --role "roles/storage.objectUser"

  3. 次の手順に沿って、Managed Lustre を設定します。

    1. Managed Lustre インスタンスを作成するの手順に沿って、Managed Lustre インスタンスを作成します。インスタンスが GKE クラスタと同じネットワークを使用していることを確認します。
    2. 既存の Managed Lustre インスタンスにアクセスするの手順に沿って、Managed Lustre インスタンスにアクセスします。

RayCluster をデプロイする

このセクションでは、サンプル リポジトリをクローニングし、マニフェストを準備して、launcher.sh スクリプトを実行します。

  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"
  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 values.additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT

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

  5. Ray クラスタをデプロイするには、launcher.sh スクリプトを実行します。

    bash launcher.sh

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

    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

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

    kubectl ray get clusters

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

    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 3 nodes (24 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=300 \
  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=64 \
  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 モデルを使用して Job を実行する構成が記述されています。GRPO トレーニング パイプラインを完了するために、このスクリプトは次のパラメータを定義します。

    • num_prompts_per_step: 16 と num_generations_per_prompt: 64: Gemma3-27b-it モデルは、プロンプトごとに大量のレスポンスを生成します。この構成では、モデルは合計 1,024 個のレスポンスを生成します（16 × 64 = 1,024）。
    • 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. apt tree ユーティリティをインストールします。

    apt install tree

  2. GRPO ジョブの健全性をモニタリングするには、Ray ヘッドノードのログを確認します。

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

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

    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_LOCATION}
gcloud storage rm -r gs://${GS_BUCKET}

次のステップ