Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

在 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：

准备环境。
设置具有 B200 或 H200 GPU 的 GKE 集群。
配置 KubeRay 以管理分布式 Ray 集群。
使用 Managed Lustre 实现高性能存储。
运行使用 NeMo RL 的 GRPO 训练作业。

准备工作

安装 Google Cloud CLI。

注意：如果您之前安装了 gcloud CLI，请确保通过运行 gcloud components update 来获得最新版本。
配置 gcloud CLI 以使用您的联合身份。

如需了解详情，请参阅使用联合身份登录 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，您需要拥有 serviceusage.services.enable 权限。如果您创建了项目，则可能已经通过 Owner 角色 (roles/owner) 获得了此权限。否则，您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。了解如何授予角色。
```
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：您的用户账号。如需查看示例，请参阅在 IAM 政策中表示员工池用户。
- ROLE：您向用户账号授予的 IAM 角色。
如果您还没有 Hugging Face 账号，请创建一个。
确保您拥有具有 read access 的 Hugging Face 令牌。
如果您还没有 Weights & Biases (Wandb) 账号，请创建一个。
创建 Wandb API 密钥。
确保您的 Google Cloud 项目具有足够的 B200 和 H200 GPU 配额。如需了解详情，请参阅规划 GPU 配额和 GPU 配额。

准备环境

在本教程中，您将使用 Cloud Shell。

前往 Google Cloud 控制台。
点击 Google Cloud 控制台窗口顶部的激活 Cloud Shell 按钮。
设置以下环境变量：
```
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="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_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 预留的名称。
- YOUR_LUSTRE_NAME：Lustre 实例的名称。
- YOUR_HF_TOKEN：您的 Hugging Face 令牌。
- YOUR_WANDB_API_KEY：您的 Wandb API 密钥。
为网络创建以下环境变量：
```
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 集群生成 VPC 网络名称的前缀。
- YOUR_GVNIC_NETWORK_PREFIX：gVNIC 网络名称的前缀。您可以使用任何前缀。
- YOUR_RDMA_NETWORK_PREFIX：远程直接内存访问 (RDMA) 网络的 prefix。您可以使用任何前缀。

设置基础架构

在本部分中，您将创建 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 
```