このページでは、プラットフォーム管理者と ML エンジニア向けに、Slurm でオーケストレートされた Cluster Director クラスタを作成し、Pathways ワークロード用に構成する方法について説明します。これらの手順を完了すると、TPU v6e スライス、CPU コントローラ ノード、共有 Google Cloud Managed Lustre ファイル システム、/home/pathways ランタイム環境を含むクラスタがプロビジョニングされます。Pathways は Cluster Director クラスタにプリインストールされていません。クラスタを作成した後、Pathways スクリプトをログインノードにダウンロードし、設定スクリプトを 1 回実行します。
始める前に
-
このページのサンプルをどのように使うかに応じて、タブを選択してください。
gcloud
Google Cloud コンソールで Cloud Shell をアクティブにします。
Google Cloud コンソールの下部で Cloud Shell セッションが開始し、コマンドライン プロンプトが表示されます。Cloud Shell はシェル環境です。Google Cloud CLI がすでにインストールされており、現在のプロジェクトの値もすでに設定されています。セッションが初期化されるまで数秒かかることがあります。
REST
このページの REST API サンプルをローカル開発環境で使用するには、gcloud CLI に指定した認証情報を使用します。
Google Cloud CLI をインストールします。
外部 ID プロバイダ（IdP）を使用している場合は、まず連携 ID を使用して gcloud CLI にログインする必要があります。
詳細については、 Google Cloud 認証ドキュメントの REST を使用して認証するをご覧ください。
始める前に
始める前に、次の前提条件を満たす必要があります。
- Google Cloud CLI をインストールします。
- Hypercompute Cluster API を有効にします。
- ユーザー アカウントと Compute Engine のデフォルトのサービス アカウントに、Cluster Director クラスタの作成と接続に必要な Identity and Access Management（IAM）ロールがあることを確認します。ロールのリストについては、Cluster Director のクイックスタートの始める前にをご覧ください。
- クラスタを作成するゾーンに TPU v6e 予約を作成します（たとえば、2 つの
2x4スライスに少なくとも 4 つのct6e-standard-4tVM）。 - スライスの自動復元のために、クラスタノードでパスワードなしの
sudoを構成します。プロジェクトに対する Compute OS 管理者ログイン（roles/compute.osAdminLogin）ロールを付与するよう管理者に依頼してください。 - Pathways のスクラッチ データとワークロード出力用の Cloud Storage バケットを作成し、プロジェクトの Storage オブジェクト ユーザー（
roles/storage.objectUser）ロールをユーザー アカウントと Compute Engine のデフォルト サービス アカウントに付与します。
ローカル環境を設定する
Google Cloud の認証情報でログインします。
gcloud auth login
ワークロードに適した値を使用して、次の環境変数を定義します。
export PROJECT=PROJECT_ID export REGION=REGION export ZONE=ZONE export CLUSTER_ID=CLUSTER_ID export RESERVATION_NAME=RESERVATION_NAME
次のように置き換えます。
PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。REGION: リソースを作成するリージョン。ZONE: リソースを作成するゾーン。CLUSTER_ID: クラスタ名。名前は最大 10 文字で、数字と英小文字のみを使用できます。RESERVATION_NAME:ZONEの TPU v6e 予約の名前。
Cluster Director クラスタを作成する
次の例では、予約から 2 つの TPU v6e 2x4 スライス（4 つの ct6e-standard-4t VM）と、Pathways コントローラ用の 1 つの CPU ノードを持つクラスタを作成します。すべてのノードが、共有 Managed Lustre ファイル システムを /home にマウントします。クラスタ定義のフィールドの詳細については、カスタム クラスタを作成するをご覧ください。
REST API を使用してクラスタを作成します。
リクエストの本文を作成します。
cat <<EOF > "${CLUSTER_ID}_create.json" { "networkResources": { "net": { "config": { "newNetwork": { "network": "projects/${PROJECT}/global/networks/${CLUSTER_ID}-net" } } } }, "storageResources": { "home": { "config": { "newLustre": { "lustre": "projects/${PROJECT}/locations/${ZONE}/instances/${CLUSTER_ID}-lustre", "filesystem": "lustre", "capacityGb": "18000", "perUnitStorageThroughput": "500" } } } }, "computeResources": { "crtpu": { "config": { "newReservedInstances": { "reservation": "projects/${PROJECT}/zones/${ZONE}/reservations/${RESERVATION_NAME}" } } }, "crcpu": { "config": { "newOnDemandInstances": { "zone": "${ZONE}", "machineType": "n2-highmem-32" } } } }, "orchestrator": { "slurm": { "loginNodes": { "count": "1", "zone": "${ZONE}", "machineType": "n2-standard-16", "enableOsLogin": true, "enablePublicIps": true, "bootDisk": { "type": "pd-balanced", "sizeGb": "100" }, "storageConfigs": [ { "id": "home", "localMount": "/home" } ] }, "nodeSets": [ { "id": "nstpu", "computeId": "crtpu", "staticNodeCount": "4", "maxDynamicNodeCount": "0", "enablePublicIps": true, "resourcePolicyConfig": { "acceleratorTopology": "2x4" }, "computeInstance": { "bootDisk": { "type": "hyperdisk-balanced", "sizeGb": "100" } }, "storageConfigs": [ { "id": "home", "localMount": "/home" } ] }, { "id": "nscpu", "computeId": "crcpu", "staticNodeCount": "1", "maxDynamicNodeCount": "0", "enablePublicIps": true, "computeInstance": { "bootDisk": { "type": "pd-balanced", "sizeGb": "100", "image": "projects/clusterdirector-public-images/global/images/family/tpu-v6e-v5p-ubuntu2204-slurm2605-v20260902" } }, "storageConfigs": [ { "id": "home", "localMount": "/home" } ] } ], "partitions": [ { "id": "tpu", "nodeSetIds": ["nstpu"] }, { "id": "cpu", "nodeSetIds": ["nscpu"] } ], "defaultPartition": "tpu", "config": { "requeueExitCodes": [1, 2, 134, 137, 139, 143] } } } } EOFリクエストを送信します。
curl -s -X POST \ -H "Content-Type: application/json" \ -H "Authorization: Bearer $(gcloud auth print-access-token)" \ -d @"${CLUSTER_ID}_create.json" \ "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters?clusterId=${CLUSTER_ID}"クラスタには次のものが含まれます。
tpuパーティション: 予約済みの TPU v6e VM の静的ノードセット。staticNodeCountは VM の数、acceleratorTopologyは各スライスのトポロジです。各ct6e-standard-4tVM には 4 つのチップがあるため、2x4スライスには 2 つの VM があります。スライスを増やす場合は、乗算します。2 つの2x4スライスには、4のstaticNodeCountが必要です。cpuパーティション: Pathways リソース マネージャー、IFRT プロキシ サーバー、MaxText クライアントを実行する 1 つのオンデマンド VM。Pathways スクリプトは、cpuという名前のパーティションにジョブを送信します。VM はn2-highmem-32マシンタイプを使用します。これは、Pathways サービスと MaxText がチェックポイントを保存する際に、最大 240 GB のメモリが必要になるためです。また、Pathways サービスがこの VM から TPU ワーカー ジョブを送信するため、TPU スライスをサポートする Slurm 26.05 のイメージも使用します。requeueExitCodes: Slurm は、これらのコードのいずれかで終了するジョブを再キューイングします。この設定により、スライス内の最初の VM だけでなく、任意の VM がプリエンプトされたときに、Slurm が TPU ワーカー ジョブを再起動できます。
クラスタの準備が完了するまで待ちます。 Google Cloud コンソールの [Cluster Director] ページで進行状況を確認するか、次のコマンドの出力で
reconcilingがfalseであることを確認します。curl -s \ -H "Authorization: Bearer $(gcloud auth print-access-token)" \ "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"
Pathways 環境を設定する
Pathways 環境は共有ファイル システムの /home/pathways に存在するため、クラスタごとに 1 回設定します。これには、Pathways コンテナ イメージ、JAX と pathwaysutils を含む Python 環境、Pathways スクリプトが含まれています。
ログインノードに接続します。他の接続方法については、クラスタに接続するをご覧ください。
gcloud compute ssh ${CLUSTER_ID}-login-001 \ --project=${PROJECT} \ --zone=${ZONE} \ --tunnel-through-iapcpuパーティションとtpuパーティションが存在し、idle状態であることを確認します。sinfo -o "%P %a %D %T %N %f"出力は次のようになります。
PARTITION AVAIL NODES STATE NODELIST FEATURES tpu* up 4 idle ${CLUSTER_ID}-nstpu-[0-3] cpu up 1 idle ${CLUSTER_ID}-nscpu-0Pathways スクリプトをログインノードにダウンロードします。
mkdir -p ~/pathways-scripts && cd ~/pathways-scripts for f in setup_env.sh start_pathways.sh run_worker_slice.sh run_maxtext.sh; do curl -fsSLO "https://raw.githubusercontent.com/GoogleCloudPlatform/slurm-gcp/tpu-2605/pathways/scripts/${f}" doneこのコマンドは、次のスクリプトをダウンロードします。
setup_env.sh: Pathways 環境の 1 回限りの設定を行います。start_pathways.sh: Pathways サービスを開始します。run_worker_slice.sh: 1 つの TPU スライスの Pathways ワーカーを起動します（start_pathways.shによって送信）。run_maxtext.sh: MaxText トレーニング ワークロードの例を実行します。
設定スクリプトを実行します。これには 5 分ほどかかります。
chmod +x *.sh ./setup_env.shこのスクリプトは、Pathways サーバーとプロキシ イメージを pull し、JAX、
pathwaysutils、MaxText を使用して Python 環境をビルドし、スクリプトを/home/pathwaysにコピーして、~/pathwaysを/home/pathwaysにリンクします。次のような出力で終了します。Environment ready: JAX 0.11.1省略可: 他のクラスタ ユーザーに Pathways 環境へのアクセス権を付与するには、
/home/pathwaysをホーム ディレクトリにリンクしてもらいます。ln -sfn /home/pathways ~/pathways
この時点で、クラスタは Pathways の準備ができています。バッチ ワークロードとインタラクティブ ワークロードは異なる設定で開始されるため、ワークロードの実行の一部として Pathways サービスを開始します。詳細については、Pathways を使用してバッチ ワークロードを実行すると Pathways を使用してインタラクティブ ワークロードを実行するをご覧ください。
問題が発生した場合は、Cloud での Pathways のトラブルシューティングをご覧ください。
クリーンアップ
このページで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、完了時に Cluster Director クラスタを削除します。
curl -s -X DELETE \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
"https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"
次のステップ
- Pathways を使用してバッチ ワークロードを実行する
- Pathways を使用してインタラクティブなワークロードを実行する
- Pathways を使用した復元力のあるトレーニング
- JAX ワークロードを Pathways に移植する
- Cloud での Pathways のトラブルシューティング