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

使用 TPU v5e 訓練模型

TPU v5e 每個 Pod 的晶片數量較少 (256 個)，經過最佳化，是 Transformer、文字轉圖像和卷積類神經網路 (CNN) 訓練、微調和服務的高價值產品。如要進一步瞭解如何使用 Cloud TPU v5e 執行服務，請參閱「使用 v5e 執行推論」。

如要進一步瞭解 Cloud TPU v5e TPU 硬體和設定，請參閱「TPU v5e」。

開始使用

以下各節說明如何開始使用 TPU v5e。

要求配額

您必須具備配額，才能使用 TPU v5e 進行訓練。隨選 TPU、預留 TPU 和 TPU Spot VM 有不同的配額類型。如要使用 TPU v5e 進行推論，則須另外取得配額。如要進一步瞭解配額，請參閱「配額」。如要申請 TPU v5e 配額，請與 Cloud 銷售人員聯絡。

建立 Google Cloud 帳戶和專案

您必須擁有 Google Cloud 帳戶和專案，才能使用 Cloud TPU。詳情請參閱「設定 Cloud TPU 環境」。

建立 Cloud TPU

最佳做法是使用 queued-resource create 指令，將 Cloud TPU v5e 佈建為排入佇列的資源。詳情請參閱「管理已排入佇列的資源」。

您也可以使用 Create Node API (gcloud compute tpus tpu-vm create) 佈建 Cloud TPU v5e。詳情請參閱「管理 TPU 資源」。

如要進一步瞭解可用於訓練的 v5e 設定，請參閱「用於訓練的 Cloud TPU v5e 類型」。

架構設定

本節說明使用 JAX 搭配 TPU v5e 訓練自訂模型的一般設定程序。如要瞭解 PyTorch 訓練選項，請參閱「在 Cloud TPU 上使用 PyTorch 訓練模型」。

如需推論設定操作說明，請參閱「v5e 推論簡介」。

定義一些環境變數：

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

JAX 設定

如果切片形狀大於 8 個晶片，一個切片中會有多個 VM。在這種情況下，您需要使用 --worker=all 旗標，在單一步驟中對所有 TPU VM 執行安裝作業，而不必使用 SSH 分別登入每個 VM：

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

指令旗標說明

TPU_NAME：使用者指派的 TPU 文字 ID，在分配佇列資源要求時建立。
PROJECT_ID： Google Cloud 專案名稱。使用現有專案，或在「設定 Google Cloud 專案」建立新專案
ZONE：如要瞭解支援的可用區，請參閱 TPU 區域和可用區文件。
worker：可存取基礎 TPU 的 TPU VM。

您可以執行下列指令，檢查裝置數量 (此處顯示的輸出內容是以 v5litepod-16 切片產生)。這段程式碼會檢查 JAX 是否能偵測到 Cloud TPU TensorCore 並執行基本作業，藉此測試所有項目是否都已正確安裝：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

輸出內容大致如下：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() 會顯示指定切片中的方塊總數。jax.local_device_count() 表示這個配量中單一 VM 可存取的晶片數量。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

輸出內容大致如下：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

請嘗試本文件中的 JAX 教學課程，開始使用 JAX 進行 v5e 訓練。

清除所用資源

工作階段結束時，請刪除 TPU 和排入佇列的資源。如要刪除排入佇列的資源，請先刪除切片，再刪除排入佇列的資源，步驟如下：

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --quiet

這兩個步驟也可用於移除處於 FAILED 狀態的佇列資源要求。

JAX/FLAX 範例

以下各節說明如何在 TPU v5e 上訓練 JAX 和 FLAX 模型的範例。

在 v5e 上訓練 ImageNet

本教學課程說明如何使用假輸入資料，在 v5e 上訓練 ImageNet。如要使用實際資料，請參閱 GitHub 上的 README 檔案。

設定

