使用 TPU7x (Ironwood) 訓練模型

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

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

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

佈建 TPU

您可以使用下列方法佈建及管理 TPU7x：

  • GKE：您可以透過 GKE 佈建及管理 TPU，做為容器化機器學習工作負載的加速器集區。使用 Google Cloud CLI 手動建立 GKE 叢集執行個體，精確自訂或擴充現有的實際工作環境 GKE 環境。詳情請參閱關於 GKE 中的 TPU。
  • GKE 和 XPK：XPK 是一項指令列工具，可簡化在 GKE 上建立叢集和執行工作負載的作業。這項服務專為機器學習從業人員設計，可讓他們佈建 TPU 並執行訓練工作，不必具備深入的 Kubernetes 專業知識。使用 XPK 快速建立 GKE 叢集，並執行工作負載進行概念驗證和測試。詳情請參閱 XPK GitHub 存放區。
  • GKE 和 TPU 叢集導向器：TPU 叢集導向器可透過「所有容量」模式預留使用，讓您完全存取所有預留容量 (無保留)，並全面掌握 TPU 硬體拓撲、使用率和健康狀態。詳情請參閱所有容量模式總覽。

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

使用加速處理套件 (XPK) 建立 GKE 叢集，進行概念驗證和測試。XPK 是一項指令列工具，旨在簡化機器學習工作負載的佈建、管理和執行作業。

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

事前準備

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

  • 確認您有 Google Cloud 啟用計費功能的專案。
  • 存取 TPU7x。詳情請洽詢您的帳戶團隊。
  • 確認您用於 XPK 的帳戶具有 XPK GitHub 存放區中列出的角色。

安裝 XPK 和依附元件

  1. 安裝 XPK。請按照 XPK GitHub 存放區中的操作說明操作。

  2. 按照管理員提供的操作說明安裝 Docker，或參閱官方安裝說明。安裝完成後，請執行下列指令來設定 Docker 並測試安裝：

    gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker

  3. 請設定下列環境變數：

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"

    更改下列內容：

    • YOUR_PROJECT_ID：您的 Google Cloud 專案 ID
    • YOUR_ZONE：要建立叢集的可用區。預覽版僅支援 us-central1-c。
    • YOUR_CLUSTER_NAME：新叢集的名稱。
    • YOUR_ACCELERATOR_TYPE：TPU 版本和拓撲。例如：tpu7x-4x4x8。如需支援的拓撲清單，請參閱「支援的設定」。
    • YOUR_RESERVATION_NAME：預留項目名稱。如要使用共用預留項目，請使用 projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME。
    • YOUR_BUCKET_NAME：Cloud Storage bucket 的名稱，也是模型訓練的輸出目錄。

  4. 如果沒有現成的 Cloud Storage bucket，請使用下列指令建立一個：

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

建立單一 NIC 單一 Slice 叢集

  1. 請按照「設定 MTU」一節的操作說明，調整網路設定。

  2. 填入 ${CLUSTER_ARGUMENTS} 變數，您會在 xpk cluster create 指令中使用這個變數：

    export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"

  3. 使用 xpk cluster create 指令，建立含有 TPU7x 節點集區的 GKE 叢集：

    xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

    將 --cluster-cpu-machine-type 旗標設為 n1-standard-8 (或更大)，可確保預設節點集區有足夠的 CPU 供系統 Pod 使用，例如 JobSet 網路鉤子，避免發生錯誤。根據預設，XPK 會使用 e2-standard-16。部分區域僅支援特定 CPU 類型，因此您可能需要變更 n1、n2 和 e2 類型。否則可能會發生配額錯誤。

  4. 新增維護作業排除時段，避免叢集升級：

    gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

    更改下列內容：

    • EXCLUSION_START_TIME：您選取的維護作業排除時間開始時間，格式為 YYYY-MM-DDTHH:MM:SSZ。
    • EXCLUSION_END_TIME：您選取的維護排除時間結束時間，格式為 YYYY-MM-DDTHH:MM:SSZ。

建構或上傳 MaxText Docker 映像檔

您可以透過 MaxText 提供的指令碼在本機建構 Docker 映像檔，也可以使用預先建構的映像檔。

在本機建構

下列指令會將本機目錄複製到容器中：

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

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

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

成功執行指令後，您應該會看到名為 maxtext_base_image 的圖片在本機建立。您可以直接在 xpk 工作負載指令中使用本機映像檔。

上傳圖片 (選用)

使用上一節的指示在本機建構 Docker 映像檔後，即可使用下列指令將 MaxText Docker 映像檔上傳至登錄檔：

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

成功執行這項指令後，您應該會在 gcr.io 中看到名為 gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME 的 MaxText 圖片。

定義 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_xpk \
        steps=30"

部署訓練工作負載

執行 xpk workload create 指令來部署訓練工作。您必須指定 --base-docker-image 標記來使用 MaxText 基礎映像檔，或是指定 --docker-image 標記和要使用的映像檔。您可以選擇加入 --enable-debug-logs 標記來啟用偵錯記錄功能。

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

叢集內的工作負載名稱不得重複。在本例中，$(date +%H%M) 會附加至工作負載名稱，確保名稱不重複。

後續步驟