运行预构建的工作负载

如果您对 Vertex AI Managed Training 感兴趣，请与您的销售代表联系以获取访问权限。

本指南介绍了如何在受管理的训练集群上使用 NVIDIA NeMo 生态系统进行端到端生成式 AI 模型开发。本文档针对以下不同但相关的工作流程提供了分步说明，每个工作流程都有自己的专用部分：

  • NVIDIA NeMo：对于基础模型开发，请按照以下说明执行大规模预训练、持续预训练 (CPT) 和监督微调 (SFT)。
  • NVIDIA NeMo-RL：对于模型对齐和偏好调优，请使用此部分应用强化学习 (RL) 等高级技术，使模型与人类指令和偏好保持一致。

无论您是从头构建模型还是优化现有模型，本文档都会引导您完成以下操作：设置环境、管理容器化作业，以及在集群上启动训练脚本。

NVIDIA NeMo

NVIDIA NeMo 框架是一个端到端平台，用于构建、自定义和部署生成式 AI 模型。本指南的这一部分专门面向专注于模型开发基础阶段的开发者和研究人员。它提供了有关如何使用 NeMo 在受管理的训练集群上执行大规模预训练、持续预训练 (CPT) 和监督微调 (SFT) 的分步说明。

本训练集群指南提供了使用 NeMo 框架运行训练作业的完整工作流。该流程分为两个主要部分：环境的初始一次性设置和启动作业的重复步骤。

设置环境

在启动作业之前，您需要准备好环境，确保您拥有容器映像和必要的训练脚本。

准备容器映像

您可以选择使用预构建的容器映像（推荐），也可以构建自定义容器映像。

使用预构建的映像（推荐）

预构建的容器映像以 .squashfs 格式提供。将适合您所在区域的相应映像复制到工作目录中。

# Example for the US region
gcloud storage cp gs://vmds-containers-us/nemo_squashfs/nemo-20250721.sqsh .
构建自定义容器（高级）

只有在预构建容器无法满足您的需求时，才应按照以下步骤操作。 此过程将指导您使用 enroot 将自定义容器映像转换为 .squashfs 格式。

第 1 步：使用 Google Cloud进行身份验证。

使用以下命令确保您的 Google Cloud 用户账号和托管映像的 Docker 注册表均已通过身份验证：

gcloud auth login
gcloud auth configure-docker us-docker.pkg.dev

第 2 步：创建转化脚本。

创建一个名为 enroot-convert.sh 的文件，并添加以下脚本内容。在运行此脚本之前，您必须更新 REMOTE_IMG 和 LOCAL_IMG 变量，使其指向您的容器映像和所选的输出路径。

#!/bin/bash

#SBATCH --gpus-per-node=8
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH --ntasks-per-node=1

# Run this script on the slurm login node:
# sbatch -N 1 enroot-convert.sh

set -x
set -e

# The remote docker image URI.
REMOTE_IMG="docker://us-docker.pkg.dev/{YOUR_CONTAINER_IMG_URI}:{YOUR_CONTAINER_IMAGE_TAG}"

# The local path to the to be imported enroot squash file.
LOCAL_IMG="${HOME}/my_nemo.sqsh"

# The path to the enroot config file.
TMP_ENROOT_CONFIG_PATH="/tmp/\$(id -u --name)/config/enroot"

# Download the docker image to each node.
srun -l -N "${SLURM_NNODES}" \
bash -c "
mkdir -p ${TMP_ENROOT_CONFIG_PATH};
echo 'machine us-docker.pkg.dev login oauth2accesstoken password $(gcloud auth print-access-token)' > ${TMP_ENROOT_CONFIG_PATH}/.credentials;
rm -f ${LOCAL_IMG};
ENROOT_CONFIG_PATH=${TMP_ENROOT_CONFIG_PATH} ENROOT_MAX_PROCESSORS=$(( $(nproc) / 2 )) enroot import -o ${LOCAL_IMG} ${REMOTE_IMG};
"

