Cloud TPU 多配量總覽

注意： 本頁內容適用於 Cloud TPU API，該 API 已停止積極開發。Cloud TPU API 只會收到錯誤修正和安全性更新。建議您使用 Google Kubernetes Engine 執行多配量作業。詳情請參閱「在 GKE 中部署 TPU 多配量」。

重要事項： 多配量僅支援 Cloud TPU v4 和後續TPU 版本。

Cloud TPU Multislice 是一項全端效能擴充技術，可讓訓練工作在單一 Pod 內使用多個 TPU 配量，或在多個 Pod 的配量上使用標準資料平行處理。使用 TPU v4 晶片時，這表示訓練工作單次執行可使用超過 4096 個晶片。如果訓練工作需要少於 4096 個晶片，單一 Slice 就能提供最佳效能。不過，多個較小的切片更容易取得，因此搭配較小的切片使用 Multislice 時，啟動時間會更快。

在多配量設定中，每個配量中的 TPU 晶片會透過晶片間互連 (ICI) 進行通訊。不同配量中的 TPU 晶片會將資料傳輸至 CPU (主機)，再由 CPU 透過資料中心網路 (DCN) 傳輸資料，藉此進行通訊。如要進一步瞭解如何使用 Multislice 擴充，請參閱「如何使用 Multislice 將 AI 訓練作業的規模擴充至數萬個 Cloud TPU 晶片」。

開發人員不必編寫程式碼，即可實作切片間 DCN 通訊。XLA 編譯器會為您產生該程式碼，並將通訊與運算重疊，以發揮最大效能。

概念

加速器類型

構成多配量的每個 TPU 配量形狀。多配量要求中的每個配量都屬於相同的加速器類型。加速器類型由 TPU 類型 (v4 以上) 和 TensorCore 數量組成。舉例來說，v5litepod-128 會指定 128 個 TensorCore 的 TPU v5e。

自動修復

當切片遇到維護事件、先占或硬體故障時，Cloud TPU 會建立新的切片。如果資源不足以建立新切片，系統會等到硬體可用時再完成建立作業。建立新的 slice 後， Multislice 環境中的所有其他 slice 都會重新啟動，以便繼續訓練。只要正確設定開機指令碼，訓練指令碼就能自動重新啟動，無需使用者介入，並從最新檢查點載入及繼續執行。

資料中心網路 (DCN)

與 ICI 相比，延遲時間較長、處理量較低的網路，用於連接多配量設定中的 TPU 配量。

群組排程

所有 TPU 切片會同時佈建，保證所有切片都佈建成功，或完全不佈建。

晶片間互連網路 (ICI)

高速、低延遲的內部連結，用於連接 TPU Pod 內的 TPU。

Multislice

兩個以上的 TPU 晶片配量，可透過 DCN 通訊。

節點

在多配量情境中，節點是指單一 TPU 配量。多配量中的每個 TPU 配量都會獲得一個節點 ID。

開機指令碼

標準 Compute Engine 開機指令碼，每次 VM 啟動或重新啟動時都會執行。如果是 Multislice，則是在 QR code 建立要求中指定。如要進一步瞭解 Cloud TPU 開機指令碼，請參閱「管理 TPU 資源」。

Tensor

用於在機器學習模型中表示多維度資料的資料結構。

Cloud TPU 容量類型

TPU 可透過不同類型的容量建立 (請參閱「TPU 的計費方式」中的「使用選項」)：

• 預留：如要使用預留項目，您必須與 Google 簽訂預留協議。建立資源時，請使用 --reserved 旗標。

• Spot：使用 Spot VM 指定先占配額。系統可能會優先處理優先順序較高的工作，因此會先占您的資源。建立資源時，請使用 --spot 旗標。

• 隨選：指定隨選配額，不需預留，也不會遭到先占。TPU 要求會排入 Cloud TPU 提供的隨選配額佇列，但無法保證資源可用性。預設為選取，不需使用旗標。

開始使用

1. 設定 Cloud TPU 環境。

2. 在 Google Cloud 控制台啟用 Cloud Shell。

   啟用 Cloud Shell

   Google Cloud 控制台底部會開啟 Cloud Shell 工作階段，並顯示指令列提示。Cloud Shell 是已安裝 Google Cloud CLI 的殼層環境，並已設定適用於您目前專案的值。工作階段可能要幾秒鐘的時間才能初始化。

如要使用 Multislice，TPU 資源必須以排入佇列的資源形式管理。

入門範例

本教學課程使用的程式碼，取自 MaxText GitHub 存放區。MaxText 是高效能、可任意擴充、開放原始碼且經過充分測試的基本 LLM，以 Python 和 Jax 編寫。MaxText 專為在 Cloud TPU 上有效率地訓練而設計。

