在 GKE 上使用 NVIDIA NeMo RL 微調及擴充強化學習模型

本教學課程說明如何在 Google Kubernetes Engine (GKE) 上，為強化學習 (RL) 協調分散式訓練環境。您可以使用 Ray 和 NVIDIA NeMo RL 架構設定分散式訓練環境，微調模型。

本教學課程著重於在 GKE 上，使用 Ray 和 NeMo RL 進行群組相對政策最佳化 (GRPO) 訓練管道。GRPO 是一種強化學習演算法，旨在提升模型的推理能力。這項演算法可節省記憶體，並透過淘汰 Critic 或價值模型，改用相對群組計算，簡化 RL 程序。

執行本教學課程前，請先完成在 GKE 上使用 Verl 微調及擴充強化學習教學課程。本教學課程使用的叢集設定和配置，與「使用 Verl 進行微調和調整 RL 大小」教學課程相同。

背景

以下各節簡要概述本教學課程中使用的概念。

增強學習 (RL)

RL 是透過經驗、探索和意見回饋來訓練模型，而不是靜態模仿。預先訓練會教導模型該說什麼，而人類回饋增強學習 (RLHF) 則會教導模型如何提供實用、安全且合乎邏輯的內容。RL 可做為基礎模型與微調模型的橋樑，適用於特定用途。

詳情請參閱「什麼是強化學習？」一文。

群組相對政策最佳化 (GRPO)

GRPO 是 DeepSeek 普及的演算法，可移除 Critic 模型，為 LLM 對齊提供記憶體效率替代方案，取代近端策略最佳化 (PPO)。GRPO 不會使用 Critic 網路，而是會針對相同提示產生一組回應，並以該組回應的平均獎勵做為基準。

詳情請參閱 GRPO。

NVIDIA NeMo RL

NeMo RL 是 NVIDIA 的開放原始碼後訓練程式庫，專為可擴充的 RL 而設計。NeMo RL 是更廣泛的 NeMo 框架生態系統的一環，可讓您在單一 GPU 上進行小規模實驗，也能在數千個 GPU 上部署多個節點。

詳情請參閱 NVIDIA NeMo RL。

GSM8k 資料集

在本教學課程中，您會使用 GSM8k 資料集，其中包含 8,500 個優質、語言多元的小學數學文字題。

模型會使用 GSM8k 和 GRPO，針對同一問題生成一組 n 個不同的回覆。GRPO 會將這些回覆與群組平均值進行比較。如果路徑持續正確且合乎邏輯，模型獲得的獎勵會比其他路徑多。隨著時間推移，模型會瞭解清楚說明步驟是獲得最高獎勵最可靠的方式，進而有效減少低成效答案的獎勵。

詳情請參閱「GSM8k」。

目標

本教學課程說明如何透過下列步驟，在 GKE 上使用 NeMo RL 設定 RL：

  1. 準備環境。
  2. 設定搭載 B200 或 H200 GPU 的 GKE 叢集。
  3. 設定 KubeRay 來管理分散式 Ray 叢集。
  4. 使用 Managed Lustre 儲存空間，享有高效能。
  5. 執行使用 NeMo RL 的 GRPO 訓練工作。

事前準備

  1. 安裝 Google Cloud CLI。

  2. 設定 gcloud CLI，使用您的聯合身分。

    詳情請參閱「使用聯合身分登入 gcloud CLI」。

  3. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  4. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需要具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要「專案建立者」角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

    • 建立 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) 取得這項權限。否則，您可以透過「服務使用情形管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。 

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

  7. 將角色授予使用者帳戶。針對下列每個 IAM 角色，執行一次下列指令： 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 和 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) 網路的前置字元。你可以使用任何前置字元。

設定基礎架構

在本節中，您將建立虛擬私有雲網路和 GKE 叢集。

Create a VPC network

  1. 為 gVNIC 介面建立虛擬私有雲網路：

    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. 建立 RDMA 專用的虛擬私有雲網路和子網路，包括八個 GPU 各自的子網路：

    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 -

準備儲存空間

在本節中，您將建立 Managed Lustre 執行個體，為 RL 工作負載佈建所需的高效能儲存空間。

  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

    在本教學課程中，您會使用兩個工作節點。如要變更工作站節點數量，請變更 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 叢集準備就緒後，您可以將 Ray 工作提交至 GKE 上執行的 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 模型執行 Job 的設定。為完成 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. 刪除虛擬私有雲對等互連：

    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

後續步驟