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

このドキュメントでは、バッチ ワークロードは、完了まで実行され、Pathways クラスタと同じ GKE クラスタ内、具体的には Pathways コントローラ コンポーネント（IFRT プロキシ サーバーと Pathways リソース マネージャー）とともにデプロイされる JAX ワークロードとして定義されます。JAX ワークロードが完了すると、Pathways クラスタ コンポーネントが終了します。このガイドでは、JAX トレーニング ワークロードを使用してこれを示します。

始める前に

インストールに必要なもの:

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

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

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

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

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

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

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=ZONE \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m MaxText.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"

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

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

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

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

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --main-only=false

completed step: 1, seconds: 0.484, TFLOP/s/device: 87.349, Tokens/s/device: 2117.382, total_weights: 2945, loss: 10.888
completed step: 2, seconds: 0.407, TFLOP/s/device: 103.699, Tokens/s/device: 2513.735, total_weights: 3253, loss: 9.697
completed step: 3, seconds: 0.248, TFLOP/s/device: 170.300, Tokens/s/device: 4128.167, total_weights: 3154, loss: 9.641
completed step: 4, seconds: 0.216, TFLOP/s/device: 195.122, Tokens/s/device: 4729.880, total_weights: 3119, loss: 9.547
completed step: 5, seconds: 0.272, TFLOP/s/device: 155.298, Tokens/s/device: 3764.512, total_weights: 2837, loss: 10.179
completed step: 6, seconds: 0.472, TFLOP/s/device: 89.489, Tokens/s/device: 2169.266, total_weights: 3069, loss: 9.776

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

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

次のステップ