第 3 步：运行脚本并验证输出。

在 Slurm 登录节点上执行脚本。

sbatch -N 1 enroot-convert.sh

作业完成后，在名为 slurm-<JOB_ID>.out 的文件中查找转换日志，并在您为 LOCAL_IMG 指定的路径中查找最终容器映像。

下载训练 recipe

训练 recipe 存储在私有 googlesource.com 代码库中。 如需通过 Git 命令行访问这些凭证，您必须先生成身份验证凭证。

  1. 生成身份验证凭证。

    访问以下网址，然后按照屏幕上的说明操作。此命令会将本地环境配置为使用代码库进行身份验证。https://www.googlesource.com/new-password

  2. 克隆代码库。

    凭证通过身份验证后，运行以下命令下载 recipe。

    git clone https://vertex-model-garden.googlesource.com/vertex-oss-training

启动训练作业

设置好环境后，您就可以启动训练作业了。

第 1 步：设置环境变量

您的作业可能需要以下环境变量：

  • 需要 HF_TOKEN 才能从 Hugging Face 下载模型和数据集。
  • 必须有 WANDB_API_KEY 才能使用 Weights & Biases 进行实验分析。
export HF_TOKEN=YOUR_HF_TOKEN
export WANDB_API_KEY=YOUR_WANDB_API_KEY

第 2 步：运行启动脚本

导航到工作目录，然后运行 run.py 脚本以启动作业。此示例使用 Llama 3.1-2b 启动演示训练作业。

# Set the working directory
export WORK_DIR=$HOME/vertex-oss-training/nemo
cd $WORK_DIR

gcloud storage cp
gs://vmds-containers-<region>/nemo_squashfs/nemo-20250721.sqsh nemo-demo.sqsh

# Launch the training job
export NEMORUN_HOME=$WORK_DIR && \
python3 run.py -e slurm --slurm-type hcc-a3m --partition a3m \
  -d $WORK_DIR -i $WORK_DIR/nemo-demo.sqsh \
  -s pretrain/llama3p1_2b_pt.py -n 2 \
  --experiment-name nemo-demo-run

启动参数

  • --slurm-type 会根据集群类型（例如 hcc-a3m、hcc-a3u、hcc-a4）进行设置。
  • --partition 必须设置为可用的分区。您可以使用 sinfo 命令检查分区名称。
  • 如果设置了 --log-dir、--cache-dir 和 --data-dir，run.py 脚本会自动将这几个目录装载到 Docker 容器。

监控作业状态和日志

启动作业后，系统会显示一个状态块：

Experiment Status for nemo-demo-run_1753123402

Task 0: nemo-demo-run
- Status: RUNNING
- Executor: SlurmExecutor on @localhost
- Job id: 75
- Local Directory: $NEMORUN_HOME/experiments/nemo-demo-run/nemo-demo-run_1753123402/nemo-demo-run

执行日志会写入状态输出中 Local Directory 字段所示的路径。例如，您可以在类似于以下路径的位置找到日志文件：

$NEMORUN_HOME/experiments/nemo-demo-run/nemo-demo-run_1753123402/nemo-demo-run/<JOB_ID>.log

常见错误和解决方案

本部分介绍了作业执行期间可能出现的常见问题，并提供了解决这些问题的推荐步骤。

分区无效错误

默认情况下，作业会尝试在常规分区上启动。如果常规分区不存在或不可用，作业将失败并显示以下错误：

sbatch: error: invalid partition specified: general
sbatch: error: Batch job submission failed: Invalid partition name specified

解决方案：

在启动命令中使用 --partition 或 -p 实参指定可用分区。如需查看可用分区列表，请在 Slurm 登录节点上运行 sinfo 命令。

sinfo

输出会显示可用的分区名称，例如本例中的 a3u：

PARTITION AVAIL TIMELIMIT NODES 状态 NODELIST
a3u* 向上 infinite 2 idle~ alice-a3u-[2-3]
a3u* 向上 infinite 2 空闲 alice-a3u-[0-1]

