Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

使用 TPU v6e 訓練模型

本文件將引導您在 Cloud TPU v6e (又稱為 Trillium) 上訓練模型，包括環境設定、效能最佳化，以及使用 JAX 和 PyTorch 的實際訓練範例。

TPU v6e (又稱 Trillium) 是 Google 第 6 代 TPU，在所有技術層面 (例如 API 和記錄) 以及整份文件中，Trillium 都會稱為 v6e。TPU v6e 每個 Pod 有 256 個晶片，架構與 v5e 高度相似。TPU v6e 經過最佳化調整，適合用於 Transformer、文字轉圖像和卷積類神經網路 (CNN) 的訓練、微調與提供。如要進一步瞭解 TPU v6e 系統架構和設定，請參閱「TPU v6e」。

如要瞭解如何在 Cloud TPU v6e 執行推論，請參閱下列教學課程：

事前準備

開始前，請先完成下列事項：

設定 TPU 的 Google Cloud 專案
安裝 Google Cloud CLI
啟用 Compute Engine API
建立 Compute Engine 服務帳戶並授予權限

詳情請參閱「設定 Cloud TPU 環境」。

確認配額和權限

確認專案具備下列配額：

Compute Engine TPU 加速器配額
Compute Engine VM 配額
IP 位址配額
Hyperdisk Balanced 和其他要使用的磁碟類型配額

注意： v6e 支援 Hyperdisk Balanced 和 Hyperdisk ML，但不支援永久磁碟 (PD)。詳情請參閱「Cloud TPU 資料的儲存空間選項」。

如果您搭配使用 Google Kubernetes Engine (GKE) 和 XPK (Accelerated Processing Kit)，則需要在 Google Cloud 控制台中取得額外權限。詳情請參閱「 Google Cloud 控制台所需的權限」。

佈建選項

您可以透過下列方法佈建及管理 TPU v6e：

Compute Engine (Google Cloud CLI)：您可以使用 Google Cloud CLI 工具，直接將 TPU VM 做為 Compute Engine 資源管理。您可以佈建個別 TPU VM，或使用代管執行個體群組 (MIG) 管理 TPU VM 群組。詳情請參閱「TPU 建立總覽」。
GKE：您可以使用 GKE 佈建及管理 TPU，做為容器化機器學習工作負載的加速器集區。詳情請參閱「GKE 中的 TPU 簡介」。
GKE 和 XPK：XPK 是指令列工具，可簡化 GKE 叢集建立和工作負載執行作業。機器學習從業人員可透過這項服務佈建 TPU 並執行訓練工作，不必具備深厚的 Kubernetes 專業知識。詳情請參閱 XPK GitHub 存放區。

使用 GKE 和 XPK 佈建 v6e TPU

如果您使用 GKE 搭配 v6e TPU，可以透過 Kubernetes 指令或 XPK 佈建 TPU，並訓練或提供模型。如要進一步瞭解如何搭配使用 GKE 和 TPU，請參閱「GKE 中的 TPU 簡介」。

Cloud TPU v6e 支援網路介面卡 (NIC) 設定，可讓您跨多個網路擴充處理量。下列各節提供指令，說明如何使用 XPK 建立支援單一 NIC 或多 NIC 的 GKE 叢集。對於大多數單一配量工作負載，單一 NIC 只要較少設定，就能提供足夠的效能。如果是 Multislice 工作負載和需要高資料擷取速度的工作負載，請使用多 NIC。

使用 XPK 建立支援單一 NIC 的叢集

對於大多數單一配量工作負載，單一 NIC 只要較少設定，就能提供足夠的效能。如果是 Multislice 工作負載和需要高資料擷取速度的工作負載，請使用多 NIC。

下列各節說明如何使用 XPK 建立支援單一 NIC 的 GKE 叢集。

安裝 XPK 並設定環境變數

