使用 GKE 和 Managed Lustre 扩缩强化学习

1. 简介

如果您希望直接运行打包的脚本，而无需逐步教程，可以在 GoogleCloudPlatform/devrel-demos 代码库中找到这些脚本。

在此 Codelab 中，您将学习如何使用 Google Kubernetes Engine (GKE) 和 Managed Lustre 部署高性能强化学习 (RL) 训练流水线。

强化学习工作负载（尤其是使用组相对策略优化 [GRPO] 等算法的工作负载）在“经验生成”期间会生成大量数据，并且需要频繁进行检查点设置。标准对象存储在这些 I/O 突发期间可能会造成瓶颈，导致昂贵的加速器处于闲置状态。

您将使用并行文件系统 Managed Lustre 来消除这些瓶颈，并实现更高的训练吞吐量。

您将执行的操作

• 为基于 GPU 的 Ray 集群配置环境变量。
• 使用 Cluster Toolkit 在 GKE 上预配 Spot GPU 集群和Managed Lustre 实例。
• 部署 KubeRay 集群并装载 Lustre 文件系统。
• 提交 NeMo-RL 训练工作负载。
• 使用 Cloud Monitoring 观察高吞吐量和低检查点延迟时间。

所需条件

• 网络浏览器，例如 Chrome。
• 启用了结算功能的 Google Cloud 项目。

此 Codelab 使用高端资源，包括 Spot GPU 和 Managed Lustre。虽然我们使用 Spot 容量来降低成本，但如果这些资源一直处于运行状态，费用可能会很高。请务必在实验结束时按照“清理”步骤操作，以免产生意外费用。

本 Codelab 适合熟悉 GKE 和存储概念的高级技术用户、平台工程师和 AI 研究人员。

预计总时长：45 至 60 分钟，外加 2 小时的培训时间

2. 准备工作

创建 Google Cloud 项目

1. 在 Google Cloud 控制台中，选择或创建 Google Cloud 项目。
2. 确保您的 Cloud 项目已启用结算功能。

启动 Cloud Shell

Cloud Shell 是在 Google Cloud 中运行的命令行环境，预加载了必要的工具。

1. 点击 Google Cloud 控制台顶部的激活 Cloud Shell。
2. 连接到 Cloud Shell 后，验证您的身份验证：

   gcloud auth list
   

3. 确认您的项目已配置：

   gcloud config get project
   

4. 如果项目未按预期设置，请进行设置：

   export PROJECT_ID=<YOUR_PROJECT_ID>
   gcloud config set project $PROJECT_ID
   

安装 Cluster Toolkit

此 Codelab 使用 Cluster Toolkit (gcluster) 部署 GKE 集群。如需了解如何设置 Cluster Toolkit，请参阅 Cluster Toolkit 设置指南。

启用 API

在 Cloud Shell 中运行以下命令，以启用所有必需的 API：

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

3. 配置环境变量

为了使此 Codelab 中的命令保持一致，请设置一些环境变量。

创建一个名为 env.sh 的文件，并使用您的配置填充该文件。您可以使用以下模板：

# Environment Variables for the RL Demo execution
export PROJECT_ID="{{'<var>'}}PROJECT_ID{{'</var>'}}"
export ZONE="us-east1-b"
export REGION="us-east1"
export CLUSTER_NAME="ray-a4-gpu-spot"
export HF_TOKEN="{{'<var>'}}YOUR_HF_TOKEN{{'</var>'}}" # Required for downloading models
export WANDB_API_KEY="{{'<var>'}}YOUR_WANDB_API_KEY{{'</var>'}}" # Optional

# Topology defaults
export NUM_NODES="8"
export GPUS_PER_NODE="8" # Fixed for A4/B200 architecture

HF_TOKEN 是您的 Hugging Face 用户访问令牌，下载模型时需要用到此令牌。您可以按照 Hugging Face 安全令牌中的说明获取令牌。

将 <YOUR_PROJECT_ID> 和 <YOUR_HF_TOKEN> 替换为您的实际值。

WANDB_API_KEY 为可选项。它用于 Weights & Biases，这是一款用于监控和直观呈现机器学习实验的工具。如果您没有账号或不想使用账号，可以将此字段留空或保持原样。

获取文件以将变量加载到当前会话中：

source env.sh

4. 使用 Cluster Toolkit 部署 GKE 集群和 Managed Lustre