词元化器下载错误

当脚本尝试下载 GPT2 词元化器时，您可能会遇到与跨设备链接相关的 OSError：

OSError: [Errno 18] Invalid cross-device link: 'gpt2-vocab.json' -> '/root/.cache/torch/megatron/megatron-gpt-345m_vocab'

解决方案：

您可以通过以下两种方法解决此问题：

  • 方法 1：重新运行作业。此错误通常是暂时性的。使用相同的 --cache-dir 重新运行作业可能会解决此问题。
  • 方法 2：手动下载词元化器文件。如果重新运行作业失败，请按以下步骤操作：
    • 下载以下两个文件：
      • gpt2-vocab.json
      • gpt2-merges.txt
    • 将下载的文件移至缓存目录（例如 <var>YOUR_CACHE_DIR</var>/torch/megatron/）中的 torch/megatron/ 子目录。
    • 按如下方式重命名文件：
      • 将 gpt2-vocab.json 重命名为 megatron-gpt-345m_vocab。
      • 将 gpt2-merges.txt 重命名为 megatron-gpt-345m_merges。

NVIDIA NeMo-RL

NVIDIA NeMo-RL 框架旨在使大语言模型与人类偏好和指令保持一致。本部分将引导您在集群上使用 NeMo-RL 执行高级对齐任务，包括监督式微调 (SFT)、偏好调优（例如直接偏好优化 [DPO]）和强化学习 (RL)。

本指南涵盖两种主要工作流：运行标准批量训练作业和使用交互式开发环境进行调试。

前提条件

在开始之前，请按照创建集群页面上的说明创建一个集群，或者使用现有的受管训练集群（如果有）。

连接到集群登录节点

如需连接到集群的登录节点，请在 Google Google Cloud 控制台中前往 Google Compute Engine 虚拟机页面，然后依次点击“SSH”>“查看 Google Cloud CLI 命令”，找到正确的 Google Cloud CLI 命令。类似于以下内容：

ssh $USER_NAME@machine-addr

示例：

ssh $USER_NAME@nic0.sliua3m1-login-001.europe-north1-c.c.infinipod-shared-dev.internal.gcpnode.com

使用预构建的 Docker 映像

已转换的 .sqsh 文件适用于预构建的容器映像。您可以为自己的区域选择一个容器，然后将其直接设置为容器映像参数，也可以将其下载到集群的文件系统中。

如需直接将其设置为容器映像参数，请使用以下路径之一。请注意，您应将 <region> 替换为您的特定区域（例如 europe、asia、us）：

/gcs/vmds-containers-<region>/nemo_rl_squashfs/nemo_rl-h20250923.sqsh

如需将映像下载到集群的 Lustre 存储空间，请使用以下命令：

gs://vmds-containers-<region>/nemo_rl_squashfs/nemo_rl-h20250923.sqsh DESTINATION

下载代码

如需通过 Git CLI 获取训练 recipe，请访问 https://www.googlesource.com/new-password。您可以使用以下命令下载 recipe：

cd $HOME
git clone https://vertex-model-garden.googlesource.com/vertex-oss-training

启动作业

第 1 步：设置环境变量。

如需从 Hugging Face 拉取模型和数据，可能需要设置 HF_TOKEN。如需使用 Weights & Biases 进行实验分析，需要设置 WANDB_API_KEY。更新以下文件中的这些变量：

要更新的文件：$HOME/vertex-oss-training/nemo_rl/configs/auth.sh

如果您不想使用 Weights & Biases，请在启动脚本中将 logger.wandb_enabled 设置为 False。

第 2 步：下载或复制容器文件到启动文件夹。

以下是一些示例。

gcloud storage cp \
  gs://vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-20250923.sqsh \
  $HOME/vertex-oss-training/nemo_rl/nemo_rl-h20250923.sqsh


# OR

