在 GKE 上使用 NVIDIA NeMo RL 微调和扩缩强化学习

本教程介绍了如何在 Google Kubernetes Engine (GKE) 上编排分布式强化学习 (RL) 训练环境。您可以使用 Ray 和 NVIDIA NeMo RL 框架设置分布式训练环境，以对模型进行微调。

本教程重点介绍如何使用 Ray 和 NeMo RL 在 GKE 上运行 Group Relative Policy Optimization (GRPO) 训练流水线。GRPO 是一种旨在提升模型推理能力的强化学习算法。这种内存高效型算法通过以下方式简化了 RL 流程：消除 Critic（即价值模型），并使用基于相对组的计算。

在运行本教程之前，请先完成在 GKE 上使用 Verl 对强化学习进行微调和扩缩教程。本教程使用的集群设置和配置与“使用 Verl 进行 RL 微调和伸缩”教程相同。

背景

以下部分简要介绍了本教程中使用的概念。

强化学习 (RL)

RL 通过经验、探索和反馈来训练模型，而不是静态模仿。虽然预训练可以教会模型说什么，但基于人类反馈的强化学习 (RLHF) 可以教会模型如何提供有用的、安全的、符合逻辑的回答。RL 可作为基础模型与针对特定用例微调的模型之间的桥梁。

如需了解详情，请参阅什么是强化学习？

群组相对政策优化 (GRPO)

GRPO 是一种由 DeepSeek 推广的算法，它通过移除 Critic 模型，为 LLM 对齐提供了一种内存高效的近端策略优化 (PPO) 替代方案。与 Critic 网络不同，GRPO 会针对同一提示生成一组回答，并使用该组回答的平均奖励作为基准。

如需了解详情，请参阅 GRPO。

NVIDIA NeMo RL

NeMo RL 是 NVIDIA 的开源后训练库，旨在实现可伸缩的 RL。作为更广泛的 NeMo 框架生态系统的一部分，NeMo RL 既支持在单个 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 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (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 权限。如果您创建了项目，则可能已经通过 Owner 角色 (roles/owner) 获得了此权限。否则，您可以通过 Service Usage Admin 角色 (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

    替换以下内容：

    • PROJECT_ID：您的项目 ID。
    • USER_IDENTIFIER：您的用户 账号。如需查看示例，请参阅 在 IAM 政策中表示员工池用户。
    • ROLE：您向用户账号授予的 IAM 角色。
  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) 网络的 prefix。您可以使用任何前缀。

设置基础架构

在本部分中，您将创建 VPC 网络和 GKE 集群。

创建 VPC 网络

  1. 为 gVNIC 接口创建 VPC 网络：

    gcloud compute networks create ${NETWORK} --subnet-mode=auto

gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
    --subnet-mode=custom

gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --region=${CONTROL_PLANE_REGION} \
    --range=192.168.0.0/24

gcloud compute firewall-rules create ${GVNIC_NETWORK_PREFIX}-internal \
    --network=${GVNIC_NETWORK_PREFIX}-net \
    --action=ALLOW \
    --rules=tcp:0-65535,udp:0-65535,icmp \
    --source-ranges=192.168.0.0/16

  2. 为 RDMA 创建 VPC 网络和子网，其中包括 8 个子网，用于 8 个 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 集群准备就绪后，您可以向 GKE 上正在运行的 Ray 集群提交 Ray 作业。NeMo RL 会在执行 RL 训练作业期间自动下载模型。

如需提交 Ray 作业，请启动互动式会话以执行该作业。

  1. 如需与 Ray 集群建立本地连接，请运行以下命令：

    kubectl ray session ray-cluster-kuberay

    此命令会在您的本地机器与 GKE 集群中的 Ray 头节点之间启动端口转发。请注意，此会话处于活跃状态时，您的终端将被占用；如需继续操作，请打开单独的终端实例。

  2. 在另一个终端中，前往 kubernetes-engine-samples/ai-ml/nemo-rl-on-gke/nemoRL/gemma3-27b-it 并修改 gemma3-27b-gsm8k.sh 文件：

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

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

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

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

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

git pull && git checkout main

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

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

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

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

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

echo ""
echo "Job submission complete."

    在 gemma3-27b-gsm8k.sh 文件中替换以下值：

    • YOUR_WANDB_API_KEY：您的 WandB API 密钥。
    • YOUR_HF_TOKEN：您的 Hugging Face 令牌。

    在此文件中，您可以查看在 GSM8k 数据集上运行使用 gemma3-27b-it 模型的作业的配置。为了完成 GRPO 训练流水线，此脚本定义了以下参数：

    • num_prompts_per_step: 16 和 num_generations_per_prompt: 32：Gemma3-27b-it 模型会针对每个提示生成大量回答。 在此配置中，模型会生成 512 个回答（16 × 32 = 512）。
    • policy.generation.colocated.enabled=False：此参数会停用同位生成功能，这意味着模型不会在与训练过程相同的节点中生成回答。在标准强化学习中，同一 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 树实用程序：

    apt update && apt install -y tree

  3. 列出 ray-head 容器的目录结构：

    tree /data/nemo_rl_gemma3_27b_3_17/

    输出类似于以下内容：

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

6 directories, 39 files

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除各个资源，或者删除包含这些资源的项目。

删除资源

  1. 删除 Slurm 集群：

    helm delete ray-cluster

  2. 删除 GKE 集群：

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

  3. 删除 Lustre 文件系统：

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

  4. 删除 VPC 对等互连：

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

  5. 删除 Lustre 专用 IP 地址范围：

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

  6. 删除 RDMA 和 gVNIC 子网：

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

for N in $(seq 0 7); do
  gcloud compute networks subnets delete ${RDMA_NETWORK_PREFIX}-sub-$N \
    --region=${CONTROL_PLANE_REGION} --quiet &
done
wait

  7. 删除防火墙规则和网络：

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

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

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

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

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

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

        sleep 3
      done

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

删除项目

删除 Google Cloud 项目：

gcloud projects delete PROJECT_ID

后续步骤