在 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 是透過經驗、探索和回饋來訓練模型，而不是靜態模仿。預先訓練會教導模型該說什麼，而人類回饋增強學習 (RLHF) 則會教導模型如何提供實用、安全且合乎邏輯的內容。RL 是基礎模型與微調模型之間的橋樑，可滿足特定用途的需求。

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

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

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

詳情請參閱「GRPO」。

NVIDIA NeMo RL

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

詳情請參閱「NVIDIA NeMo RL」。

GSM8k 資料集

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

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

詳情請參閱「GSM8k」。

目標

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

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

事前準備

  1. 登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自體驗產品的實際應用成效。新客戶還能獲得價值 $300 美元的免費抵免額，能用於執行、測試及部署工作負載。

  2. 安裝 Google Cloud CLI。

  3. 若您採用的是外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

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

    gcloud init

  5. 建立或選取 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 專案名稱。

  6. 確認 Google Cloud 專案已啟用計費功能。

  7. 啟用必要的 API (如有尚未啟用的 API)：

    啟用 API 所需的角色

    如要啟用 API，您必須具備 serviceusage.services.enable 權限。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。

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

  8. 安裝 Google Cloud CLI。

  9. 若您採用的是外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

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

    gcloud init

  11. 建立或選取 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 專案名稱。

  12. 確認 Google Cloud 專案已啟用計費功能。

  13. 啟用必要的 API (如有尚未啟用的 API)：

    啟用 API 所需的角色

    如要啟用 API，您必須具備 serviceusage.services.enable 權限。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com
    14.

  14. 將角色授予使用者帳戶。針對下列每個 IAM 角色，執行一次下列指令： roles/container.admin, roles/iam.serviceAccountAdmin, roles/storage.admin, roles/servicenetworking.networksAdmin, roles/compute.networkAdmin, roles/lustre.admin 

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

    更改下列內容：

    • PROJECT_ID：專案 ID。
    • USER_IDENTIFIER：使用者帳戶 ID。 例如：myemail@example.com。
    • ROLE：授予使用者帳戶的 IAM 角色。
  15. 如果還沒有帳戶，請建立 Hugging Face 帳戶。
  16. 確認您有 read access 的 Hugging Face 權杖。
  17. 如果沒有 Weights & Biases (Wandb) 帳戶，請建立一個。
  18. 建立 Wandb API 金鑰。
  19. 請確認 Google Cloud 專案有足夠的 B200 和 H200 GPU 配額。詳情請參閱「規劃 GPU 配額」和「GPU 配額」。
  20. 確認 Google Cloud 專案有足夠的 Managed Lustre 配額。詳情請參閱「 Managed Lustre 配額和限制」。

準備環境

在本教學課程中，您會使用 Cloud Shell。

  1. 前往 Google Cloud 控制台。

  2. 按一下 Google Cloud 控制台視窗頂端的「Activate Cloud Shell」(啟用 Cloud Shell) 按鈕。

  3. 請設定下列環境變數：

    export CONTROL_PLANE_REGION="YOUR_REGION"
export NODE_ZONE="YOUR_ZONE"
export LUSTRE_ZONE="YOUR_LUSTRE_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GPU_TYPE="YOUR_GPU_TYPE"
export MACHINE_TYPE="YOUR_MACHINE_TYPE"
export CPU_MACHINE_TYPE="YOUR_CPU_MACHINE_TYPE"
export KSA_NAME="generic-ksa"
export NAMESPACE="default"
export RESERVATION="YOUR_RESERVATION_NAME"
export LUSTRE_NAME="YOUR_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_ZONE：節點所在的可用區。選取提供 NVIDIA B200 或 H200 GPU 的區域。
    • YOUR_LUSTRE_ZONE：Managed Lustre 執行個體的可用區。選取提供 Managed Lustre 的區域。為提升效能、縮短延遲時間及避免網路輸出費用，請使用與 YOUR_ZONE 相同的可用區。
    • 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_CPU_MACHINE_TYPE：CPU 預設節點集區的 Compute Engine 機型。使用叢集可用區支援的 highmem 機型，且至少要有 16 個 vCPU，例如 n4-highmem-16 或 n2d-highmem-16。
    • YOUR_RESERVATION_NAME：GPU 預留項目的名稱。
    • YOUR_LUSTRE_NAME：Lustre 執行個體的名稱。
    • YOUR_HF_TOKEN：您的 Hugging Face 權杖。
    • YOUR_WANDB_API_KEY：您的 Wandb API 金鑰。

  4. 為網路建立下列環境變數：

    export NETWORK_PREFIX="YOUR_NETWORK_NAME"
