使用 TPU 运行流水线

本页面介绍了如何在具有 TPU 的 Dataflow 上运行 Apache Beam 流水线。使用 TPU 的作业按 Dataflow 价格页面中指定的费用收费。

如需详细了解如何将 TPU 与 Dataflow 搭配使用，请参阅 Dataflow 对 TPU 的支持。

可选：进行特定预留以使用加速器

虽然您可以按需使用 TPU，但我们强烈建议您使用 Dataflow TPU 并专门针对 Google Cloud预留。这有助于确保您可以使用可用的加速器，并缩短工作器启动时间。使用 TPU 预留的流水线不需要额外的 TPU 配额。

如果您不进行预留，而是选择按需使用 TPU，请在运行流水线之前预配 TPU 配额。

可选：预配 TPU 配额

您可以使用按需容量的 TPU，也可以使用预留的 TPU。如果您想按需使用 TPU，必须先预配 TPU 配额。如果您使用专门针对的预留，则可以跳过此部分。

如需在不预留的情况下按需使用 TPU，请按以下步骤查看 TPU 的 Compute Engine API 配额的限制和当前用量：

控制台

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. 选择维度（例如位置）属性，然后输入 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 流水线会为每个虚拟机核心启动一个 Apache Beam SDK 进程。TPU 机器类型具有大量 vCPU 核心，但只有一个进程可以在 TPU 设备上执行计算。此外，TPU 设备可能会被某个进程在整个进程生命周期内预留。因此，在运行 Dataflow TPU 流水线时，您必须限制工作器内并行性。如需限制工作器的并行性，请遵循以下指南：

• 如果您的应用场景涉及对模型运行推理，请使用 Beam RunInference API。如需了解详情，请参阅 Beam 中的大型语言模型推理。
• 如果您无法使用 Beam RunInference API，请使用 Beam 的多进程共享对象将某些操作限制为单个进程。
• 如果您无法使用上述建议，并且希望每个 worker 只启动一个 Python 进程，请设置 --experiments=no_use_multiple_sdk_containers 流水线选项。
• 对于具有 100 多个 vCPU 的工作器，请使用 --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 存储桶。
• REGION：Dataflow 区域，例如 us-central1。
• TPU_TYPE：受支持的 TPU 类型，例如 tpu-v5-lite-podslice。如需查看类型和拓扑的完整列表，请参阅支持的 TPU 加速器。
• TPU_TOPOLOGY：TPU 拓扑，例如 1x1。
• MACHINE_TYPE：相应的机器类型，例如 ct5lp-hightpu-1t。
• DISK_SIZE_GB：每个工作器虚拟机的启动磁盘大小，例如 100。
• IMAGE：Docker 映像的 Artifact Registry 路径。
• NUMBER_OF_THREADS：可选。工作器框架线程数。

验证 Dataflow 作业

如需确认作业使用具有 TPU 的工作器虚拟机，请按以下步骤操作：

1. 在 Google Cloud 控制台中，依次前往 Dataflow > 作业页面。

   转到作业

2. 选择一个作业。

3. 点击作业指标标签页。

4. 在自动扩缩部分中，确认至少有一个当前工作器虚拟机。

5. 在侧边作业信息窗格中，检查 machine_type 是否以 ct 开头。例如 ct6e-standard-1t。这表示 TPU 用量。

排查 Dataflow 作业问题

如果您在使用 TPU 运行 Dataflow 作业时遇到问题，请参阅排查 Dataflow TPU 作业问题。

后续步骤

• 详细了解 Dataflow 上的 TPU 支持。
• 了解 Beam 中的大型模型推理。