事前構築済みのワークロードを実行する

Vertex AI マネージド トレーニングに関心をお持ちの場合は、営業担当者にお問い合わせください。

このガイドでは、マネージド トレーニング クラスタで NVIDIA NeMo エコシステムを使用して、エンドツーエンドの生成 AI モデル開発を行う方法について説明します。このガイドでは、次の個別の関連ワークフローについて、それぞれ専用のセクションで手順を説明します。

  • NVIDIA NeMo: 基盤モデルの開発では、次の手順に沿って大規模な事前トレーニング、継続的事前トレーニング(CPT)、教師ありファインチューニング(SFT)を行います。
  • NVIDIA NeMo-RL: モデルの調整と好みのチューニングには、このセクションを使用して、強化学習(RL)などの高度な手法を適用し、モデルを人間の指示や好みに合わせます。

このドキュメントでは、モデルを一から構築する場合や、既存のモデルを改良する場合でも、環境の設定、コンテナ化されたジョブの管理、クラスタでのトレーニング スクリプトの起動について説明します。

NVIDIA NeMo

NVIDIA NeMo フレームワークは、生成 AI モデルの構築、カスタマイズ、デプロイを行うためのエンドツーエンドのプラットフォームです。このガイドのこのセクションは、モデル開発の基礎段階に重点を置くデベロッパーと研究者を対象としています。このガイドでは、NeMo を使用して Managed Training クラスタで大規模な事前トレーニング、継続的な事前トレーニング(CPT)、教師ありファインチューニング(SFT)を行うための手順を説明します。

このトレーニング クラスタガイドでは、NeMo フレームワークでトレーニング ジョブを実行するための完全なワークフローについて説明します。このプロセスは、環境の初回設定とジョブの起動の繰り返し手順という 2 つの主要な部分に分かれています。

環境を設定する

ジョブを起動する前に、コンテナ イメージと必要なトレーニング スクリプトを用意して、環境を準備する必要があります。

コンテナ イメージを準備する

コンテナ イメージには、ビルド済みイメージを使用する(推奨)か、カスタム イメージをビルドするかの 2 つのオプションがあります。

ビルド済みのイメージを使用する(推奨)

ビルド済みコンテナ イメージは .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 に指定したパスで最終的なコンテナ イメージを確認します。

トレーニング レシピをダウンロードする

トレーニング レシピは非公開の googlesource.com リポジトリに保存されます。Git コマンドラインでアクセスするには、まず認証情報を作成する必要があります。

  1. 認証情報を生成します

    次の URL にアクセスし、画面上の手順に沿って操作します。これにより、リポジトリで認証を行うようにローカル環境が構成されます。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>/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-a3mhcc-a3uhcc-a4 など)に基づいて設定されます。
  • --partition は利用可能なパーティションに設定する必要があります。パーティション名は sinfo コマンドで確認できます。
  • run.py スクリプトは、--log-dir--cache-dir--data-dir などの複数のディレクトリが設定されている場合、それらを 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 STATE NODELIST
a3u* アップ 無限 2 idle~ alice-a3u-[2-3]
a3u* アップ 無限 2 idle alice-a3u-[0-1]

トークナイザーのダウンロード エラー

スクリプトが GPT2 トークナイザーをダウンロードしようとすると、クロスデバイス リンクに関連する OSError が発生することがあります。

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

ソリューション:

この問題を解決するには、次の 2 つのオプションがあります。

  • オプション 1: ジョブを再実行します。このエラーは一時的なものであることがよくあります。同じ --cache-dir を使用してジョブを再実行すると、問題が解決する場合があります。
  • オプション 2: トークナイザー ファイルを手動でダウンロードします。ジョブの再実行が失敗した場合は、次の手順を行います。
    • 次の 2 つのファイルをダウンロードします。
      • gpt2-vocab.json
      • gpt2-merges.txt
    • ダウンロードしたファイルをキャッシュ ディレクトリ内の torch/megatron/ サブディレクトリ(<var>YOUR_CACHE_DIR</var>/torch/megatron/ など)に移動します。
    • ファイルの名前を次のように変更します。
      • gpt2-vocab.jsonmegatron-gpt-345m_vocab に名前を変更する。
      • gpt2-merges.txtmegatron-gpt-345m_merges に名前を変更する。

NVIDIA NeMo-RL

NVIDIA NeMo-RL フレームワークは、大規模言語モデルを人間の好みや指示に合わせるように設計されています。このセクションでは、クラスタで NeMo-RL を使用して、教師ありファインチューニング(SFT)、プリファレンス チューニング(Direct Preference Optimization(DPO)など)、強化学習(RL)などの高度なアライメント タスクを実行する方法について説明します。

このガイドでは、標準のバッチ トレーニング ジョブの実行と、デバッグ用のインタラクティブ開発環境の使用という 2 つの主なワークフローについて説明します。

前提条件

始める前に、クラスタを作成するページの手順に沿ってクラスタを作成するか、既存の Managed Training クラスタがある場合はそれを使用します。

クラスタのログインノードに接続する

クラスタのログインノードに接続するには、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> は、特定のリージョン(europeasiaus など)に置き換えてください。

/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 でトレーニング レシピにアクセスするには、https://www.googlesource.com/new-password にアクセスしてください。レシピは次のコマンドでダウンロードできます。

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

ジョブの起動

ステップ 1: 環境変数を設定します。

Hugging Face からモデルとデータを pull するには、HF_TOKEN の設定が必要になることがあります。テスト分析に Weights & Biases を使用するには、WANDB_API_KEY を設定する必要があります。次のファイルでこれらの変数を更新します。

更新するファイル: $HOME/vertex-oss-training/nemo_rl/configs/auth.sh

Weights & Biases を使用しない場合は、起動スクリプトで logger.wandb_enabledFalse に設定します。

ステップ 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 を使用した迅速なインタラクティブ開発には、次の 2 つのオプションがあります。

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

次のステップ

事前構築済みのワークロードを実行すると、クラスタの運用ステータスが検証されます。次のステップは、独自のカスタム トレーニング アプリケーションを実行することです。