使用 Pathways 執行批次工作負載

在本文件中，批次工作負載是指執行至完成的 JAX 工作負載，且部署在與 Pathways 叢集相同的 GKE 叢集中，特別是與 Pathways 控制器元件 (IFRT Proxy 伺服器和 Pathways 資源管理工具) 搭配使用。JAX 工作負載完成後，Pathways 叢集元件就會終止。本指南將以 JAX 訓練工作負載為例，示範如何執行這項操作。

事前準備

請確認您已備妥以下項目：

使用 Maxtext 建構訓練映像檔

MaxText 是 Google 開發的開放原始碼大型語言模型 (LLM) 專案。MaxText 以 JAX 編寫，具備高效能和可擴充性，可在 Google Cloud TPU 和 GPU 上高效執行。

如要使用 OSS 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

這個指令會將 MaxText Kubernetes 映像檔推送至 gcr.io/$PROJECT_ID/${USER}_runner。您可以使用這個 Docker 映像檔，透過 Pathways 後端在 TPU 上執行訓練。

使用 Cluster Toolkit 執行批次工作負載

如果 PATH 環境變數中沒有 gcluster 二進位檔，請將包含 gcluster 的目錄新增至 PATH：

export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

將 CLUSTER_TOOLKIT_PATH 替換為包含 gcluster 二進位檔的目錄絕對路徑，例如 $HOME/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=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"

如要進一步瞭解工作提交選項，請參閱 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
  • LOCATION：GKE 叢集的區域或可用區
  • USER：您的 Google Cloud 使用者 ID
  • BUCKET_NAME：輸出內容的 Cloud Storage bucket
  • RUN_NAME：使用者指派的名稱，用於識別工作流程執行作業

使用 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

如要在工作負載完成前取消，請使用 gcluster job cancel 指令：

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

後續步驟