安裝 XPK。請按照 XPK GitHub 存放區中的指示操作。
設定叢集的環境變數：
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=1
```
請設定下列環境變數：
- CLUSTER_NAME：叢集名稱。
- ZONE：要建立 TPU 叢集的可用區。如要進一步瞭解支援的可用區，請參閱「區域和可用區」。
- PROJECT_ID： Google Cloud 專案 ID。
- ACCELERATOR_TYPE：TPU 類型 (也稱為加速器類型) 會指定要建立的 Cloud TPU 版本和大小。例如：v6e-256。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- NUM_SLICES：叢集的 TPU slice 數量。每個切片都有 ACCELERATOR_TYPE 中指定的晶片數量。如果是單一配量叢集，請將 NUM_SLICES 設為 1。如果是多配量叢集，請根據工作負載的擴充性需求指定配量數量。叢集中的晶片總數為 ACCELERATOR_TYPE 中的晶片數量乘以 NUM_SLICES。

建立叢集

選擇下列其中一個選項來建立叢集。建議使用 MTU 為 8,896 的自訂網路，以達到最佳效能。詳情請參閱「設定 MTU」。

自訂網路

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

為網路和防火牆名稱設定環境變數：
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
更改下列內容：
- NETWORK_NAME：網路名稱。
- FIREWALL_NAME：網路防火牆規則的名稱。

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

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

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

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

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

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

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

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-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=e2-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。

使用 XPK 建立支援多 NIC 的叢集

如果是 Multislice 工作負載或其他需要高網路頻寬的工作負載 (例如資料擷取)，可以使用多 NIC 來提升效能。使用多 NIC 時，每個 TPU VM 都會分配到額外的網路介面，每個介面都會連線至不重複的虛擬私有雲網路，進而提高整體網路處理量。對於大多數單一配量工作負載，單一 NIC 只要較少設定，就能提供足夠的效能。

下列各節說明如何使用 XPK 建立支援多 NIC 的 GKE 叢集。

安裝 XPK 並設定環境變數

安裝 XPK。請按照 XPK GitHub 存放區中的指示操作。
為叢集和主要網路設定環境變數：
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export REGION=REGION
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}

export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}
```
請設定下列環境變數：
- CLUSTER_NAME：叢集名稱。
- REGION：要建立 TPU 叢集的區域。
- ZONE：要建立 TPU 叢集的可用區。如要進一步瞭解支援的可用區，請參閱「區域和可用區」。
- PROJECT_ID： Google Cloud 專案 ID。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。例如，v6e-256。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- NUM_SLICES：叢集的 TPU slice 數量。每個切片都有 ACCELERATOR_TYPE 中指定的晶片數量。如果是單一配量叢集，請將 NUM_SLICES 設為 1。如果是多配量叢集，請根據工作負載的擴充性需求指定配量數量。叢集中的晶片總數為 ACCELERATOR_TYPE 中的晶片數量乘以 NUM_SLICES。

建立主要網路資源

建立主要網路，並將最大傳輸單位 (MTU) 設為 8,896：

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

使用 MTU 為 8,896 的自訂網路，可提供更佳效能。詳情請參閱「設定 MTU」。

建立主要子網路：

gcloud compute networks subnets create ${SUBNET_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --range=10.11.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

為主要網路建立防火牆規則，允許主要網路上的 tcp、icmp 和 udp 流量：

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

為主要網路建立 Cloud Router：

gcloud compute routers create ${ROUTER_NAME_1} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_1} \
    --region=${REGION}

為主要網路設定網路位址轉譯 (NAT)。下列指令可允許叢集流量連至網際網路：

gcloud compute routers nats create ${NAT_CONFIG_1} \
    --router=${ROUTER_NAME_1} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

建立次要網路資源

建立次要網路：

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

為次要網路建立子網路：

gcloud compute networks subnets create ${SUBNET_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --range=10.10.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

建立防火牆規則，允許新網路中的流量：

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --allow tcp,icmp,udp \
    --source-ranges 10.10.0.0/18 \
    --project=${PROJECT_ID}

為次要網路建立 Cloud Router：

gcloud compute routers create ${ROUTER_NAME_2} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_2} \
    --region=${REGION}

為 Cloud Router 建立網路位址轉譯 (NAT) 設定：

gcloud compute routers nats create ${NAT_CONFIG_2} \
    --router=${ROUTER_NAME_2} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

建立叢集

為叢集和節點集區引數設定環境變數，以便使用您建立的網路和子網路：

export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

這些引數會設定叢集，使用您為多 NIC 支援建立的兩個網路。

建立叢集。下列指令會佈建隨選容量：
```
xpk cluster create \
    --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE}  \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
    --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
    --create-vertex-tensorboard
```
如要使用預留容量，請將 --on-demand 替換為 --reservation=RESERVATION_NAME。如要使用 TPU Spot VM，請將 --on-demand 替換為 --spot。

注意： --cluster-cpu-machine-type 旗標會指定叢集中 CPU 主機 VM 的機型。雖然 e2-standard-8 是許多工作負載的建議預設值，但如果工作負載對資料預先處理等工作有特定的 CPU 或記憶體需求，您可以選取其他機器類型。

驗證多 NIC 設定

建立支援多 NIC 的叢集後，您可以建立 XPK 工作負載並新增 --command ifconfig 旗標，驗證兩個 NIC 是否都在使用中。

使用下列指令，在Google Cloud 控制台記錄中顯示 ifconfig 指令的輸出內容。您必須指定 --base-docker-image maxtext_base_image 旗標，才能使用 MaxText 基本映像檔 (如以下範例所示)，或指定 --docker-image 旗標和要使用的映像檔。

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"

如要啟用偵錯記錄檔或使用 Vertex AI TensorBoard，請在指令中新增下列選用引數：

