執行預先建構的工作負載

如要使用 Vertex AI Managed Training，請與銷售代表聯絡以取得存取權。

本指南說明如何在代管訓練叢集上使用 NVIDIA NeMo 生態系統，進行端對端生成式 AI 模型開發。本指南提供下列不同但相關工作流程的逐步操作說明，每個工作流程都有專屬章節：

  • NVIDIA NeMo：如要開發基礎模型，請按照這些操作說明進行大規模預先訓練、持續預先訓練 (CPT) 和監督式微調 (SFT)。
  • NVIDIA NeMo-RL：如要進行模型校正和偏好設定微調，請使用本節套用強化學習 (RL) 等進階技術，讓模型符合人類指令和偏好設定。

無論是從頭建構模型或改良現有模型，本文都會引導您設定環境、管理容器化工作，以及在叢集上啟動訓練指令碼。

NVIDIA NeMo

NVIDIA NeMo 架構是端對端平台，可建構、自訂及部署生成式 AI 模型。本指南的這一節專為開發人員和研究人員而設計，著重於模型開發的基礎階段。這份指南提供逐步操作說明，教您如何在 Managed Training 叢集上使用 NeMo 執行大規模預先訓練、持續預先訓練 (CPT) 和監督式微調 (SFT)。

本訓練叢集指南提供完整的工作流程，說明如何使用 NeMo 架構執行訓練工作。這個程序主要分為兩部分：環境的初始一次性設定，以及啟動作業的重複步驟。

設定環境

啟動作業前，請先準備環境，確保您擁有容器映像檔和必要的訓練指令碼。

安裝 NeMo-Run

pip install git+https://github.com/NVIDIA/NeMo-Run.git

準備容器映像檔

容器映像檔有兩種選擇：使用預建映像檔 (建議) 或建構自訂映像檔。

使用預先建構的映像檔 (建議)

預先建構的容器映像檔採用 .squashfs 格式。將您所在區域適用的圖片複製到工作目錄。

# Example for the US region
gcloud storage cp gs://vmds-containers-us/vmds_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 指定的路徑中找到最終容器映像檔。

下載訓練食譜

訓練食譜會儲存在私人 googlesource.com 存放區。 如要透過 Git 指令列存取這些憑證，請先產生驗證憑證。

  1. 產生驗證憑證。

    前往下列網址，然後按照畫面上的指示操作。 這會設定本機環境，以便透過存放區進行驗證。https://www.googlesource.com/new-password

  2. 複製存放區。

    憑證通過驗證後，請執行下列指令下載食譜。

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

啟動訓練工作

設定環境後，即可啟動訓練工作。

步驟 1：設定環境變數

工作可能需要下列環境變數：

  • 如要從 Hugging Face 下載模型和資料集，必須提供 HF_TOKEN。
  • 如要使用 Weights & Biases 進行實驗分析，必須先安裝 WANDB_API_KEY。
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>/vmds_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 指令檢查分割區名稱。
  • run.py 指令碼會自動將多個目錄掛接到 Docker 容器，包括 --log-dir、--cache-dir 和 --data-dir (如有設定)。

監控工作狀態和記錄

啟動工作後，畫面會顯示狀態區塊：

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 節點 狀態 NODELIST
a3u* 向上鍵 無限 2 idle~ alice-a3u-[2-3]
a3u* 向上鍵 無限 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
    • 將下載的檔案移至快取目錄中的 torch/megatron/ 子目錄 (例如 <var>YOUR_CACHE_DIR</var>/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 Virtual Machine」頁面，依序點選「SSH」>「View Google Cloud CLI command」(查看 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>/vmds_nemo_rl_squashfs/nemo_rl-h20250923.sqsh

如要將映像檔下載至叢集的 Lustre 儲存空間，請使用下列指令：

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

下載程式碼

如要透過 git CLI 存取訓練配方，請前往 https://www.googlesource.com/new-password。 您可以使用下列指令下載方案：

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 終端機中，透過來源 bash_utils.sh 匯入 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

後續步驟

執行預先建構的工作負載，即可驗證叢集的運作狀態。 下一步是執行自訂訓練應用程式。