Cloud TPU 性能指南
排查 TPU 性能问题时，第一步是分析模型。如需详细了解如何捕获性能配置文件，请参阅在 Cloud TPU 上剖析模型。
TPU 模型性能
本部分介绍可能会降低模型性能的常见问题，以及如何解决这些问题。
模型受限于输入
TPU 执行计算的速度非常快。为了确保 TPU 不处于空闲状态，请务必确保将稳定的数据流加载到 TPU。实现这一点的方式取决于您加载和预处理数据集的方式。例如，您可以使用 tf.data.TFRecordset() 和
num_parallel_reads参数并行读取数据文件。由于分片，批次大小太小（跨核心拆分批次）
TPU 运行时在 TPU 设备（例如 v2-8 或 v3-8）的所有 8 个核心上拆分批量。如果指定的全局批量大小为 128，则每个核心会收到批量大小 16 (128 / 8)。
为获得最佳内存用量，请使用 TPU 内存中可容纳的最大批量大小。每个 TPU 核心都使用二维 8 X 128 向量寄存器来处理矩阵乘法。通常，您的批次大小应该能被 8 或 128 整除。
内存管理调优
您可以使用
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 编译器会填充张量。填充有两个不足：
- 填充的张量不能充分利用 TPU 核心。
- 填充增加了张量所需的片上内存存储，并可能导致内存不足错误。
虽然填充是由 XLA 编译器在必要时自动执行的，但您可以使用内存查看器工具来确定填充量。您可以通过选择非常适合 TPU 的张量维度来避免填充。
张量维度
为了达到峰值 FLOP，矩阵乘法维度应大于您使用的 TPU 版本的 MXU 大小。对于 v6e，MXU 大小为 256 x 256；对于 v6e 之前的版本，MXU 大小为 128 x 128。如需了解详情，请参阅 Cloud TPU 系统架构。
批次大小
XLA 编译器将存储在 TPU HBM 内存中的张量的大小向上舍入，以便更高效地执行计算。此填充操作在硬件级别以透明方式进行，不会影响结果。但是，在某些情况下，填充可能会导致内存使用量和执行时间显著增加。
TPU 运行时会在内存中存放张量，以便最大限度地提高计算效率并减少填充。为了最大限度地减少内存开销并提高计算效率，必须满足以下条件之一：
总批量大小应为 64 的倍数（每个 TPU 核心 8 个），并且特征维度大小应为 128 的倍数。
总批量大小应为 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)`