Pathways でバッチ ワークロードを実行する

このガイドでは、Pathways コントローラ コンポーネントと同じクラスタで完了まで実行される JAX トレーニング ワークロードを使用して、バッチ ワークロードを実行する方法について説明します。

バッチ ワークロードは、Google Kubernetes Engine（GKE）クラスタまたは Slurm を使用する Cluster Director クラスタで実行できます。

  • GKE では、JAX ワークロードが完了すると、Pathways クラスタ コンポーネントが終了します。
  • Cluster Director では、Pathways コントローラ コンポーネントは個別の Slurm ジョブとして実行されるため、JAX ワークロードが完了した後も、停止するまで実行が継続されます。

このページでは、両方のプラットフォームの手順について説明します。各セクションで、[GKE] タブまたは [Cluster Director] タブを選択します。

始める前に

始める前に

GKE

次の前提条件を満たしていることを確認してください。

Cluster Director

次の前提条件を満たしていることを確認してください。

MaxText を使用してトレーニング イメージをビルドする

MaxText は、Google が開発したオープンソースの大規模言語モデル（LLM）プロジェクトです。JAX で記述されており、 Google Cloud TPU と GPU で効率的に実行できるように設計されています。

GKE

オープンソースの GitHub リポジトリから最新バージョンの安定版 JAX を使用して MaxText Docker イメージをビルドするには、次のコマンドを実行します。

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

次のように置き換えます。

  • PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
  • USER: 実際のGoogle Cloud ユーザー ID。

このコマンドは、MaxText Kubernetes イメージを gcr.io/PROJECT_ID/USER_runner に push します。この Docker イメージを使用すると、Pathways バックエンドを使用して TPU でトレーニングを実行できます。

Cluster Director

トレーニング イメージをビルドする必要はありません。setup_env.sh は /home/pathways の Pathways 環境に MaxText をインストールし、run_maxtext.sh がそれを実行します。

バッチ ワークロードを実行する

GKE

  1. gcluster バイナリが PATH 環境変数にない場合は、gcluster を含むディレクトリを PATH に追加します。

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    CLUSTER_TOOLKIT_PATH は、gcluster バイナリを含むディレクトリの絶対パス（$HOME/cluster-toolkit など）に置き換えます。

  2. gcluster job submit コマンドを使用して、ビルド済みの MaxText Docker イメージを送信します。

    gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.train /deps/src/maxtext/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

    次のように置き換えます。

    • BUCKET_NAME: 出力用の Cloud Storage バケット。
    • WORKLOAD: ワークロードを識別する一意の名前（DNS ラベルの制限により 22 文字以下）。
    • CLUSTER: GKE クラスタの名前。
    • PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
    • LOCATION: GKE クラスタのリージョンまたはゾーン。
    • WORKLOAD_NODEPOOL_COUNT: TPU スライス ノードプールの数。
    • COMPUTE_TYPE: TPU マシンタイプ（例: ct6e-standard-4t）。TPU の各バージョンでサポートされている TPU タイプの詳細については、TPU のバージョンをご覧ください。
    • TOPOLOGY: TPU の配置トポロジ（例: 2x4）。
    • USER: 実際のGoogle Cloud ユーザー ID。
    • RUN_NAME: ワークフロー実行を識別するユーザー割り当ての名前。

    ジョブ送信オプションの詳細については、Cluster Toolkit ジョブ送信ガイドをご覧ください。

  3. gcluster job logs コマンドを使用して、ワークロードの進行状況を確認します。

    gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

    出力は次のようになります。

    completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

  4. 省略可: ワークロードが完了する前にキャンセルするには、gcluster job cancel コマンドを使用します。

    gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION

Cluster Director

~/pathways からクラスタのログインノードで次のコマンドを実行します。スクリプトは互いを検索し、そのディレクトリにログを書き込みます。

  1. Pathways サービスを開始します。

    cd ~/pathways
SLICE_SHAPE=TOPOLOGY \
    MIN_SLICES=MIN_SLICES \
    MAX_SLICES=MAX_SLICES \
    GCS_BUCKET=gs://BUCKET_NAME \
    sbatch start_pathways.sh

    次のように置き換えます。

    • TOPOLOGY: 各スライスの TPU トポロジ。TPU ノードセットの acceleratorTopology と一致する必要があります（例: 2x4）。
    • MIN_SLICES: トレーニングを開始して続行するために必要な、実行中のスライスの最小数。
    • MAX_SLICES: 使用するスライスの最大数。通常はノードセット内のスライスの数。
    • BUCKET_NAME: Pathways のスクラッチ データとチェックポイント用の Cloud Storage バケット。

    TPU パーティションがクラスタのデフォルト パーティションでない場合は、PARTITION をその名前に設定します。

    ログ pathways-service-JOB_ID.out（JOB_ID は Slurm ジョブ ID）に次のような出力が表示されると、サービスの準備が完了します。

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
 Running TPU Slices : 2 / 2 (min=2)

  2. MaxText トレーニング ワークロードを送信します。サービスジョブの直後に送信できます。サービスが準備完了になるまで最大 1 時間待機します。

    sbatch run_maxtext.sh

    MaxText ジョブは、Pathways サービスに少なくとも MIN_SLICES 個の実行中のスライスが存在するまで待機してから、トレーニングを開始します。ワークロードは、合成データで Llama 3.1 8B をトレーニングします。スライス設定とバケットはサービスから取得されるため、このジョブでは設定しないでください。次のオプション変数を設定できます。

    • STEPS: トレーニング ステップの数。デフォルトは 30 です。
    • CHECKPOINT_PERIOD: チェックポイント間のステップ数。デフォルトは 5 です。
    • BASE_OUTPUT_DIR: MaxText がチェックポイントを書き込む Cloud Storage パス。デフォルトは gs://BUCKET_NAME/maxtext-runs です。
    • REDUCED_LAYERS=1: 2x2 スライスに対して、32 レイヤーすべてではなく、8 レイヤーのバリアントをトレーニングします。

    ワークロードの進行状況を確認します（JOB_ID は、MaxText ジョブの Slurm ジョブ ID に置き換えます）。MaxText はログを標準エラーに書き込みます。

    tail -f maxtext-llama31-JOB_ID.err

    各トレーニング ステップで、次のような行が出力されます。

    completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238

    サービスを再起動せずに、後で run_maxtext.sh を再度送信できます。

  3. 完了したら、ワークロードを停止してから Pathways サービスを停止します。

    scancel -u ${USER} -n maxtext-llama31
scancel -u ${USER} -n pathways-service

    サービスをキャンセルすると、Pathways ワーカーのジョブもキャンセルされます。TPU VM は静的ノードであるため、実行を継続し、他の Slurm ジョブで使用できるようになります。

問題が発生した場合は、Cloud での Pathways のトラブルシューティングをご覧ください。

次のステップ