--enable-debug-logs \
--use-vertex-tensorboard

在 Google Cloud 控制台記錄中查看 XPK 工作負載的輸出內容，確認 eth0 和 eth1 的 MTU 都設為 8,896。

設定 JAX 或 PyTorch

下列資源說明如何在 TPU 上設定 JAX 或 PyTorch，具體取決於您使用的佈建和管理方法：

GKE Autopilot：準備 TPU 應用程式
GKE Standard：準備工作負載
GKE 和 XPK：XPK README
使用 JAX 的單一主機 Cloud TPU：使用 JAX 在 Cloud TPU VM 上執行計算
使用 JAX 的多主機 Cloud TPU：在 TPU 配量上執行 JAX 程式碼
使用 PyTorch 的 Cloud TPU：Cloud TPU 上的 PyTorch

如要設定 XPK 並搭配 MaxText 執行，請參閱「使用 XPK 大規模執行 MaxText 」。

使用 SkyPilot

您可以搭配使用 Cloud TPU v6e 和 SkyPilot。SkyPilot 是開放原始碼架構，可簡化執行、管理及擴充 AI 工作負載的程序。您可以在 SkyPilot 中新增 v6e 相關位置和定價資訊。詳情請參閱 SkyPilot TPU v6e 範例。

訓練範例

下列各節提供在 Cloud TPU v6e 訓練 MaxText、MaxDiffusion 和 PyTorch 模型的範例。

這些範例已使用下列軟體版本進行測試：

Python 3.10 以上版本
夜間軟體版本：
- 每晚 JAX 0.4.32.dev20240912
- Nightly LibTPU 0.1.dev20240912+nightly
穩定軟體版本：
- JAX + JAX Lib 0.4.37 版

在 Cloud TPU v6e 上訓練 MaxText 和 MaxDiffusion

以下各節說明 MaxText 和 MaxDiffusion 模型的訓練生命週期。

大致步驟如下：

建構工作負載基礎映像檔。
使用 XPK 執行工作負載。
1. 建構工作負載的訓練指令。
2. 部署工作負載。
追蹤工作負載並查看指標。
如果不需要 XPK 工作負載，請刪除。
刪除不再需要的叢集。

建構底層映像檔

安裝 MaxText 或 MaxDiffusion，並建構 Docker 映像檔：

複製要使用的存放區，並切換至該存放區的目錄：

MaxText：

git clone https://github.com/google/maxtext.git
cd maxtext
git checkout maxtext-v0.2.0

MaxDiffusion：

git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

將 Docker 設為使用 Google Cloud CLI：
```
gcloud auth configure-docker
```

使用下列指令或 JAX AI 圖像建構 Docker 映像檔。如要進一步瞭解 JAX AI 圖像，請參閱「JAX AI 圖像」。

MaxText：

bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable

MaxDiffusion：

bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

在 gcloud CLI 有效設定中設定專案 ID：
```
gcloud config set project ${PROJECT_ID}
```
如果啟動工作負載的機器未在本機建構映像檔，請上傳映像檔。

