使用 TPU7x (Ironwood) 訓練模型

本文說明如何佈建 TPU7x 資源，並提供使用 MaxText 和 Cluster Toolkit 部署訓練工作負載的範例。

TPU7x 是 Ironwood 系列的第一個版本，也是第七代 TPU。 Google CloudIronwood 世代的設計目標是支援大規模 AI 訓練和推論。詳情請參閱 TPU7x。

如需針對 TPU7x 最佳化的更多範例，請參閱 GitHub 上的 Ironwood TPU 訓練配方。

使用 MaxText 和 Cluster Toolkit 部署訓練工作負載

使用 Cluster Toolkit 建立可供正式環境使用的 GKE 叢集，並執行訓練工作負載。

以下各節說明如何使用 MaxText 和 Cluster Toolkit 部署訓練工作負載。

事前準備

開始之前，請先完成下列步驟：

  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 叢集」一文中的操作說明，部署 GKE TPU7x 叢集。請確認叢集節點集區拓撲符合訓練工作負載需求 (例如下方範例 MaxText 工作負載的 4x4x8)。

建構及上傳 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 模型是 MaxText 架構中的設定，專為訓練約有 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"

部署訓練工作負載

使用 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) 會附加至工作負載名稱，確保名稱不重複。

後續步驟