在 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 训练作业。

准备工作

  • 登录您的 Google Cloud 账号。如果您是 Google Cloud新手,请 创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  • 安装 Google Cloud CLI。

  • 如果您使用的是外部身份提供方 (IdP),则必须先使用联合身份登录 gcloud CLI

  • 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  • 创建或选择 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 项目名称。

  • 验证是否已为您的 Google Cloud 项目启用结算功能

  • 启用所需的 API:

    启用 API 所需的角色

    如需启用 API,您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin),该角色包含 serviceusage.services.enable 权限。了解如何授予角色

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
  • 安装 Google Cloud CLI。

  • 如果您使用的是外部身份提供方 (IdP),则必须先使用联合身份登录 gcloud CLI

  • 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  • 创建或选择 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 项目名称。

  • 验证是否已为您的 Google Cloud 项目启用结算功能

  • 启用所需的 API:

    启用 API 所需的角色

    如需启用 API,您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin),该角色包含 serviceusage.services.enable 权限。了解如何授予角色

    gcloud services enable container.googleapis.com storage.googleapis.com compute.googleapis.com
  • 向您的用户账号授予角色。对以下每个 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:用户账号的标识符。 例如,myemail@example.com
    • ROLE:您向用户账号授予的 IAM 角色。

准备环境

在本教程中,您将使用 Cloud Shell

  1. 前往 Google Cloud 控制台

  2. 点击 Google Cloud 控制台窗口顶部的激活 Cloud Shell 按钮。

  3. 设置以下环境变量:

    export PROJECT_ID=$(gcloud config get project)
    export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
    export CONTROL_PLANE_REGION=CONTROL_PLANE_REGION
    export NODE_ZONE=NODE_ZONE
    export CLUSTER_NAME=CLUSTER_NAME
    export GPU_TYPE=GPU_TYPE
    export MACHINE_TYPE=MACHINE_TYPE
    export KSA_NAME=generic-ksa
    export NAMESPACE=default
    export RESERVATION=RESERVATION
    export LUSTRE_NAME=LUSTRE_NAME
    export HF_TOKEN=YOUR_HF_TOKEN
    export WANDB_API_KEY=YOUR_WANDB_API_KEY
    

    替换以下值:

    • CLUSTER_NAME:GKE 集群的名称。
    • CONTROL_PLANE_REGION:GKE 集群控制平面的 Compute Engine 区域。
    • NODE_ZONE:节点的可用区。选择提供 NVIDIA B200 或 H200 GPU 的可用区
    • GPU_TYPE:您在 Compute Engine 容量预留中预留的加速器。必须是以下值之一:
      • nvidia-b200:NVIDIA B200 (180 GB)
      • nvidia-h200-141gb:NVIDIA H200 (141 GB)
    • MACHINE_TYPE:要使用的机器类型:
      • 对于 NVIDIA B200 (180 GB) GPU,请使用 a4-highgpu-8g 或更高版本。
      • 对于 NVIDIA H200 (141 GB) GPU,请使用 a3-ultragpu-8g 或更高版本。
    • RESERVATION:GPU 预留的名称。
    • LUSTRE_NAME:Lustre 实例的名称。
    • YOUR_HF_TOKEN:您的 Hugging Face 令牌。
    • YOUR_WANDB_API_KEY:您的 Wandb API 密钥。
  4. 为网络创建以下环境变量:

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

    替换以下值:

    • NETWORK-NAME:GKE 的网络名称。
    • GVNIC-NAME:gVNIC 网络名称的前缀。您可以使用任何所需的前缀。
    • RDMA-NAME:远程直接内存访问 (RDMA) 网络的网络前缀。您可以使用任何所需的前缀。

设置基础架构

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

创建 VPC 网络

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

    gcloud compute networks create ${NETWORK} --subnet-mode=auto
    gcloud compute networks create ${GVNIC_NETWORK_PREFIX}-net \
        --subnet-mode=custom
    gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-sub \
        --network=${GVNIC_NETWORK_PREFIX}-net \
        --region=${CONTROL_PLANE_REGION} \
        --range=192.168.0.0/24
    gcloud compute firewall-rules create ${GVNIC_NETWORK_PREFIX}-internal \
        --network=${GVNIC_NETWORK_PREFIX}-net \
        --action=ALLOW \
        --rules=tcp:0-65535,udp:0-65535,icmp \
        --source-ranges=192.168.0.0/16
    
  2. 为 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 > network-mapping-updated.yaml
    kubectl apply -f network-mapping-updated.yaml
    

准备存储空间

在本部分中,您将创建一个 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 > lustre-pv-updated.yaml
      kubectl apply -f lustre-pv-updated.yaml
      
    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
    cd kubernetes-engine-samples
    
  2. 导航到工作目录:

    cd ai-ml/nemo-rl-on-gke/nemoRL
    
  3. 检查 values.yaml 清单:

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #     http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    image:
      repository: "nvcr.io/nvidia/nemo-rl"
      tag: "v0.5.0" 
      pullPolicy: Always
    
    nameOverride: "kuberay"
    fullnameOverride: ""
    
    common:
      containerEnv: {}
    
    configMap:
      fluentbit:
        data:
          fluent-bit.conf: |
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/worker-*
                Tag               ray-worker
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/raylet*
                Tag               raylet
            [INPUT]
                Name              tail
                Path              /tmp/ray/session_latest/logs/*
                Exclude_Path      /tmp/ray/session_latest/logs/debug_state.txt,/tmp/ray/session_latest/logs/raylet*,/tmp/ray/session_latest/logs/worker-*
                Tag               ray-misc
            [OUTPUT]
                Name              stackdriver
                Match             *
                resource          gce_instance
                labels_key        labels
    
    # --- Head Node Configuration ---
    head:
      enableInTreeAutoscaling: false
      serviceAccountName: ""
      rayStartParams:
        dashboard-host: '0.0.0.0'
      template:
        metadata:
          annotations:
            gke-gcsfuse/volumes: "true"
            networking.gke.io/default-interface: 'eth0'
      containerEnv:
      - name: RAY_GROUP
        value: "head"
      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                3/3   Running 0      33h
    ray-cluster-kuberay-worker-grp-0-worker-gkbxw 3/3   Running 0      33h
    ray-cluster-kuberay-worker-grp-0-worker-kdg62 3/3   Running 0      33h
    
  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. 修改 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: 16num_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 树实用程序:

    apt update && apt install -y tree
    

    输出类似于以下内容:

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

清理

为避免产生费用,请删除以下资源:

helm delete ray-cluster
gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${CONTROL_PLANE_REGION} \
    --quiet
gcloud lustre instances delete ${LUSTRE_NAME} --location=${NODE_ZONE} --quiet
gcloud services vpc-peerings delete \
    --service=servicenetworking.googleapis.com \
    --network=${NETWORK}
gcloud compute addresses delete ${LUSTRE_NAME}-range --global --quiet
gcloud compute firewall-rules delete ${GVNIC_NETWORK_PREFIX}-internal --quiet
gcloud compute networks subnets delete ${GVNIC_NETWORK_PREFIX}-sub \
    --region=${CONTROL_PLANE_REGION} --quiet
gcloud compute networks delete ${GVNIC_NETWORK_PREFIX}-net --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
gcloud compute networks delete ${RDMA_NETWORK_PREFIX}-net --quiet
gcloud compute networks delete ${NETWORK} --quiet

后续步骤