export NETWORK="${NETWORK_PREFIX}-nemo-rl"
export GVNIC_NETWORK_PREFIX="YOUR_GVNIC_NETWORK_PREFIX"
export RDMA_NETWORK_PREFIX="YOUR_RDMA_NETWORK_PREFIX"

    替換下列值：

    • YOUR_NETWORK_NAME：用於產生 GKE 叢集虛擬私有雲網路名稱的前置字串。
    • YOUR_GVNIC_NETWORK_PREFIX：gVNIC 網路名稱的前置字串。你可以使用任何前置字元。
    • YOUR_RDMA_NETWORK_PREFIX：遠端直接記憶體存取 (RDMA) 網路的前置字元。你可以使用任何前置字元。

  5. 如果尚未安裝，請為 kubectl CLI 安裝 KubeRay 外掛程式 (Beta 版)：

    KUBECTL_RAY_VERSION="v1.3.2"
curl -LO "https://github.com/ray-project/kuberay/releases/download/${KUBECTL_RAY_VERSION}/kubectl-ray_${KUBECTL_RAY_VERSION}_linux_amd64.tar.gz"
tar -xvf "kubectl-ray_${KUBECTL_RAY_VERSION}_linux_amd64.tar.gz"
sudo cp kubectl-ray /usr/local/bin/
sudo chmod +x /usr/local/bin/kubectl-ray
rm -f "kubectl-ray_${KUBECTL_RAY_VERSION}_linux_amd64.tar.gz" kubectl-ray

    這個外掛程式提供指令，可管理 Kubernetes 上的 Ray 工作負載，例如開啟互動式工作階段及提交工作。

設定基礎架構

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

建立虛擬私有雲網路

  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=${CPU_MACHINE_TYPE} \
    --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/64f0eeb0dc1d980eaef8939584f832bb10e3d482/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 執行個體，為強化學習工作負載佈建所需的高效能儲存空間。

  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=${LUSTRE_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=${LUSTRE_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}/${LUSTRE_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
git -C kubernetes-engine-samples checkout 87fa0575c977b96f6d3af2cab58c9128f3cf0dc3

  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 \
  --set head.resources.requests.cpu=12 \
  --set head.resources.requests.memory=96Gi \
  --set head.resources.limits.cpu=12 \
  --set head.resources.limits.memory=96Gi

    在本教學課程中，您會使用兩個工作節點。如要變更 worker 節點數量，請變更 REPLICA_COUNT 值。

  5. 確認代管 Ray 主要和工作節點的 Pod 正在執行：

    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

    Pod 大約需要 10 分鐘才能啟動。

  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 資料集執行 Gemma 3 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. 在初始終端機視窗中提交工作：

    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 || true

  3. 刪除 Lustre 檔案系統：

    gcloud lustre instances delete ${LUSTRE_NAME} \
    --location=${LUSTRE_ZONE} \
    --quiet || true

  4. 刪除虛擬私有雲對等互連：

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

  5. 刪除 Lustre 私人 IP 位址範圍：

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

  6. 刪除 RDMA 和 gVNIC 子網路：

    gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
    --region=${CONTROL_PLANE_REGION} \
    --quiet || true

for N in $(seq 0 7); do
  gcloud compute networks subnets delete ${RDMA_NETWORK_PREFIX}-sub-$N \
    --region=${CONTROL_PLANE_REGION} --quiet || true &
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

  # Skip if empty, named 'default', or not starting with a lowercase letter (invalid resource name)
  if [[ "${NW}" == "default" || ! "${NW}" =~ ^[a-z] ]]; then
    echo "Skipping protected/invalid network: '${NW}'"
    continue
  fi

  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

後續步驟