使用 Pathways 執行批次工作負載

本指南使用 JAX 訓練工作負載，在與 Pathways 控制器元件相同的叢集中執行至完成，示範如何執行批次工作負載。

您可以在 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。您可以使用這個 Docker 映像檔，透過 Pathways 後端在 TPU 上執行訓練。

Cluster Director

您不需要建構訓練映像檔：setup_env.sh 已在 /home/pathways 的 Pathways 環境中安裝 MaxText，並由 run_maxtext.sh 執行。

執行批次工作負載

GKE

  1. 如果 PATH 環境變數中沒有 gcluster 二進位檔，請將包含 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 bucket。
    • 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 bucket。

    如果 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 訓練工作負載。您可以在服務工作完成後立即提交，系統最多會等待一小時，直到服務準備就緒。

    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：訓練 8 層變體，而非完整的 32 層，適用於 2x2 個切片。

    追蹤工作負載的進度 (將 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 worker 工作。TPU VM 是靜態節點，因此會持續運作，並可供其他 Slurm 工作使用。

如果遇到問題，請參閱「排解 Cloud 上的 Pathways 問題」。

後續步驟