shardings.py 中的程式碼旨在協助您開始實驗不同的平行處理選項。例如資料平行處理、完全分區資料平行處理 (FSDP) 和張量平行處理。程式碼可從單一配量擴充至多配量環境。

ICI 平行處理

ICI 是指連結單一配量中 TPU 的高速互連網路。ICI sharding 對應至 slice 內部的 sharding。shardings.py 提供三個 ICI 平行處理參數：

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

您為這些參數指定的值，會決定每種平行處理方法的 shard 數量。

這些輸入內容必須受到限制，因此 ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism 等於切片中的晶片數量。

下表顯示 v4-8 中四個可用晶片的 ICI 平行處理使用者輸入內容範例：

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
4 向 FSDP	1	4	1
4 向 Tensor 平行處理	1	1	4
雙向 FSDP + 雙向 Tensor 平行處理	1	2	2

請注意，在大多數情況下，ici_data_parallelism 應保留為 1，因為 ICI 網路速度夠快，幾乎一律會優先使用 FSDP，而非資料平行處理。

這個範例假設您熟悉如何在單一 TPU 切片上執行程式碼，例如使用 JAX 在 Cloud TPU VM 上執行運算。這個範例說明如何在單一切片上執行 shardings.py。

1. 設定環境：

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   變數說明

   • QR_ID：使用者指派的待處理資源 ID。
   • TPU_NAME：使用者指派的 TPU 名稱。
   • PROJECT： Google Cloud 專案名稱
   • ZONE：指定要建立資源的可用區。
   • NETWORK_NAME：虛擬私有雲網路名稱。
   • SUBNETWORK_NAME：虛擬私有雲網路中的子網路名稱
   • RUNTIME_VERSION：Cloud TPU 軟體版本。
   • ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。
   • EXAMPLE_TAG_1, EXAMPLE_TAG_2 …：用來識別網路防火牆有效來源或目標的標記
   • SLICE_COUNT：配量數。最多 256 個配量。
   • STARTUP_SCRIPT：如果指定開機指令碼，指令碼會在 TPU slice 佈建或重新啟動時執行。

2. 為 gcloud 建立安全殼層金鑰。建議您留空密碼 (執行下列指令後按兩次 Enter 鍵)。如果系統提示 google_compute_engine 檔案已存在，請取代現有版本。

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. 佈建 TPU：

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

   Google Cloud CLI 不支援所有建立 QR code 的選項，例如標記。詳情請參閱「建立 QR」。

   控制台

   1. 前往 Google Cloud 控制台的「TPUs」頁面：

      前往 TPU

   2. 按一下「Create TPU」(建立 TPU)。

   3. 在「Name」(名稱) 欄位中，輸入 TPU 的名稱。

   4. 在「Zone」(可用區) 方塊中，選取要建立 TPU 的可用區。

   5. 在「TPU type」(TPU 類型) 方塊中，選取加速器類型。加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。

   6. 在「TPU 軟體版本」方塊中，選取軟體版本。建立 Cloud TPU VM 時，TPU 軟體版本會指定要安裝的 TPU 執行階段版本。詳情請參閱「TPU 軟體版本」。

   7. 按一下「啟用佇列」切換鈕。

   8. 在「已加入佇列的資源名稱」欄位中，輸入已加入佇列的資源要求名稱。

   9. 點按「建立」，建立排入佇列的資源要求。

4. 請等到佇列資源處於 ACTIVE 狀態，這表示 worker 節點處於 READY 狀態。排入佇列的資源開始佈建後，視資源大小而定，可能需要 1 到 5 分鐘才能完成。您可以使用 gcloud CLI 或 Google Cloud 控制台，查看已排入佇列的資源要求狀態：

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   控制台

   1. 前往 Google Cloud 控制台的「TPUs」頁面：

      前往 TPU

   2. 按一下「已加入佇列的資源」分頁標籤。

   3. 按一下排入佇列的資源要求名稱。

5. 使用 SSH 連線至 TPU VM：

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

6. 將 MaxText (包括 shardings.py) 複製到 TPU VM：

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. 安裝 Python 3.10：

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

8. 建立並啟用虛擬環境：

   注意： 為每個專案建立虛擬環境，是管理 Python 依附元件的最佳做法。詳情請參閱 venv 說明文件。

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

9. 在 MaxText 存放區目錄中，執行設定指令碼，在 TPU 節點上安裝 JAX 和其他依附元件。設定指令碼需要幾分鐘才能執行完畢。

   $ bash setup.sh

