Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

GKE と Managed Lustre を使用して強化学習をスケーリングする

1. はじめに

ステップバイステップのチュートリアルを使用せずにパッケージ化されたスクリプトを直接実行する場合は、GoogleCloudPlatform/devrel-demos リポジトリでスクリプトを確認してください。

この Codelab では、Google Kubernetes Engine（GKE）と Managed Lustre を使用して、強化学習（RL）用の高性能トレーニングパイプラインをデプロイする方法について学習します。

強化学習ワークロード（特に Group Relative Policy Optimization（GRPO）などのアルゴリズムを使用するワークロード）は、「エクスペリエンス生成」中に大量のデータを生成し、頻繁なチェックポイント作成を必要とします。標準オブジェクトストレージでは、このような I/O バースト中にボトルネックが発生し、高価なアクセラレータがアイドル状態になる可能性があります。

並列ファイルシステムである Managed Lustre を使用して、これらのボトルネックを解消し、トレーニングのスループットを向上させます。

演習内容

GPU ベースの Ray クラスタの環境変数を構成します。
Cluster Toolkit を使用して、GKE に スポット GPU クラスタと Managed Lustre インスタンスをプロビジョニングします。
KubeRay クラスタをデプロイし、Lustre ファイルシステムをマウントします。
NeMo-RL トレーニングワークロードを送信します。
Cloud Monitoring を使用して、高スループットと低チェックポイントレイテンシをモニタリングします。

GKE、KubeRay、Managed Lustre のアーキテクチャ図

必要なもの

ウェブブラウザ（Chrome など）。
課金を有効にした Google Cloud プロジェクト

この Codelab は、GKE とストレージのコンセプトに精通している高度な技術ユーザー、プラットフォームエンジニア、AI 研究者を対象としています。

推定合計所要時間: 45 ～ 60 分 + 2 時間のトレーニング時間

2. 始める前に

Google Cloud プロジェクトの作成

Google Cloud コンソールで、Google Cloud プロジェクトを選択または作成します。
Cloud プロジェクトで課金が有効になっていることを確認します。

Cloud Shell の起動

Cloud Shell は、必要なツールがプリロードされた Google Cloud で動作するコマンドライン環境です。

Google Cloud コンソールの上部にある「Cloud Shell をアクティブにする」アイコンをクリックします。
Cloud Shell に接続したら、認証を確認します。
```
gcloud auth list
```
プロジェクトが構成されていることを確認します。
```
gcloud config get project
```
プロジェクトが想定どおりに設定されていない場合は、設定します。
```
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

<YOUR_PROJECT_ID> と <YOUR_HF_TOKEN> は実際の値に置き換えます。

ファイルを読み込んで、現在のセッションに変数を読み込みます。

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 クラスタを作成するをご覧ください）。

Managed Lustre を有効にする: gke-a4.yaml で managed-lustre と lustre-pvc のモジュールセクションのコメントを解除します。
RayOperator アドオンを有効にする: gke-a4.yaml の gke_cluster モジュール設定で enable_ray_operator: true を設定します。

インフラストラクチャをデプロイする

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

デプロイが完了するまで待ちます。Cluster Toolkit は、1 つの調整されたデプロイで 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 という名前のファイルを作成します。これは、nvidia-b200 アクセラレータタイプを使用して、/lustre に Lustre ボリュームをマウントする KubeRay ヘッドノードとワーカーノードを指定します。

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. 強化学習ワークロードを送信する

このステップでは、NeMo-RL GRPO トレーニングジョブを Ray クラスタに送信します。

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

ランタイム環境の構成を作成する

環境変数を Ray ジョブに渡すための JSON ファイルを作成します。

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 ワーカーが長時間ブロックされることはありません。

チェックポイントの速度を確認するには、保存されたチェックポイントを示すログ行を探します。

Cloud コンソールで Lustre 指標を表示する

Lustre インスタンスの指標を表示するには:

Google Cloud コンソールで、Managed Service for Lustre を検索します。
インスタンス名（${CLUSTER_NAME}-lustre や rl-demo-gpu-lustre など）をクリックします。
[モニタリング] タブをクリックします。

ここでは、次のことを確認できます。

スループット（バイト/秒）: チェックポイント作成中にスパイクが発生していることを確認します。
容量: チェックポイントで使用されている容量をモニタリングします。

Lustre のパフォーマンスグラフ Lustre は非常に高速な書き込みが可能で、チェックポイントを最小限の時間で書き込むことができます

8. リソースをクリーンアップする

Cloud Shell で次のコマンドを実行して、プロビジョニングされたすべてのインフラストラクチャ（GKE クラスタ、GPU ノードプール、Managed Lustre インスタンス、VPC ネットワーク）を 1 つのステップで破棄します。

gcluster destroy "${CLUSTER_NAME}"

このコマンドはフォアグラウンドで同期的に実行され、デプロイによって管理されるすべてのインフラストラクチャを破棄し、進行状況ログをターミナルに出力します。コマンドが完全に終了するまで待ってから、Cloud Shell セッションを閉じます。

9. 完了

GKE と Managed Lustre を使用して強化学習をスケーリングする Codelab を完了しました。

学習した内容

Cluster Toolkit を使用して、スポットインスタンスと Managed Lustre ストレージで GKE GPU クラスタをプロビジョニングする方法。
KubeRay クラスタをデプロイして Lustre ストレージをマウントする方法。
NeMo-RL GRPO トレーニングワークロードを送信する方法。
トレーニング中にストレージパフォーマンスをモニタリングする方法。

次のステップ

NVIDIA NeMo-RL のその他の機能を確認する。
Google Cloud AI Hypercomputer の詳細を確認する。
Managed Service for Lustre のドキュメントを確認する。