/gcs/vmds-containers-<region>/vmds_nemo_rl_squashfs/nemo_rl-h20250923.sqsh \
  $HOME/vertex-oss-training/nemo_rl/nemo_rl-h20250923.sqsh

cd $HOME/vertex-oss-training/nemo_rl/

# Where region is either `us`, `asia`, or `europe`

第 3 步：准备或克隆 NeMo-RL 代码库。

如果尚未克隆 NeMo-RL 代码，请克隆一份。请注意，如果您之前已在未添加 --recursive 标志的情况下克隆过代码库，则可能需要使用 git submodule update --init --recursive。

git clone https://github.com/NVIDIA-NeMo/RL --recursive

第 4 步：启动训练作业。

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --job_script algorithms/dpo.sh

其中：

  • --cluster-type 根据集群类型设置：
    • A3-Mega：hcc-a3m
    • A3-Ultra：hcc-a3u
    • A4：hcc-a4
    • A3H：hcc-a3h
  • 应相应地设置 --partition，其中 sinfo 可用于检查 Slurm 分区。

作业开始后，系统会在当前位置创建一个以其 SLURM 作业 ID 命名的新目录。您可以在其中找到属于相应作业的所有日志和检查点。更确切地说，其中包含以下目录和文件：

  • checkpoints/ → 此目录已装载到 NeMo-RL 容器中，包含训练中的所有检查点。
  • ray-logs/ → 此目录包含来自 Ray 头节点和 Ray 工作器的日志。
  • nemo_rl_output.log → 此文件包含您提交的作业的 Slurm 日志。
  • attach.sh（仅限交互式作业）→ 这是一个 Bash 脚本，可让您附加到交互式作业。如果作业成功启动，则可能需要几分钟时间才能创建此文件。

使用 NeMo-RL 进行开发

互动式设置

您可以通过以下两种方式使用 NeMo-RL 快速进行交互式开发。

nemorlinteractive

这是一个简单的辅助命令，可让您从集群中选择一个 GPU 节点（假设为节点 5），然后将您带到所选节点内正在运行的 NeMo-RL 容器。此命令有助于单节点工作流。

如需使用 nemorlinteractive，您需要按照以下前提步骤操作：

  1. 在 configs/auth.sh 文件中提供您希望加载到作业中的所有身份验证令牌（例如 HF 和 WandB）。

  2. 根据以下准则设置 CLUSTER_TYPE 环境变量：

    export CLUSTER_TYPE="hcc-a3m" # --> if you have A3-Mega cluster
export CLUSTER_TYPE="hcc-a3u" # --> if you have A3-Ultra cluster
export CLUSTER_TYPE="hcc-a4"  # --> If you have A4 cluster
export CLUSTER_TYPE="hcc-a3h" # --> If you have A3H cluster

  3. 通过获取 bash_utils.sh，在 bash 终端中导入 nemorlinteractive：

    source bash_utils.sh

  4. 运行 nemorlinteractive 命令。例如：

    # Assuming you want to take the compute node number 5.
nemorlinteractive 5

互动式启动

此选项可让您在多个计算节点上以交互方式运行工作负载。交互式作业最适合用于调试和验证用例。这些工作负载会无限期预留节点，直到开发者决定调试已结束并希望释放资源。

此选项需要遵循以下步骤：

在 configs/auth.sh 文件中提供您希望加载到作业中的所有身份验证令牌（例如 HF 和 WandB）。

sbatch -N <num_nodes> launch.sh --cluster_type hcc-a3m --interactive

  • 等待 2-5 分钟，您应该会看到 <job_id>/attach.sh 已创建。

  • 如需监控启动检查的进度，请检查 <job_id>/nemo_rl_output.log 以查看启动脚本的进度，并检查 <job_id>/ray_logs/ 以查看 Ray 头部和工作节点的启动进度。

  • 连接到交互式作业。借助此脚本，即使连接中断，您也可以重新连接：

bash <job_id>/attach.sh

后续步骤

运行预构建的工作负载可验证集群的运行状态。下一步是运行您自己的自定义训练应用。