在此步骤中，您将使用 Cluster Toolkit (gcluster) 部署一个具有 Spot GPU 的 GKE 集群，并使用 Lustre CSI 驱动程序和预配置的 PersistentVolumeClaim (lustre-pvc) 自动预配 Managed Lustre 存储空间。

准备蓝图

在部署之前，请查看 examples/gke-a4/gke-a4.yaml 蓝图（如需了解详情，请参阅创建 A4 集群）：

1. 启用 Managed Lustre：取消 gke-a4.yaml 中 managed-lustre 和 lustre-pvc 模块部分的注释。
2. 启用 RayOperator 插件：在 gke-a4.yaml 中，将 gke_cluster 模块设置下的 enable_ray_operator: true 设置为 1。

部署基础设施

设置 Cloud Shell 访问权限的授权 CIDR，并使用 gcluster deploy 进行部署：

export AUTHORIZED_CIDR="$(curl -s ifconfig.me)/32"

gcluster deploy examples/gke-a4/gke-a4.yaml \
  --vars project_id=${PROJECT_ID},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},static_node_count=${NUM_NODES},authorized_cidr=${AUTHORIZED_CIDR},spot=true

在 --vars 中设置 spot=true 可使用 Spot 容量动态预配 GPU 节点池，而无需手动修改 reservation_affinity。请注意，竞价型容量可能会被抢占。对于生产环境，请省略 spot=true 以使用按需容量或配置预留。

等待部署完成。Cluster Toolkit 会在一个协调的部署中自动预配 VPC 网络、专用服务访问通道 (PSA) 对等互连、Managed Lustre 文件系统、Lustre CSI 驱动程序、RayOperator 插件和 Kubernetes 存储空间声明 (lustre-pvc)。

5. 在 GKE 上部署 Ray 集群

在此步骤中，您将在 GKE 节点上部署 KubeRay 集群，并使用由 Cluster Toolkit 自动预配的 PersistentVolumeClaim (lustre-pvc) 装载 Lustre 文件系统。

创建 RayCluster 配置

创建一个名为 ray-cluster.yaml 的文件。此规范指定了 KubeRay 头节点和工作器节点，使用 nvidia-b200 加速器类型并将 Lustre 卷装载到 /lustre。

cat << EOF > ray-cluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ${CLUSTER_NAME}
  namespace: default
spec:
  rayVersion: '2.54.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-head
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          ports:
          - containerPort: 6379
            name: gcs-server
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "32"
              memory: "1000Gi"
            requests:
              cpu: "8"
              memory: "64Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
  workerGroupSpecs:
  - groupName: gpu-worker-group
    replicas: ${NUM_NODES}
    minReplicas: ${NUM_NODES}
    maxReplicas: ${NUM_NODES}
    rayStartParams: {}
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-worker
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          resources:
            limits:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
            requests:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
        - name: dshm
          emptyDir:
            medium: Memory
EOF

连接到集群

确保您的 Cloud Shell 会话已通过 GKE 集群的身份验证：

gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --region ${REGION} \
  --project ${PROJECT_ID}

应用 RayCluster 配置

应用 Ray 集群配置：

kubectl apply -f ray-cluster.yaml

验证集群状态

监控 pod 的创建情况：

kubectl get pods -w

等待头 Pod 和工作器 Pod 变为 Running。

6. 提交强化学习工作负载

在此步骤中，您将向 Ray 集群提交 NeMo-RL GRPO 训练作业。

连接到 Ray 信息中心

如需提交作业和查看指标，您需要连接到 Ray 信息中心。由于信息中心位于 GKE 中，因此请使用端口转发从 Cloud Shell 访问它：

# Run this in a separate Cloud Shell tab or in the background
kubectl port-forward service/${CLUSTER_NAME}-head-svc 8265:8265 &

创建执行脚本

创建一个名为 run_nemo_rl.sh 的文件。此脚本将在 Ray 集群工作器上执行。我们使用 cat << EOF 来填充您之前设置的环境变量。

cat << EOF > run_nemo_rl.sh
#!/bin/bash
set -ex

# Override job runtime conflicts (NeMo-RL passes os.environ to ray.init)
export RAY_OVERRIDE_JOB_RUNTIME_ENV=1

echo "--- Running on Ray Cluster ---"
cd /opt/nemo-rl

# Ensure directories exist on the high-speed Lustre drive
mkdir -p /lustre/huggingface_cache
mkdir -p /lustre/nemo_rl_qwen_72b_ds_cp

