使用 TPU 运行流水线

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

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

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

虽然您可以按需使用 TPU，但我们强烈建议您将 Dataflow TPU 与特定目标的 Google Cloud Platform 预留搭配使用。这有助于确保您可以使用可用的加速器，并缩短工作器启动时间。使用 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 的多进程共享对象将某些操作限制为单个进程。
• 如果您无法采纳上述建议，并且希望每个工作器只启动一个 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

使用 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 作业问题。

后续步骤

• 不妨试用快速入门示例：在 TPU 上运行 Dataflow。
• 详细了解 Dataflow 上的 TPU 支持。
• 了解 Beam 中的大型模型推理。