使用 TPU7x (Ironwood) 訓練模型

本文說明如何使用 Cluster Toolkit 佈建 TPU7x 資源，並使用 MaxText (JAX) 或 TorchTitan on TPU (PyTorch) 部署訓練工作負載。

TPU7x 是 Ironwood 系列的第一個版本，也是 Google Cloud的第七代 TPU。Ironwood 世代專為大規模 AI 訓練和推論而設計。詳情請參閱「TPU7x」。

如需更多針對 TPU7x 最佳化的範例，請參閱 GitHub 上的下列存放區：

事前準備

開始之前，請確認您符合下列條件：

接著完成下列設定步驟：

  1. 設定專案和叢集的環境變數：

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"

    更改下列內容：

    • YOUR_PROJECT_ID：您的 Google Cloud 專案 ID。
    • YOUR_ZONE：要在哪個可用區建立叢集。
    • YOUR_CLUSTER_NAME：新叢集的名稱。
    • YOUR_BUCKET_NAME：用於訓練輸出的 Cloud Storage bucket 名稱。

  2. 如果沒有 Cloud Storage bucket，請建立一個：

    gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

部署 TPU7x (Ironwood) 叢集

按照「部署 GKE TPU 7x 叢集」中的操作說明，部署 Google Kubernetes Engine (GKE) TPU7x 叢集。請確保叢集節點集區拓撲符合訓練工作負載需求，例如 MaxText 範例的 4x4x8，或 TorchTitan 範例的 2x2x4。

使用 MaxText 部署訓練工作負載

以下各節說明如何使用 Cluster Toolkit 建構容器映像檔，並提交 MaxText 訓練工作。

建構及上傳 MaxText Docker 映像檔

建構及推送 MaxText Docker 映像檔：

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

定義 MaxText 訓練指令

準備指令，在 Docker 容器中執行訓練指令碼。

MaxText 1B 模型設定會訓練約 10 億個參數的語言模型。這個模型可用於測試小型晶片規模。效能未達最佳狀態。

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_training \
        steps=30"

提交 MaxText 訓練工作

使用 gcluster job submit 提交訓練工作。如要進一步瞭解工作提交選項，請參閱 Cluster Toolkit 工作指南：

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

叢集中的工作負載名稱不得重複。這個範例會將 $(date +%H%M) 附加至工作負載名稱。

使用 TorchTitan 部署訓練工作負載

TorchTitan 會使用 TorchTPU，在 Cloud TPU 執行分散式 PyTorch 訓練。詳情請參閱「TorchTitan on TPU README」。

後續步驟