10. 執行下列指令，在 TPU 切片上執行 shardings.py。

    $ python3 -m pedagogical_examples.shardings \
      --ici_fsdp_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

    您可以在記錄中查看結果。您的 TPU 每秒應可達到約 260 TFLOP，或高達 90%以上的 FLOP 利用率！在本例中，我們選取的批次大小接近 TPU 高頻寬記憶體 (HBM) 的容量上限。

11. 歡迎探索 ICI 的其他分片策略，例如嘗試下列組合：

    $ python3 -m pedagogical_examples.shardings \
      --ici_tensor_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

12. 完成後，請刪除排入佇列的資源和 TPU 節點。您應從設定 Slice 的環境執行這些清理步驟 (先執行 exit 結束 SSH 工作階段)。刪除作業會在兩到五分鐘內完成。如果使用 gcloud CLI，可以搭配選用的 --async 旗標，在背景執行這項指令。

    重要事項： 使用完畢後，必須刪除已加入佇列的資源。否則，這些資源會繼續消耗排入佇列的資源配額，這項配額與 TPU 配額分開追蹤。

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    控制台

    1. 前往 Google Cloud 控制台的「TPUs」頁面：

       前往「TPU」

    2. 按一下「Queued resources」(已排入佇列的資源) 分頁標籤。

    3. 找出要取消的排入佇列資源要求，然後勾選旁邊的核取方塊。

    4. 按一下「刪除」圖示 。

使用 DCN 平行處理的多切片資料分割

shardings.py 指令碼會採用三個參數，指定 DCN 平行處理，對應於每種資料平行處理的分片數量：

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

這些參數的值必須受到限制，使 dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism 等於切片數量。

以兩個切片為例，請使用 --dcn_data_parallelism = 2。

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	切片數量
雙向資料平行處理	2	1	1	2

dcn_tensor_parallelism 應一律設為 1，因為 DCN 不適合這類資料分割。對於 v4 晶片上的典型 LLM 工作負載，dcn_fsdp_parallelism 也應設為 1，因此 dcn_data_parallelism 應設為切片數量，但這取決於應用程式。

如果增加切片數量 (假設切片大小和每個切片的批次大小不變)，資料平行處理量就會增加。

在 Multislice 環境中執行 shardings.py

您可以使用 multihost_runner.py，在多配量環境中執行 shardings.py，也可以在每個 TPU VM 上執行 shardings.py。這裡我們使用 multihost_runner.py。下列步驟與 MaxText 存放區中的「入門指南：在多個切片上快速實驗」非常相似，但我們在這裡執行 shardings.py，而不是 train.py 中較複雜的大型語言模型。

multihost_runner.py 工具經過最佳化，適合快速實驗，可重複使用相同的 TPU。由於 multihost_runner.py 指令碼依賴長時間執行的 SSH 連線，因此不建議用於任何長時間執行的工作。如要執行較長時間的工作 (例如數小時或數天)，建議使用 multihost_job.py。

在本教學課程中，我們使用「執行器」一詞，表示執行 multihost_runner.py 指令碼的機器。我們使用「worker」一詞，表示構成切片的 TPU VM。您可以在本機或與切片位於相同專案的任何 Compute Engine VM 上執行 multihost_runner.py。不支援在工作人員上執行 multihost_runner.py。

multihost_runner.py 會自動使用 SSH 連線至 TPU 工作站。

在這個範例中，您會在兩個 v5e-16 切片上執行 shardings.py，總共四個 VM 和 16 個 TPU 晶片。您可以修改範例，在更多 TPU 上執行。

設定環境

1. 在執行器機器上複製 MaxText：

   $ git clone https://github.com/AI-Hypercomputer/maxtext

2. 前往存放區目錄。

   $ cd maxtext

3. 為 gcloud 建立安全殼層金鑰，建議您將密碼留空 (執行下列指令後，按兩次 Enter 鍵)。如果系統提示google_compute_engine檔案已存在，請選取不要保留現有版本。

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. 新增環境變數，將 TPU 節點數設為 2。

     $ export SLICE_COUNT=2
     

