使用 TPU 執行管道

本頁說明如何在 Dataflow 上使用 TPU 執行 Apache Beam 管道。使用 TPU 的工作會產生費用，詳情請參閱 Dataflow 定價頁面。

如要進一步瞭解如何搭配使用 TPU 和 Dataflow，請參閱「Dataflow 支援 TPU」。

選用：預留特定加速器

雖然您可以隨選使用 TPU，但我們強烈建議您搭配專門指定的 Google Cloud預訂使用 Dataflow TPU。 這有助於確保您能存取可用的加速器，並縮短工作人員啟動時間。使用 TPU 預留項目的管道不需要額外的 TPU 配額。

如果您未預訂，而是選擇隨選使用 TPU，請在執行管道前佈建 TPU 配額。

選用：佈建 TPU 配額

您可以透過預留項目或隨選容量使用 TPU。如要隨選使用 TPU，必須先佈建 TPU 配額。如果使用特定目標預訂，則可略過本節。

如要隨需使用 TPU，而不需預訂，請按照下列步驟檢查 Compute Engine API 的 TPU 配額限制和目前用量：

主控台

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

   前往配額頁面

2. 在「篩選器」方塊中，執行下列操作：

   1. 請參閱下表，根據 TPU 版本和機器類型選取並複製配額的屬性。舉例來說，如果您打算建立機器類型開頭為 ct5lp- 的隨選 TPU v5e 節點，請輸入 Name: TPU v5 Lite PodSlice chips。

      TPU 版本，機器類型開頭為	隨選執行個體配額的屬性和名稱
      TPU v5e、 ct5lp-	Name: TPU v5 Lite PodSlice chips
      TPU v5p、 ct5p-	Name: TPU v5p chips
      TPU v6e、 ct6e-	Dimensions (e.g. location): tpu_family:CT6E

   2. 選取「Dimensions (e.g. location)」(維度，例如位置) 屬性，然後輸入 region:，並在後面加上您打算啟動管道的區域名稱。舉例來說，如果您打算使用 us-west4-a 區域，請輸入 region:us-west4。TPU 配額是區域配額，因此同一區域內的所有可用區都會共用相同的 TPU 配額。

設定自訂容器映像檔

如要在 Dataflow 管道中與 TPU 互動，您必須在管道執行階段環境中提供可在 XLA 裝置上運作的軟體。您必須根據管道需求安裝 TPU 程式庫，並根據使用的 TPU 裝置設定環境變數。

如要自訂容器映像檔，請將 Apache Beam 安裝到具有必要 TPU 程式庫的現成基本映像檔中。或者，您也可以將 TPU 軟體安裝到發布的映像檔，並搭配 Apache Beam SDK 版本使用。

如要提供自訂容器映像檔，請使用 sdk_container_image 管道選項。詳情請參閱「在 Dataflow 中使用自訂容器」。

使用 TPU 加速器時，您需要在容器映像檔中設定下列環境變數。

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

視您使用的加速器而定，下表中的變數也必須設定。

類型	拓撲	必要 Dataflow worker_machine_type	其他環境變數
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	TPU_ACCELERATOR_TYPE=v5litepod-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	TPU_ACCELERATOR_TYPE=v5litepod-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	TPU_ACCELERATOR_TYPE=v5litepod-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice	1x1	ct6e-standard-1t	TPU_ACCELERATOR_TYPE=v6e-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice	2x2	ct6e-standard-4t	TPU_ACCELERATOR_TYPE=v6e-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice	2x4	ct6e-standard-8t	TPU_ACCELERATOR_TYPE=v6e-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	TPU_ACCELERATOR_TYPE=v5p-8 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

自訂容器映像檔的 Dockerfile 範例如下：

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

使用 TPU 執行工作

使用 TPU 執行 Dataflow 工作時，請注意下列事項：

• 由於 TPU 容器可能很大，為避免磁碟空間不足，請使用 --disk_size_gb 管道選項，將預設開機磁碟大小增加至 50 GB，或容器映像檔所需的適當大小。
• 限制工作站內平行處理。

TPU 和工作站平行處理

在預設設定中，Dataflow Python 管道會為每個 VM 核心啟動一個 Apache Beam SDK 程序。TPU 機器類型具有大量 vCPU 核心，但只有一個程序可以在 TPU 裝置上執行運算。此外，程序可能會在整個生命週期內預留 TPU 裝置。因此，執行 Dataflow TPU 管道時，您必須限制工作站內的平行處理量。如要限制工作站的平行處理量，請按照下列指引操作：

• 如果您的用途是針對模型執行推論，請使用 Beam RunInference API。詳情請參閱「Beam 中的大型語言模型推論」。
• 如果無法使用 Beam RunInference API，請使用 Beam 的多程序共用物件，將特定作業限制為單一程序。
• 如果無法使用上述建議，且偏好每個工作站只啟動一個 Python 程序，請設定 --experiments=no_use_multiple_sdk_containers 管道選項。
• 如果工作站的 vCPU 數量超過 100 個，請使用 --number_of_worker_harness_threads 管道選項減少執行緒數量。請參閱下表，確認您的 TPU 類型是否使用超過 100 個 vCPU。

下表列出每種 TPU 設定的每個工作站總運算資源。

TPU 類型	拓撲	機器類型	TPU 晶片	vCPU	RAM (GB)
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	1	24	48
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	4	112	192
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	8	224	384
tpu-v6e-slice	1x1	ct6e-standard-1t	1	44	176
tpu-v6e-slice	2x2	ct6e-standard-4t	4	180	720
tpu-v6e-slice	2x4	ct6e-standard-8t	8	360	1440
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	4	208	448

使用 TPU 執行管道

如要使用 TPU 執行 Dataflow 工作，請使用下列指令。

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

更改下列內容：

• PIPELINE：管道的原始碼檔案。
• PROJECT： Google Cloud 專案名稱。
• BUCKET：Cloud Storage bucket。
• ：Dataflow 區域，例如 us-central1。REGION
• TPU_TYPE：支援的 TPU 類型，例如 tpu-v5-lite-podslice。如需完整清單，請參閱「支援的 TPU 加速器」。
• TPU_TOPOLOGY：TPU 拓撲，例如 1x1。
• MACHINE_TYPE：對應的機器類型，例如 ct5lp-hightpu-1t。
• DISK_SIZE_GB：每個工作站 VM 的開機磁碟大小，例如 100。
• IMAGE：Docker 映像檔的 Artifact Registry 路徑。
• NUMBER_OF_THREADS：選用。工作站輔助執行緒的數量。

驗證 Dataflow 工作

如要確認工作是否使用搭載 TPU 的工作站 VM，請按照下列步驟操作：

1. 在 Google Cloud 控制台中，依序前往「Dataflow」>「Jobs」(工作) 頁面。

   前往「Jobs」(工作) 頁面

2. 選取職務。

3. 按一下「工作指標」分頁標籤。

4. 在「Autoscaling」(自動調度資源) 區段中，確認至少有一個「Current workers」(目前的工作人員) VM。

5. 在側邊的「Job info」(工作資訊) 窗格中，確認 machine_type 以 ct 開頭。例如：ct6e-standard-1t。這表示 TPU 用量。

排解 Dataflow 工作問題

如果使用 TPU 執行 Dataflow 工作時發生問題，請參閱排解 Dataflow TPU 工作問題。

後續步驟

• 進一步瞭解 Dataflow 的 TPU 支援。
• 瞭解 Beam 中的大型模型推論。