echo "Launching NeMo-RL GRPO training..."
uv run python examples/run_grpo_math.py \
  --config examples/configs/grpo_math_70B_megatron.yaml \
  policy.model_name='Qwen/Qwen2.5-72B-Instruct' \
  policy.megatron_cfg.converter_type='Qwen2ForCausalLM' \
  logger.wandb_enabled=False \
  cluster.num_nodes=${NUM_NODES} \
  cluster.gpus_per_node=${GPUS_PER_NODE} \
  logger.wandb.name='nemo-rl-grpo-test1' \
  grpo.max_num_steps=20 \
  grpo.num_generations_per_prompt=8 \
  grpo.num_prompts_per_step=32 \
  policy.train_global_batch_size=256 \
  checkpointing.enabled=True \
  checkpointing.save_period=2 \
  checkpointing.keep_top_k=2 \
  checkpointing.metric_name=null \
  checkpointing.checkpoint_dir=/lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 \
  data.dataset_name='DeepScaler'
EOF
chmod +x run_nemo_rl.sh

创建 Ray 忽略文件

创建一个 .rayignore 文件，以防止 Ray 上传大型或不必要的目录：

cat << EOF > .rayignore
cluster-toolkit/
.git/
*.sh.log
EOF

创建运行时环境配置

创建一个 JSON 文件，以将环境变量传递给 Ray 作业：

cat << EOF > ray_runtime_env_nemo.json
{
  "env_vars": {
    "HF_TOKEN": "${HF_TOKEN}",
    "WANDB_API_KEY": "${WANDB_API_KEY}",
    "HF_HOME": "/lustre/huggingface_cache",
    "GLOO_SOCKET_IFNAME": "eth0",
    "NCCL_SOCKET_IFNAME": "eth0"
  }
}
EOF

提交作业

使用 Ray CLI 将作业提交到信息中心端点。如果在 Cloud Shell 中找不到 ray 命令，您可以使用 pip install ray 安装该命令：

ray job submit \
    --address="http://localhost:8265" \
    --working-dir . \
    --runtime-env ray_runtime_env_nemo.json \
    -- bash run_nemo_rl.sh

您会在 Cloud Shell 终端中看到流式日志。作业将加载模型、初始化 Ray 工作器，并开始 GRPO 训练循环。

7. 监控训练表现

在此步骤中，您将观察 Lustre 文件系统在训练和设置检查点期间的性能。

查看训练日志

随着训练的进行，您会看到日志表明正在将检查点保存到 /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1。请注意，检查点创建是异步进行的，不会长时间阻塞 Ray worker。

如需查看检查点创建速度，请查找指示已保存检查点的日志行。

在 Cloud 控制台中查看 Lustre 指标

如需查看 Lustre 实例的指标，请执行以下操作：

1. 在 Google Cloud 控制台中，搜索 Managed Service for Lustre。
2. 点击您的实例名称（例如 ${CLUSTER_NAME}-lustre 或 rl-demo-gpu-lustre）。
3. 点击监控标签页。

您可以在此处查看：

• 吞吐量（字节/秒）：查看在检查点创建期间出现的峰值。
• 容量：监控检查点占用的空间量。

Lustre 能够以极高的速度写入，在最短的时间内写入检查点

8. 清理资源

此步骤对于避免持续产生费用至关重要。如果 GPU 集群和 Managed Lustre 持续运行，可能会产生高昂的费用。

在 Cloud Shell 中运行以下命令，以一步销毁所有已配置的基础设施（GKE 集群、GPU 节点池、Managed Lustre 实例和 VPC 网络）：

gcluster destroy "${CLUSTER_NAME}"

此命令会在前台同步运行，拆除部署管理的所有基础架构，并将进度日志输出到终端。等待命令完全运行完毕，然后再关闭 Cloud Shell 会话。

9. 恭喜

您已成功完成 使用 GKE 和 Managed Lustre 扩缩强化学习 Codelab！

您学到的内容

• 如何使用 Cluster Toolkit 预配具有 Spot 实例和 Managed Lustre 存储空间的 GKE GPU 集群。
• 如何部署 KubeRay 集群并挂载 Lustre 存储空间。
• 如何提交 NeMo-RL GRPO 训练工作负载。
• 如何在训练期间观察存储性能。

后续步骤

• 探索更多 NVIDIA NeMo-RL 功能。
• 详细了解 Google Cloud AI Hypercomputer。
• 查看 Managed Service for Lustre 文档。