本教程介绍了如何在 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:
- 准备环境。
- 设置具有 B200 或 H200 GPU 的 GKE 集群。
- 配置 KubeRay 以管理分布式 Ray 集群。
- 使用 Managed Lustre 实现高性能存储。
- 运行使用 NeMo RL 的 GRPO 训练作业。
准备工作
- 登录您的 Google Cloud 账号。如果您是 Google Cloud新手,请 创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
-
安装 Google Cloud CLI。
-
如果您使用的是外部身份提供方 (IdP),则必须先使用联合身份登录 gcloud CLI。
-
如需初始化 gcloud CLI,请运行以下命令:
gcloud init -
选择或创建项目所需的角色
- 选择项目:选择项目不需要特定的 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 项目名称。
启用所需的 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 -
选择或创建项目所需的角色
- 选择项目:选择项目不需要特定的 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 项目名称。
启用所需的 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.admingcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
替换以下内容:
PROJECT_ID:您的项目 ID。USER_IDENTIFIER:用户账号的标识符。 例如,myemail@example.com。ROLE:您向用户账号授予的 IAM 角色。
- 如果您还没有 Hugging Face 账号,请创建一个。
- 确保您拥有 Hugging Face 令牌。
- 确保您的项目具有足够的 B200 和 H200 GPU 配额。如需了解详情,请参阅规划 GPU 配额和 GPU 配额。
准备环境
在本教程中,您将使用 Cloud Shell。
前往 Google Cloud 控制台。
点击 Google Cloud 控制台窗口顶部的激活 Cloud Shell 按钮。
设置以下环境变量:
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或更高版本。
- 对于 NVIDIA B200 (180 GB) GPU,请使用
RESERVATION:GPU 预留的名称。LUSTRE_NAME:Lustre 实例的名称。YOUR_HF_TOKEN:您的 Hugging Face 令牌。YOUR_WANDB_API_KEY:您的 Wandb API 密钥。
为网络创建以下环境变量:
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 网络
为 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为 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。
创建 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}获取集群的凭据:
gcloud container clusters get-credentials ${CLUSTER_NAME} \ --location=${CONTROL_PLANE_REGION}创建 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安装 NCCL RDMA 安装程序:
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-rdma-installer.yaml
配置网络映射
将以下清单保存为
network-mapping.yaml:应用清单:
envsubst < network-mapping.yaml > network-mapping-updated.yaml kubectl apply -f network-mapping-updated.yaml
准备存储空间
在本部分中,您将创建一个 Managed Lustre 实例,该实例可提供 RL 工作负载所需的高性能存储空间。
为专用服务访问通道分配 IP 地址范围:
gcloud compute addresses create ${LUSTRE_NAME}-range \ --global --purpose=VPC_PEERING \ --prefix-length=20 --network=${NETWORK}连接对等互连:
gcloud services vpc-peerings connect \ --service=servicenetworking.googleapis.com \ --ranges=${LUSTRE_NAME}-range \ --network=${NETWORK}创建 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使用 Managed Lustre CSI 驱动程序访问现有 Managed Lustre 实例
提取 Managed Lustre 实例的 IP 地址。
export LUSTRE_IP=$(gcloud lustre instances describe ${LUSTRE_NAME} \ --location=$NODE_ZONE --format="value(mountPoint)" | awk -F'@' '{print $1}')检查
lustre-pv.yaml清单。应用清单:
envsubst < lustre-pv.yaml > lustre-pv-updated.yaml kubectl apply -f lustre-pv-updated.yaml检查
lustre-pvc.yaml清单。应用清单:
kubectl apply -f lustre-pvc.yaml
部署 RayCluster
在本部分中,您将克隆示例代码库、准备清单并部署 Ray 集群:
克隆示例代码库:
git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git cd kubernetes-engine-samples导航到工作目录:
cd ai-ml/nemo-rl-on-gke/nemoRL检查
values.yaml清单:根据您在本教程中使用的加速器,将
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 信息中心。工作器节点运行训练作业。
- NVIDIA B200 (180 GB):
部署 Ray 集群:
export REPLICA_COUNT=2 helm install ray-cluster . \ --set additionalWorkerGroups.worker-grp-0.replicas=$REPLICA_COUNT在本教程中,您将使用两个工作器节点。如果您想更改工作器节点的数量,请更改
REPLICA_COUNT值。验证工作器节点和头节点是否正在运行:
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验证 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 作业,请启动互动式会话以执行该作业。
如需与 Ray 集群建立本地连接,请运行以下命令:
kubectl ray session ray-cluster-kuberay此命令会在您的本地机器与 GKE 集群中的 Ray 头节点之间启动端口转发。请注意,此会话处于活跃状态时,您的终端将被占用;如需继续操作,请打开单独的终端实例。
修改
gemma3-27b-gsm8k.sh文件:在
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:此形参会停用同位生成功能,这意味着模型不会在与训练过程相同的节点中生成响应。在标准 RL 中,相同的 GPU 会同时处理训练和生成。在此 NeMo RL 设置中,您可以使用policy.generation.colocated.resources形参专门管理用于 vLLM 推理的特定节点,而集群的其余部分则专注于繁重的训练计算。通过分离这些工作负载,您可以防止内存密集型训练缓冲区与计算密集型推理工作负载之间发生资源争用。
如需提交作业,请运行以下命令:
bash gemma3-27b-it/gemma3-27b-gsm8k.sh作业运行期间,输出会显示训练结果、时间安排和效果指标。
监控 GRPO 作业的运行状况
Ray 完成作业后,NeMo RL 会将检查点存储在配置的路径中。
如需检查 GRPO 作业的输出,请创建与
ray-head容器的 SSH 会话:kubectl exec -it $(kubectl get pods -l ray.io/node-type=head -o name) -c ray-head -- bash在
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