建立環境變數：
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id
```
環境變數說明
- PROJECT_ID：您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
- TPU_NAME：TPU 的名稱。
- ZONE：建立 TPU VM 的可用區。如要進一步瞭解支援的可用區，請參閱「TPU 區域和可用區」。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- RUNTIME_VERSION：Cloud TPU 軟體版本。
- SERVICE_ACCOUNT：服務帳戶的電子郵件地址。如要查看，請前往 Google Cloud 控制台的「服務帳戶」頁面。
  例如： tpu-service-account@PROJECT_ID.iam.gserviceaccount.com
- QUEUED_RESOURCE_ID：使用者指派的待處理資源要求文字 ID。

建立 TPU 資源：

gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --node-id=$TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --runtime-version=$RUNTIME_VERSION \
   --service-account=$SERVICE_ACCOUNT

當排入佇列的資源處於 ACTIVE 狀態時，您就能透過 SSH 連線至 TPU VM：

gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
   --project=$PROJECT_ID \
   --zone=$ZONE

當 QueuedResource 處於 ACTIVE 狀態時，輸出內容會與下列內容相似：

 state: ACTIVE

安裝最新版本的 JAX 和 jaxlib：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

複製 ImageNet 模型並安裝對應的需求條件：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="git clone https://github.com/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"

如要生成假資料，模型需要資料集的維度資訊。這項資訊可從 ImageNet 資料集的中繼資料收集：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"

訓練模型

完成上述所有步驟後，即可訓練模型。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

刪除 TPU 和排入佇列的資源

工作階段結束時，請刪除 TPU 和排入佇列的資源。

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

Hugging Face FLAX 模型

在 FLAX 中實作的 Hugging Face 模型可直接在 Cloud TPU v5e 上運作。本節提供執行熱門模型的操作說明。

在 Imagenette 上訓練 ViT

本教學課程說明如何在 Cloud TPU v5e 上，使用 Fast AI Imagenette 資料集訓練 HuggingFace 的 Vision Transformer (ViT) 模型。

ViT 模型是第一個成功在 ImageNet 上訓練 Transformer 編碼器的模型，與卷積網路相比，成效相當出色。詳情請參閱「ViT 總覽」。

設定

建立環境變數：
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id
```
環境變數說明
- PROJECT_ID：您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
- TPU_NAME：TPU 的名稱。
- ZONE：建立 TPU VM 的可用區。如要進一步瞭解支援的可用區，請參閱「TPU 區域和可用區」。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- RUNTIME_VERSION：Cloud TPU 軟體版本。
- SERVICE_ACCOUNT：服務帳戶的電子郵件地址。如要查看，請前往 Google Cloud 控制台的「服務帳戶」頁面。
  例如： tpu-service-account@PROJECT_ID.iam.gserviceaccount.com
- QUEUED_RESOURCE_ID：使用者指派的待處理資源要求文字 ID。

建立 TPU 資源：

gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --node-id=$TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --runtime-version=$RUNTIME_VERSION \
   --service-account=$SERVICE_ACCOUNT

當排入佇列的資源處於 ACTIVE 狀態時，您就能透過 SSH 連線至 TPU VM：

gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
   --project=$PROJECT_ID \
   --zone=$ZONE

當排入佇列的資源處於 ACTIVE 狀態時，輸出內容會與下列內容相似：

 state: ACTIVE

安裝 JAX 及其程式庫：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

下載 Hugging Face 存放區並安裝必要項目：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.19 && sed -i "s/torchvision==0.12.0+cpu/torchvision==0.22.1/" examples/flax/vision/requirements.txt && pip install -r examples/flax/vision/requirements.txt && pip install tf-keras'

下載 Imagenette 資料集：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

訓練模型

使用預先對應的緩衝區 (4 GB) 訓練模型。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

刪除 TPU 和排入佇列的資源

工作階段結束時，請刪除 TPU 和排入佇列的資源。

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

ViT 基準測試結果

訓練指令碼是在 v5litepod-4、v5litepod-16 和 v5litepod-64 上執行。下表列出不同加速器類型的處理量。

加速器類型	v5litepod-4	v5litepod-16	v5litepod-64
訓練週期	3	3	3
全域批次大小	32	128	512
處理量 (樣本數/秒)	263.40	429.34	470.71

訓練寶可夢擴散模型

本教學課程說明如何在 Cloud TPU v5e 上，使用寶可夢資料集訓練 HuggingFace 的 Stable Diffusion 模型。