注意： 請確認您使用的是 Bash 4.2 以上版本。
1. 設定 CLOUD_IMAGE_NAME 環境變數：
```
export CLOUD_IMAGE_NAME=${USER}_runner
```
2. 上傳圖片：
```
bash dependencies/scripts/docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

使用 XPK 執行工作負載

如果您不使用 MaxText 或 MaxDiffusion 設定的預設值，請設定下列環境變數：

export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192

建構模型指令碼。在後續步驟中，這段指令碼會複製為訓練指令。

請先不要執行模型指令碼。

MaxText

MaxText 是高效能、高擴充性的開放原始碼 LLM，以純 Python 和 JAX 編寫，並以 Google Cloud TPU 和 GPU 為訓練和推論目標。

JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

Gemma2

Gemma 是 Google DeepMind 開發的一系列開放權重 LLM，採用 Gemini 研究成果和技術。

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

Mixtral 8x7b

Mixtral 是 Mistral AI 開發的最先進 AI 模型，採用稀疏混合專家 (MoE) 架構。

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

MaxDiffusion

MaxDiffusion 是一組參考實作，包含各種以純 Python 和 JAX 編寫的潛在擴散模型，可在 XLA 裝置 (包括 Cloud TPU 和 GPU) 上執行。Stable Diffusion 是潛在文字轉圖像模型，可根據任何輸入文字生成逼真的圖像。

您需要安裝特定 Git 分支版本，才能執行 MaxDiffusion，如下列訓練指令碼所示。

git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

匯出下列變數：
```
export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
```
環境變數說明
- CLUSTER_NAME：叢集名稱。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- NUM_SLICES：TPU 配量數量。
- YOUR_MODEL_SCRIPT：要以訓練指令執行的模型指令碼。
使用上一步建立的指令碼執行模型。您必須指定 --base-docker-image 旗標，才能使用 MaxText 基本映像檔，或指定 --docker-image 旗標和要使用的映像檔。

您可以選擇新增下列選用旗標：
- 加入 --enable-debug-logs 旗標，即可啟用偵錯記錄檔功能。詳情請參閱「在 MaxText 上對 JAX 進行偵錯」。
- 您可以建立 Vertex AI Experiments，並加入 --use-vertex-tensorboard 旗標，將資料上傳至 Vertex AI TensorBoard。詳情請參閱「使用 Vertex AI 監控 MaxText 上的 JAX」。
```
xpk workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"
```
輸出內容包含追蹤工作負載的連結。開啟連結並點選「記錄」分頁標籤，即可即時追蹤工作負載。

在 MaxText 上對 JAX 進行偵錯

使用 XPK 補充指令，診斷叢集或工作負載無法執行的原因：

XPK 工作負載清單
XPK 檢查器
建立 XPK 工作負載時，請使用 --enable-debug-logs 旗標，在工作負載記錄中啟用詳細記錄功能

使用 Vertex AI 監控 MaxText 上的 JAX

如要使用 TensorBoard，您的 Google Cloud 使用者帳戶必須具備aiplatform.user角色。執行下列指令來授予這個角色：

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

透過 Vertex AI 代管的 TensorBoard 查看純量和設定檔資料。

將所用可用區的資源管理 (CRUD) 要求從 600 提高至 5000。如果工作負載較小，使用的 VM 不到 16 個，這可能不是問題。

安裝 Vertex AI 的依附元件，例如 cloud-accelerator-diagnostics：

# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

使用 --create-vertex-tensorboard 旗標建立叢集，如「建立 Vertex AI TensorBoard」一文所述。您也可以對現有叢集執行這項指令。
使用 --use-vertex-tensorboard 旗標和選用的 --experiment-name 旗標，在執行 XPK 工作負載時建立 Vertex AI Experiments。如需完整步驟清單，請參閱「建立 Vertex AI Experiments，將資料上傳至 Vertex AI TensorBoard」。

記錄中包含 Vertex AI TensorBoard 的連結，如下所示：

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

您也可以在 Google Cloud 控制台中找到 Vertex AI TensorBoard 連結。前往 Google Cloud 控制台的「Vertex AI Experiments」頁面。從下拉式選單中選取適當的區域。

TensorBoard 目錄也會寫入您使用 ${BASE_OUTPUT_DIR} 指定的 Cloud Storage bucket。

刪除 XPK 工作負載

使用 xpk workload delete 指令，根據工作前置字串或工作狀態刪除一或多個工作負載。如果您傳送的 XPK 工作負載不再需要執行，或有工作卡在佇列中，這個指令可能會很有用。

刪除叢集

使用 xpk cluster delete 指令刪除叢集：

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

MaxDiffusion 基準測試結果

我們在 v6e-4、v6e-16 和兩個 v6e-16 上執行 MaxDiffusion 的訓練指令碼。下表列出測得的處理量。

	v6e-4	v6e-16	兩個 v6e-16
訓練步數	0.069	0.073	0.13
全域批次大小	8	32	64
處理量 (樣本數/秒)	115.9	438.4	492.3

在 Cloud TPU 上使用 PyTorch 訓練模型

您可以直接使用 TorchTPU 訓練模型，或重複使用 TorchTitan 參考程式庫。如要開始使用 v6e 工作負載，請參閱 TorchTitan Trillium 方案。如要進一步瞭解 Cloud TPU 上的 PyTorch 生態系統，請參閱「Cloud TPU 上的 PyTorch」。

使用 TPU v6e 訓練模型

事前準備

確認配額和權限

佈建選項

使用 GKE 和 XPK 佈建 v6e TPU

使用 XPK 建立支援單一 NIC 的叢集

安裝 XPK 並設定環境變數

建立叢集

自訂網路

預設網路

使用 XPK 建立支援多 NIC 的叢集

安裝 XPK 並設定環境變數

建立主要網路資源

建立次要網路資源

建立叢集

驗證多 NIC 設定

設定 JAX 或 PyTorch

使用 SkyPilot

訓練範例

在 Cloud TPU v6e 上訓練 MaxText 和 MaxDiffusion

建構底層映像檔

使用 XPK 執行工作負載

MaxText

Gemma2

Mixtral 8x7b

MaxDiffusion

環境變數說明

在 MaxText 上對 JAX 進行偵錯

使用 Vertex AI 監控 MaxText 上的 JAX

刪除 XPK 工作負載

刪除叢集

MaxDiffusion 基準測試結果

在 Cloud TPU 上使用 PyTorch 訓練模型