Cloud TPU 性能指南

排查 TPU 性能问题时，第一步是分析模型。如需详细了解如何捕获性能配置文件，请参阅在 Cloud TPU 上剖析模型。

TPU 模型性能

本部分介绍可能会降低模型性能的常见问题，以及如何解决这些问题。

  1. 模型受限于输入

    TPU 执行计算的速度非常快。为了确保 TPU 不处于空闲状态，请务必确保将稳定的数据流加载到 TPU。实现这一点的方式取决于您加载和预处理数据集的方式。例如，您可以使用 tf.data.TFRecordset() 和 num_parallel_reads 参数并行读取数据文件。

  2. 由于分片，批次大小太小（跨核心拆分批次）

    TPU 运行时在 TPU 设备（例如 v2-8 或 v3-8）的所有 8 个核心上拆分批量。如果指定的全局批量大小为 128，则每个核心会收到批量大小 16 (128 / 8)。

    为获得最佳内存用量，请使用 TPU 内存中可容纳的最大批量大小。每个 TPU 核心都使用二维 8 X 128 向量寄存器来处理矩阵乘法。通常，您的批次大小应该能被 8 或 128 整除。

  3. 内存管理调优

    您可以使用 TPU_PREMAPPED_BUFFER_SIZE 环境变量来微调低级运行时行为。

  • 说明： TPU_PREMAPPED_BUFFER_SIZE 用于设置预映射并固定供 TPU 运行时用于数据传输（例如 DMA）的主机内存缓冲区的大小（以字节为单位）。默认值为 4294967296 字节。 该值必须是 2^12（4KB = 4 * 1024 字节 = 4096 = 2^12）的倍数。

    以下示例展示了有效的 TPU_PRE_MAPPED_BUFFER_SIZE 值。

        17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
    40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).

  • 影响：增大此大小可能会提高主机与 TPU 设备之间的数据传输性能，尤其对于具有大型张量或频繁进行主机-设备通信的工作负载。不过，它也会增加固定主机内存量，从而减少可供其他进程使用的内存。

    缓冲大小

    如果预映射的缓冲区区域不够大，无法在程序运行时分配内存，工作负载将失败并返回类似于以下的 RESOURCE_EXHAUSTED 错误：

    “从预映射区域分配缓冲区失败，错误为：RESOURCE_EXHAUSTED。 尝试分配 allocation_size。这是不可能的。可用大小为 available_size。”

    如果缓冲区过大，TPU 初始化可能需要更长时间（可能超过 15 秒），这会让人感觉 TPU 卡住了。

    如需诊断此问题，请检查 TPU 运行时日志。这些日志详细记录了正在执行的操作，包括缓冲区的预映射。您可以在 /tmp/tpu_logs/tpu_driver.INFO 中找到日志，也可以通过设置环境变量 TPU_STDERR_LOG_LEVEL=0 将日志直接输出到控制台。此设置将生成类似于以下内容的输出：

     I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
 I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
 I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
 I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
 ```

This output will tell you how long it took to initialize the TPU and
the size of the premapped buffer.

  • 用法：如果预映射的缓冲区过小或过大，您可以使用以下环境变量手动设置缓冲区大小。

    TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
pre-mapped buffer region.
TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
a single buffer that can be allocated from the pre-mapped region.

    例如，您可以执行以下操作：

     export TPU_PREMAPPED_BUFFER_SIZE=4294967296

    设置缓冲区大小，并：

     export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
 ```
 to enable it.

 This export sets the size to the default.

  • 指南：如果您怀疑主机-设备数据传输是瓶颈，请调整 TPU_PREMAPPED_BUFFER_SIZE 的值。监控主机内存使用情况和模型性能，以找到最佳平衡点。对于大多数使用情形，默认值通常就足够了。

XLA 编译器优化

XLA 是一款机器学习编译器，可以为 TPU、CPU、GPU 和其他平台生成二进制文件。虽然 XLA 是标准 TensorFlow 代码库的一部分，但它也可用于 PyTorch 和 JAX 模型。适用于 Cloud TPU 的模型被转换为 XLA 图，然后 XLA 将其编译为 TPU 可执行文件。如需详细了解 XLA，请参阅 XLA：优化机器学习编译器。

内边距

为了高效地使用 TPU 内存，请设计数据结构以将其平铺为 128 x 8 个数据块。如果矩阵计算的数据填满了 128 x 8 个数据块，则 XLA 编译器会填充张量。填充有两个不足：

  1. 填充的张量不能充分利用 TPU 核心。
  2. 填充增加了张量所需的片上内存存储，并可能导致内存不足错误。

虽然填充是由 XLA 编译器在必要时自动执行的，但您可以使用内存查看器工具来确定填充量。您可以通过选择非常适合 TPU 的张量维度来避免填充。

张量维度

为了达到峰值 FLOP，矩阵乘法维度应大于您使用的 TPU 版本的 MXU 大小。对于 v6e，MXU 大小为 256 x 256；对于 v6e 之前的版本，MXU 大小为 128 x 128。如需了解详情，请参阅 Cloud TPU 系统架构。

批次大小

XLA 编译器将存储在 TPU HBM 内存中的张量的大小向上舍入，以便更高效地执行计算。此填充操作在硬件级别以透明方式进行，不会影响结果。但是，在某些情况下，填充可能会导致内存使用量和执行时间显著增加。

TPU 运行时会在内存中存放张量，以便最大限度地提高计算效率并减少填充。为了最大限度地减少内存开销并提高计算效率，必须满足以下条件之一：

  1. 总批量大小应为 64 的倍数（每个 TPU 核心 8 个），并且特征维度大小应为 128 的倍数。

  2. 总批量大小应为 1024 的倍数（每个 TPU 核心 128 个），并且特征维度大小应为 8 的倍数。

使用 1024 作为批量大小以及 128 的倍数作为特征维度可以获得最佳效率，但这并非适用于所有模型。

融合

融合是 XLA 编译器用于优化程序的通用技术。融合操作是指合并多个需要联合执行的操作。

例如，请考虑以下系列操作：

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

此代码大致相当于以下伪代码：

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result[i] = tmp[i] * z[i];
    }

通过融合，数组访问将同时进行：

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

在此例中，内存往返次数减少，XLA 无需为“tmp”分配任何空间。

融合是一项关键优化措施，可从多方面对 Cloud TPU 形成有利影响：

  • 无需在主内存中存储中间结果（通常很慢），可减少内存传输。
  • 可以更充分地利用硬件设备，以免这些设备被闲置。
  • 由于同时需要的活动缓冲区减少，可以降低模型的内存利用量。

广播

当两个形状不同但可兼容的张量合并时，系统将以隐式方式进行广播。

例如，tf.add(vector, matrix) 要求将矢量广播到矩阵的形状。操作结果具有与矩阵相同的形状。如需了解详情，请参阅广播数组指南。

虽然广播通常可与其使用方融合，但强制执行广播可能会导致性能低下，同时增加内存使用量。

在以下示例中，矢量与矩阵的相加操作中隐含的广播不能与 argmax 融合，因而导致广播物化：

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`