Stable Diffusion 模型是潛在文字轉圖像模型，可根據任何輸入文字生成擬真圖像。詳情請參閱下列資源：

設定

為儲存空間 bucket 名稱設定環境變數：
```
export GCS_BUCKET_NAME=your_bucket_name
```

設定模型輸出的儲存空間 bucket：

Google Cloud CLI storage buckets create gs://$GCS_BUCKET_NAME \
     --project=$PROJECT_ID \
     --location=$ZONE

建立環境變數：
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west1-c
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id
```
環境變數說明
- PROJECT_ID：您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
- TPU_NAME：TPU 的名稱。
- ZONE：建立 TPU VM 的可用區。如要進一步瞭解支援的可用區，請參閱「TPU 區域和可用區」。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- RUNTIME_VERSION：Cloud TPU 軟體版本。
- SERVICE_ACCOUNT：服務帳戶的電子郵件地址。如要查看，請前往 Google Cloud 控制台的「服務帳戶」頁面。
  例如： tpu-service-account@PROJECT_ID.iam.gserviceaccount.com
- QUEUED_RESOURCE_ID：使用者指派的待處理資源要求文字 ID。

建立 TPU 資源：

gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --node-id=$TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --runtime-version=$RUNTIME_VERSION \
   --service-account=$SERVICE_ACCOUNT

排入佇列的資源處於 ACTIVE 狀態後，您就能透過 SSH 連線至 TPU VM：

gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
   --project=$PROJECT_ID \
   --zone=$ZONE

當排入佇列的資源處於 ACTIVE 狀態時，輸出內容會與下列內容相似：

 state: ACTIVE

安裝 JAX 及其程式庫。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

下載 HuggingFace 存放區並安裝必要項目。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
     --project=$PROJECT_ID \
     --zone=$ZONE \
     --worker=all \
     --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install -U -r examples/text_to_image/requirements_flax.txt && pip install tensorflow==2.17.1 clu && pip install tensorboard==2.17.1'

訓練模型

使用預先對應的緩衝區 (4 GB) 訓練模型。

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project=$PROJECT_ID --worker=all --command="
    git clone https://github.com/google/maxdiffusion
    cd maxdiffusion
    pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    pip3 install -r requirements.txt
    pip3 install .
    pip3 install gcsfs
    export LIBTPU_INIT_ARGS=''
    python -m src.maxdiffusion.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run \
    jax_cache_dir=gs://${GCS_BUCKET_NAME} activations_dtype=bfloat16 weights_dtype=bfloat16 \
    per_device_batch_size=1 precision=DEFAULT dataset_save_location=gs://${GCS_BUCKET_NAME} \
    output_dir=gs://${GCS_BUCKET_NAME}/ attention=flash"

清除所用資源

工作階段結束時，請刪除 TPU、排入佇列的資源和 Cloud Storage bucket。

刪除 TPU：

gcloud compute tpus tpu-vm delete $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --quiet

刪除排入佇列的資源：

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --quiet

刪除 Cloud Storage bucket：

gcloud storage rm -r gs://$GCS_BUCKET_NAME

擴散基準測試結果

訓練指令碼在 v5litepod-4、v5litepod-16 和 v5litepod-64 上執行。下表列出處理量。

加速器類型	v5litepod-4	v5litepod-16	v5litepod-64
訓練步驟	1500	1500	1500
全域批次大小	32	64	128
處理量 (樣本數/秒)	36.53	43.71	49.36

在 Cloud TPU 上使用 PyTorch 訓練模型

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

使用 TPU v5e 訓練模型

開始使用

要求配額

建立 Google Cloud 帳戶和專案

建立 Cloud TPU

架構設定

JAX 設定

指令旗標說明

清除所用資源

JAX/FLAX 範例

在 v5e 上訓練 ImageNet

設定

環境變數說明

訓練模型

刪除 TPU 和排入佇列的資源

Hugging Face FLAX 模型

在 Imagenette 上訓練 ViT

設定

環境變數說明

訓練模型

刪除 TPU 和排入佇列的資源

ViT 基準測試結果

訓練寶可夢擴散模型

設定

環境變數說明

訓練模型

清除所用資源

擴散基準測試結果

在 Cloud TPU 上使用 PyTorch 訓練模型