使用 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 上建立叢集和執行工作負載的作業。機器學習從業人員無須具備深入的 Kubernetes 專業知識，即可透過這項服務佈建 TPU 並執行訓練工作。使用 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 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_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 叢集

選擇下列其中一個選項來建立叢集。建議使用 MTU 為 8,896 的自訂網路，以獲得最佳效能。

自訂網路

如要建立 MTU 為 8,896 的自訂網路，並用於叢集，請按照下列步驟操作：

  1. 為網路和防火牆名稱設定環境變數：

    export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME

    更改下列內容：

    • NETWORK_NAME：網路名稱。
    • FIREWALL_NAME：網路防火牆規則的名稱。

  2. 建立 MTU 為 8,896 的自訂網路：

    gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

  3. 建立防火牆規則，允許網路上的 TCP、ICMP 和 UDP 流量：

    gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

  4. 為 XPK 叢集引數設定環境變數，以使用您建立的網路：

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

  5. 建立 XPK 叢集。下列指令會佈建隨需容量：

    xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

    如要使用預留容量，請將 --on-demand 替換為 --reservation=RESERVATION_NAME。如要使用 TPU Spot VM，請將 --on-demand 替換為 --spot。

預設網路

如果您不需要高 MTU 網路，可以建立使用預設虛擬私有雲網路的叢集。下列指令會佈建隨需容量：

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

如要使用預留容量，請將 --on-demand 替換為 --reservation=RESERVATION_NAME。如要使用 TPU Spot VM，請將 --on-demand 替換為 --spot。

建構或上傳 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

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

成功執行指令後，您應該會看到名為 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) 會附加至工作負載名稱，確保名稱不重複。

後續步驟