5. 使用 queued-resources create 指令或 Google Cloud 控制台建立多切片環境。

   gcloud

   注意： 請先確認您有相應的配額，再選取 --reserved、--spot 或預設隨選配額。如要瞭解配額類型，請參閱「配額政策」。

   下列指令說明如何建立 v5e Multislice TPU。如要使用其他 TPU 版本，請指定其他 accelerator-type 和 runtime-version。

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   控制台

   1. 前往 Google Cloud 控制台的「TPUs」頁面：

      前往「TPU」

   2. 按一下「建立 TPU」。

   3. 在「Name」(名稱) 欄位中，輸入 TPU 的名稱。

   4. 在「Zone」(可用區) 方塊中，選取要建立 TPU 的可用區。

   5. 在「TPU type」(TPU 類型) 方塊中，選取加速器類型。加速器類型會指定要建立的 Cloud TPU 版本和大小。Multislice 僅支援 Cloud TPU v4 和後續 TPU 版本。如要進一步瞭解 TPU 版本，請參閱「TPU 版本」。

   6. 在「TPU 軟體版本」方塊中，選取軟體版本。建立 Cloud TPU VM 時，TPU 軟體版本會指定要安裝在 TPU VM 上的 TPU 執行階段版本。詳情請參閱「TPU 軟體版本」。

   7. 按一下「啟用佇列」切換鈕。

   8. 在「Queued resource name」(待處理資源名稱) 欄位中，輸入待處理資源要求的名稱。

   9. 按一下「將此設為多配量 TPU」核取方塊。

   10. 在「Slice count」(切片數量) 欄位中，輸入要建立的切片數量。

   11. 點按「建立」，建立排入佇列的資源要求。

6. 排入佇列的資源開始佈建後，最多可能需要五分鐘才能完成，視排入佇列的資源大小而定。等待佇列資源進入 ACTIVE 狀態。您可以使用 gcloud CLI 或 Google Cloud 控制台，查看已排入佇列的資源要求狀態：

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

   輸出內容應如下所示：

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v5litepod-16             ACTIVE
   ...

   控制台

   1. 前往 Google Cloud 控制台的「TPUs」頁面：

      前往「TPU」

   2. 按一下「Queued resources」(已排入佇列的資源) 分頁標籤。

   3. 按一下佇列資源要求的名稱。

   如果 QR 狀態維持在 WAITING_FOR_RESOURCES 或 PROVISIONING 狀態超過 15 分鐘，請與 Google Cloud 帳戶代表聯絡。

7. 安裝依附元件：

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

8. 使用 multihost_runner.py 在每個工作站上執行 shardings.py。

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   記錄檔中會顯示每秒約 230 TFLOP 的效能。

   如要進一步瞭解如何設定平行處理作業，請參閱「使用 DCN 平行處理作業進行多切片分片」和 shardings.py。

9. 完成後，請清除 TPU 和排入佇列的資源。刪除作業需要 2 到 5 分鐘才能完成。如果您使用 gcloud CLI，可以透過選用的 --async 旗標，在背景執行這項指令。

將工作負載擴充至 Multislice

在多配量環境中執行模型前，請先對程式碼進行下列變更：

• 建立網格時，請使用 jax.experimental.mesh_utils.create_hybrid_device_mesh， 而非 jax.experimental.mesh_utils.create_device_mesh。

改用 Multislice 時，只需變更這些程式碼。為達到高效能，DCN 必須對應至資料平行、完全資料分割平行或管道平行軸。如要進一步瞭解成效考量和資料分割策略，請參閱「使用 Multislice 進行資料分割，盡可能提高成效」。

如要驗證程式碼是否可存取所有裝置，可以確認 len(jax.devices()) 等於 Multislice 環境中的晶片數量。舉例來說，如果您使用四個 v4-16 切片，每個切片有八個晶片，因此 len(jax.devices()) 應會傳回 32。

為 Multislice 環境選擇切片大小

如要線性加速，請新增與現有切片大小相同的切片。舉例來說，如果您使用 v4-512 切片，Multislice 會新增第二個 v4-512 切片，並將全域批次大小加倍，效能約可提升兩倍。詳情請參閱「Sharding With Multislice for Maximum Performance」。

在多個切片上執行工作

在 Multislice 環境中執行自訂工作負載有三種不同的方法：

1. 使用實驗執行工具指令碼 multihost_runner.py
2. 使用正式版執行器指令碼 multihost_job.py
3. 手動做法

實驗執行工具指令碼

multihost_runner.py 指令碼會將程式碼發布至現有的 Multislice 環境，並在每個主機上執行指令、將記錄複製回來，以及追蹤每個指令的錯誤狀態。multihost_runner.py 指令碼記錄在 MaxText README 中。

由於 multihost_runner.py 會維持永久 SSH 連線，因此僅適用於規模適中、執行時間相對較短的實驗。您可以根據工作負載和硬體設定，調整multihost_runner.py教學課程中的步驟。

正式版執行器指令碼

如要確保正式版作業在硬體故障和其他搶占情況下仍能正常運作，建議直接整合 Create Queued Resource API。請使用 multihost_job.py 做為工作範例，透過適當的開機指令碼觸發 Created Queued Resource API 呼叫，以便執行訓練並在先占時繼續執行。multihost_job.py 指令碼記載於 MaxText README。

