使用 Pathways 執行批次工作負載

就本文而言，批次工作負載是指執行完畢的 JAX 工作負載，且部署在與 Pathways 叢集相同的 GKE 叢集內，具體來說，是與 Pathways 控制器元件 (IFRT Proxy 伺服器和 Pathways 資源管理工具) 部署在同一位置。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。您可以使用這個 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 bucket
  • 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

後續步驟