Pathways を使用して Cluster Director クラスタを作成する

このページでは、プラットフォーム管理者と ML エンジニア向けに、Slurm でオーケストレートされた Cluster Director クラスタを作成し、Pathways ワークロード用に構成する方法について説明します。これらの手順を完了すると、TPU v6e スライス、CPU コントローラ ノード、共有 Google Cloud Managed Lustre ファイル システム、/home/pathways ランタイム環境を含むクラスタがプロビジョニングされます。Pathways は Cluster Director クラスタにプリインストールされていません。クラスタを作成した後、Pathways スクリプトをログインノードにダウンロードし、設定スクリプトを 1 回実行します。

始める前に

始める前に

始める前に、次の前提条件を満たす必要があります。

  • Google Cloud CLI をインストールします。
  • Hypercompute Cluster API を有効にします。
  • ユーザー アカウントと Compute Engine のデフォルトのサービス アカウントに、Cluster Director クラスタの作成と接続に必要な Identity and Access Management（IAM）ロールがあることを確認します。ロールのリストについては、Cluster Director のクイックスタートの始める前にをご覧ください。
  • クラスタを作成するゾーンに TPU v6e 予約を作成します（たとえば、2 つの 2x4 スライスに少なくとも 4 つの ct6e-standard-4t VM）。
  • スライスの自動復元のために、クラスタノードでパスワードなしの 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 にマウントします。クラスタ定義のフィールドの詳細については、カスタム クラスタを作成するをご覧ください。

  1. 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-4t VM には 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 ワーカー ジョブを再起動できます。

  2. クラスタの準備が完了するまで待ちます。 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 スクリプトが含まれています。

  1. ログインノードに接続します。他の接続方法については、クラスタに接続するをご覧ください。

    gcloud compute ssh ${CLUSTER_ID}-login-001 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --tunnel-through-iap

  2. cpu パーティションと 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-0

  3. Pathways スクリプトをログインノードにダウンロードします。

    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 トレーニング ワークロードの例を実行します。

  4. 設定スクリプトを実行します。これには 5 分ほどかかります。

    chmod +x *.sh
./setup_env.sh

    このスクリプトは、Pathways サーバーとプロキシ イメージを pull し、JAX、pathwaysutils、MaxText を使用して Python 環境をビルドし、スクリプトを /home/pathways にコピーして、~/pathways を /home/pathways にリンクします。次のような出力で終了します。

    Environment ready: JAX 0.11.1

  5. 省略可: 他のクラスタ ユーザーに 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}"

次のステップ