由於 multihost_job.py 必須為每次執行佈建資源，因此疊代週期不如 multihost_runner.py 快速。

手動方法

建議您使用或調整 multihost_runner.py 或 multihost_job.py，在 Multislice 設定中執行自訂工作負載。不過，如果您偏好直接使用 QR 指令佈建及管理環境，請參閱「管理多重切片環境」。

管理 Multislice 環境

如要手動佈建及管理 QR，而不使用 MaxText 存放區中提供的工具，請參閱下列各節。

建立排入佇列的資源

gcloud

1. 使用下列指令建立排入佇列的資源要求：

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --node-count=${SLICE_COUNT} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --network=${NETWORK_NAME} \
       --subnetwork=${SUBNETWORK_NAME} \
       --tags=${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \
       --metadata=startup-script="${STARTUP_SCRIPT}" \
       [--reserved|--spot]

請務必確認配額充足，再選取 --reserved、--spot 或預設隨選配額。如要瞭解配額類型，請參閱「配額政策」。

curl

1. 建立名為 queued-resource-req.json 的檔案，然後將下列 JSON 複製到檔案中。

   {
   "guaranteed": { "reserved": true },
   "tpu": {
       "node_spec": [
       {
       "parent": "projects/your-project-number/locations/your-zone",
           "node": {
           "accelerator_type": "accelerator-type",
           "runtime_version": "tpu-vm-runtime-version",
           "network_config": {
               "network": "your-network-name",
               "subnetwork": "your-subnetwork-name",
               "enable_external_ips": true
           },
           "tags" : ["example-tag-1"]
           "metadata": {
               "startup-script": "your-startup-script"
           }
       },
       "multi_node_params": {
           "node_count": slice-count,
           "node_id_prefix": "your-queued-resource-id"
       }
       }
       ]
   }
   }

   替換下列值：

   • your-project-number：您的 Google Cloud 專案編號
   • your-zone：要建立佇列資源的可用區
   • accelerator-type：單一配量版本和大小。Multislice 僅支援 Cloud TPU v4 和更新的 TPU 版本。
   • tpu-vm-runtime-version：要使用的 TPU VM 執行階段版本。
   • your-network-name：選用，佇列資源要附加的網路
   • your-subnetwork-name：選用，要附加佇列資源的子網路
   • example-tag-1：選用，任意標記字串
   • your-startup-script：排入佇列的資源分配時執行的開機指令碼
   • slice-count：Multislice 環境中的 TPU 配量數量
   • your-queued-resource-id：使用者提供的已加入佇列資源 ID

   如要進一步瞭解所有可用選項，請參閱 REST Queued Resource API 說明文件。

   注意： 雖然 node_spec 是清單，但僅支援一個 node_spec 項目。

   如要使用 Spot 容量，請替換：

   "guaranteed": { "reserved": true }"spot": {}

   移除該行即可使用預設的隨選容量。

2. 使用 JSON 酬載提交已加入佇列的資源建立要求：

   $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
   -H "Content-Type: application/json" \
   -d @queuedresourcereq.json \
   https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

   替換下列值：

   • your-project-id：您的 Google Cloud 專案 ID
   • your-zone：要建立待處理資源的可用區
   • your-queued-resource-id：使用者提供的已加入佇列資源 ID

回覆應類似下方內容：

{
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}

使用 name 屬性字串值結尾的 GUID 值，取得已排入佇列的資源要求相關資訊。

控制台

1. 前往 Google Cloud 控制台的「TPUs」頁面：

   前往 TPU

2. 按一下「建立 TPU」。

3. 在「Name」(名稱) 欄位中，輸入 TPU 的名稱。

4. 在「Zone」(可用區) 方塊中，選取要建立 TPU 的可用區。

5. 在「TPU type」(TPU 類型) 方塊中，選取加速器類型。加速器類型會指定要建立的 Cloud TPU 版本和大小。多配量僅支援 Cloud TPU v4 和後續 TPU 版本。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。

6. 在「TPU 軟體版本」方塊中，選取軟體版本。建立 Cloud TPU VM 時，TPU 軟體版本會指定要安裝的 TPU 執行階段版本。詳情請參閱「TPU 軟體版本」。

7. 按一下「啟用佇列」切換鈕。

8. 在「Queued resource name」(待處理資源名稱) 欄位中，輸入待處理資源要求的名稱。

9. 按一下「將此設為多配量 TPU」核取方塊。

10. 在「Slice count」(切片數量) 欄位中，輸入要建立的切片數量。

11. 按一下「建立」，即可建立已加入佇列的資源要求。

擷取佇列中資源的狀態

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

如果佇列資源處於 ACTIVE 狀態，輸出內容會如下所示：

...
state:
    state: ACTIVE
...

curl

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

如果待處理資源處於 ACTIVE 狀態，輸出內容會如下所示：

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

控制台

1. 前往 Google Cloud 控制台的「TPUs」頁面：

   前往「TPU」頁面

2. 按一下「Queued resources」(已排入佇列的資源) 分頁標籤。

3. 按一下排入佇列的資源要求名稱。

佈建 TPU 後，您也可以前往「TPUs」(TPU) 頁面，找到 TPU 並點選對應的待處理資源要求名稱，查看待處理資源要求的詳細資料。

在極少數情況下，您可能會發現排入佇列的資源處於 FAILED 狀態，而部分切片則處於 ACTIVE 狀態。如果發生這種情況，請刪除已建立的資源，並在幾分鐘後重試，或聯絡Google Cloud 支援團隊。

透過 SSH 連線並安裝依附元件

在 TPU 節點上執行 JAX 程式碼： 說明如何使用單一節點中的 SSH 連線至 TPU VM。如要透過 SSH 連線至多切片環境中的所有 TPU VM 並安裝依附元件，請使用下列 gcloud 指令：

  $ gcloud compute tpus queued-resources ssh ${QR_ID} \
        --zone=${ZONE} \
        --node=all \
        --worker=all \
        --command="command-to-run" \
        --batch-size=4

這個 gcloud 指令會使用 SSH，將指定指令傳送至 QR 中的所有 worker 和節點。指令會分批傳送，每批四個，並同時傳送。目前批次完成執行後，系統就會傳送下一批指令。如果其中一個指令失敗，處理作業就會停止，系統也不會再傳送任何批次。詳情請參閱佇列資源 API 參考資料。如果使用的切片數量超過本機電腦的執行緒限制 (也稱為批次處理限制)，就會發生死結。舉例來說，假設本機電腦的批次處理限制為 64。如果嘗試在超過 64 個 slice (例如 100 個) 上執行訓練指令碼，SSH 指令會將 slice 分成多個批次。這會對第一批 64 個切片執行訓練指令碼，並等待指令碼完成，再對剩餘的 36 個切片執行指令碼。不過，前 64 個切片必須等待其餘 36 個切片開始執行指令碼，才能完成作業，因此會造成死結。

為避免這種情況，您可以在指令碼指令 (使用 --command 旗標指定) 後方加上連字號 (&)，在各個 VM 的背景執行訓練指令碼。這麼做時，在第一個切片批次上啟動訓練指令碼後，控制項會立即返回 SSH 指令。接著，SSH 指令可以開始在剩餘的 36 個 slice 批次上執行訓練指令碼。在背景執行指令時，您需要適當管道傳送 stdout 和 stderr 串流。如要提高同一 QR 中的平行度，可以使用 --node 參數選取特定切片。

網路設定

請執行下列步驟，確保 TPU 節點可以互相通訊。 在每個切片上安裝 JAX。詳情請參閱「在 TPU 切片上執行 JAX 程式碼」。確認 len(jax.devices()) 等於 Multislice 環境中的晶片數量。如要執行這項操作，請在每個 slice 上執行以下指令：

  $ python3 -c 'import jax; print(jax.devices())'

如果您在四個 v4-16 切片上執行這段程式碼，每個切片有八個晶片，四個切片總共 32 個晶片 (裝置)，應由 jax.devices() 傳回。

列出排入佇列的資源

gcloud

您可以使用 queued-resources list 指令查看佇列資源的狀態：

$ gcloud compute tpus queued-resources list --zone=${ZONE}

輸出看起來類似以下內容：

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

控制台

1. 前往 Google Cloud 控制台的「TPUs」頁面：

   前往「TPU」

2. 按一下「已加入佇列的資源」分頁標籤。

在已佈建的環境中啟動工作

您可以透過 SSH 連線至每個切片中的所有主機，並在所有主機上執行下列指令，手動執行工作負載。

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

重設 QR code

您可以使用 ResetQueuedResource API 重設 ACTIVE QR 中的所有 VM。重設 VM 會強制清除機器的記憶體，並將 VM 重設為初始狀態。本機儲存的資料不會受到影響，重設後系統會叫用啟動指令碼。如要重新啟動所有 TPU，ResetQueuedResource API 相當實用。例如，當訓練卡住時，重設所有 VM 會比偵錯更容易。

所有 VM 的重設作業會同時執行，ResetQueuedResource 作業需要一到兩分鐘才能完成。如要叫用 API，請使用下列指令：

$ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

正在刪除排入佇列的資源

如要在訓練課程結束時釋出資源，請刪除已加入佇列的資源。刪除作業會在兩到五分鐘內完成。如果您使用 gcloud CLI，可以搭配選用的 --async 旗標，在背景執行這項指令。

gcloud

$ gcloud compute tpus queued-resources \
    delete ${QR_ID} --zone=${ZONE} --force [--async]

控制台

1. 前往 Google Cloud 控制台的「TPUs」頁面：

   前往「TPU」

2. 按一下「Queued resources」(已排入佇列的資源) 分頁標籤。

3. 找出要取消的排入佇列資源要求，然後勾選旁邊的核取方塊。

4. 按一下「刪除」圖示 。

自動復原故障

如果發生中斷情形，Multislice 會自動修復受影響的切片，並在之後重設所有切片。受影響的切片會由新切片取代，其餘正常的切片則會重設。如果沒有可用運算資源來分配替代切片，訓練就會停止。

如要在中斷後自動繼續訓練，您必須指定開機指令碼，檢查並載入上次儲存的檢查點。每當重新分配切片或重設 VM 時，系統都會自動執行開機指令碼。您可以在傳送至建立 QR 要求 API 的 JSON 酬載中，指定開機指令碼。

下列開機指令碼 (用於「建立 QR」) 可讓您在 MaxText 訓練期間，從儲存在 Cloud Storage bucket 的檢查點自動復原失敗狀態，並繼續訓練：

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 -m MaxText.train MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

請先複製 MaxText 存放區，再嘗試這項操作。

剖析和偵錯

單一切片和多切片環境的剖析作業相同。詳情請參閱「剖析 JAX 程式」。

取得最佳培訓成效

以下各節說明如何最佳化 Multislice 訓練。

使用 Multislice 分片，盡可能提升效能

如要在多配量環境中達到最高效能，必須考量如何跨多個配量進行資料分割。通常有三種選擇 (資料平行處理、完全分區資料平行處理和管道平行處理)。我們不建議跨模型維度進行啟用分片 (有時稱為張量平行處理)，因為這需要過多的切片間頻寬。對於所有這些策略，您可以在過去適用於您的切片中保留相同的分片策略。

建議您先從純資料平行處理開始。使用完全分區資料平行處理，有助於釋出記憶體用量。缺點是配量之間的通訊會使用 DCN 網路，導致工作負載變慢。只有在根據批次大小 (如下方分析) 判斷有必要時，才使用管道平行處理。

資料平行處理的使用時機

如果工作負載運作良好，但您想透過多個切片擴充來提升效能，純資料平行處理會是理想的做法。

如要在多個切片中實現強大的擴充功能，透過 DCN 執行 all-reduce 作業所需的時間，必須少於執行反向傳遞作業所需的時間。DCN 用於切片之間的通訊，是工作負載輸送量的限制因素。

每顆 v4 TPU 晶片的尖峰效能為每秒 275 * 10¹² FLOPS。

每個 TPU 主機有四個晶片，每個主機的網路頻寬上限為 50 Gbps。

也就是說，算術強度為 4 * 275 * 10¹² FLOPS / 50 Gbps = 22000 FLOPS / bit。

模型會為每個步驟的每個參數使用 32 到 64 位元的 DCN 頻寬。如果使用兩個配量，模型會使用 32 位元的 DCN 頻寬。如果使用超過兩個切片，編譯器會執行完整的隨機排序 all-reduce 作業，每個步驟的每個參數最多會使用 64 位元的 DCN 頻寬。每個參數所需的 FLOPS 數量會因模型而異。具體來說，對於以 Transformer 為基礎的語言模型，正向和反向傳遞所需的 FLOPS 數量約為 6 * B * P，其中：

• B 是以權杖為單位的批次大小
• P 是參數數量

每個參數的 FLOPS 數量為 6 * B，反向傳遞期間每個參數的 FLOPS 數量為 4 * B。

為確保多個切片之間的強大擴充能力，請確保運算強度超過 TPU 硬體的算術強度。如要計算運算強度，請將反向傳遞期間每個參數的 FLOPS 數量，除以每個步驟中每個參數的網路頻寬 (以位元為單位)：Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

因此，如果使用兩個切片，以 Transformer 為基礎的語言模型會是： Operational intensity = 4 * B / 32

如果使用超過兩個切片：Operational intensity = 4 * B/64

這表示 Transformer 語言模型的最小批次大小介於 176k 和 352k 之間。由於 DCN 網路可能會短暫捨棄封包，因此最好保留大量錯誤空間，只有在每個 Pod 的批次大小至少為 35 萬 (兩個 Pod) 至 70 萬 (多個 Pod) 時，才部署資料平行處理。

如果是其他模型架構，您需要估算每個切片的向後傳遞執行時間 (可使用剖析器計時或計算 FLOPS)。接著，您可以將這個結果與 DCN 上 all-reduce 的預期執行時間進行比較，並準確估算資料平行處理是否適合您。

使用完全分區資料平行處理 (FSDP) 的時機

完全分片資料平行處理 (FSDP) 會結合資料平行處理 (將資料分片到各個節點) 和權重分片 (將權重分片到各個節點)。在正向和反向傳遞的每項作業中，權重都會全部收集，因此每個切片都有所需的權重。梯度不會使用 all-reduce 同步處理，而是會在產生時 reduce-scattered。這樣一來，每個切片只會取得負責權重的梯度。

與資料平行處理類似，FSDP 需要根據切片數量，線性調整全域批次大小。隨著切片數量增加，FSDP 會降低記憶體壓力。這是因為每個切片的權重和最佳化工具狀態數量會減少，但代價是網路流量增加，且集體延遲導致封鎖的可能性更高。

在實務上，如果您要增加每個切片的批次大小、儲存更多啟用項目，以盡量減少反向傳遞期間的重新具體化作業，或是增加神經網路中的參數數量，最好使用跨切片的 FSDP。

FSDP 中的全體收集和全體縮減作業與 DP 中的作業類似，因此您可以按照上一節所述方式，判斷 FSDP 工作負載是否受 DCN 效能限制。

管道平行處理的使用時機

如果其他平行處理策略需要全域批量大於偏好的最大批量，才能達到高效能，則管道平行處理會變得相關。管道平行處理可讓組成管道的切片「共用」批次。不過，管道平行處理有兩個重大缺點：

1. 這會導致「pipeline bubble」，也就是晶片因等待資料而閒置。
2. 這項技術需要微批量，會降低有效批量、算術強度，最終降低模型 FLOP 使用率。

只有在其他平行處理策略需要過大的全域批次大小時，才應使用管道平行處理。在嘗試管道平行處理之前，建議先進行實驗，看看在達到高效能 FSDP 所需的批次大小時，每個樣本的收斂速度是否會變慢。FSDP 通常能提高模型 FLOP 使用率，但如果每個樣本的收斂速度隨著批次大小增加而變慢，管道平行化可能仍是較好的選擇。大多數工作負載都能容忍足夠大的批次大小，因此無法從管道平行處理中受益，但您的工作負載可能不同。

如有必要，建議您將管道平行處理與資料平行處理或 FSDP 結合使用。這樣一來，您就能盡量縮短管道深度，同時增加每個管道的批次大小，直到 DCN 延遲時間對處理量的影響降低為止。具體來說，如果您有 N 個切片，請考慮深度為 2 的管道和 N/2 個資料平行處理副本，然後考慮深度為 4 的管道和 N/4 個資料平行處理副本，以此類推，直到每個管道的批次大小足以讓 DCN 集合隱藏在反向傳遞的算術運算之後。這樣可以盡量減少管道平行處理造成的速度減慢，同時突破全域批次大小限制。

多配量最佳做法

以下各節說明 Multislice 訓練的最佳做法。

資料載入

在訓練期間，我們會從資料集重複載入批次資料，並提供給模型。為了避免 TPU 閒置，請務必使用有效率的非同步資料載入器，將批次資料分片到各主機。MaxText 目前的資料載入器會讓每個主機載入相同數量的樣本。這個解決方案適用於文字，但需要在模型內重新分 shard。此外，MaxText 尚未提供確定性快照功能，無法讓資料疊代器在先占前後載入相同資料。

查核點機制

Orbax 檢查點程式庫提供檢查點 JAX PyTrees 的原始物件，可儲存至本機儲存空間或 Google Cloud 儲存空間。我們提供參考整合，將同步檢查點設定到 MaxText 中，請參閱 checkpointing.py。

支援的設定

以下各節說明 Multislice 支援的切片形狀、自動化調度管理、架構和平行處理。

形狀

所有切片都必須具有相同的形狀 (例如相同的 AcceleratorType)。不支援異質切片形狀。

自動化調度管理

GKE 支援自動調度管理。詳情請參閱「GKE 中的 TPU」。

架構

Multislice 僅支援 JAX 和 PyTorch 工作負載。

平行處理工作數量

建議使用者搭配資料平行處理功能測試 Multislice。如要進一步瞭解如何使用 Multislice 導入管線平行處理，請洽詢Google Cloud 帳戶代表。

支援和意見回饋

歡迎提供意見！如要分享意見回饋或要求支援，請使用 Cloud TPU 支援或意見回饋表單與我們聯絡。