本頁面由 Cloud Translation API 翻譯而成。

使用 TPU 執行管道

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

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

選用：預留特定加速器

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

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

選用：佈建 TPU 配額

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

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

控制台

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

前往配額頁面

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

請參閱下表，根據 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

選取「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 管道選項。
如果使用 --number_of_worker_harness_threads 管道選項可提升效能，您就能進一步減少執行緒數量。

下表列出每種 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