使用 TensorFlow 在 TPU 上进行训练¶
如果你不想阅读长篇大论，只想直接查看 TPU 的代码示例，可以查看我们的 TPU 示例笔记本！
什么是 TPU？¶
TPU 是 张量处理单元（Tensor Processing Unit）。它们是由谷歌设计的硬件，用于大幅加速神经网络中的张量计算，类似于 GPU。TPU 可以用于网络训练和推理。它们通常通过谷歌的云服务访问，但在 Google Colab 和 Kaggle Kernels 中也可以免费访问小型 TPU。
因为 所有 🤗 Transformers 中的 TensorFlow 模型都是 Keras 模型，本文档中的大多数方法也适用于任何 Keras 模型的 TPU 训练！不过，有些内容是特定于 HuggingFace 的 Transformers 和 Datasets 生态系统的，我们会在需要时特别指出这些内容。
可以使用哪些类型的 TPU？¶
新用户经常对不同类型的 TPU 和访问方式感到困惑。首先需要了解的关键区别是 TPU 节点（TPU Node） 和 TPU 虚拟机（TPU VM） 之间的区别。
当你使用 TPU 节点 时，实际上是间接访问远程 TPU。你需要一个单独的虚拟机来初始化网络和数据管道，然后将它们转发到远程节点。在 Google Colab 中使用 TPU 时，实际上是在使用 TPU 节点。
使用 TPU 节点可能会有一些令人意外的行为，特别是对于不熟悉它的人！具体来说，由于 TPU 位于与运行 Python 代码的机器物理上不同的系统中，因此你的数据不能保存在本地机器上——从本地机器内部存储加载数据的数据管道将完全失败！相反，数据必须存储在 Google Cloud Storage 中，以便数据管道在远程 TPU 节点上运行时仍能访问这些数据。
如果你可以将所有数据存储在内存中作为 np.ndarray 或 tf.Tensor，则可以在使用 Colab 或 TPU 节点时直接调用 fit() 而不需要上传数据到 Google Cloud Storage。
🤗 Hugging Face 小贴士 🤗： 方法 Dataset.to_tf_dataset() 及其高级包装器 model.prepare_tf_dataset() 在 TPU 节点上会失败。原因是虽然它们创建了一个 tf.data.Dataset，但它不是一个“纯” tf.data 管道，而是使用 tf.numpy_function 或 Dataset.from_generator() 从底层的 HuggingFace Dataset 流式传输数据。这个 HuggingFace Dataset 存储在本地磁盘上，远程 TPU 节点无法读取。
第二种访问 TPU 的方式是通过 TPU 虚拟机。使用 TPU 虚拟机时，你可以直接连接到 TPU 附带的机器，就像在 GPU 虚拟机上训练一样。TPU 虚拟机通常更容易使用，尤其是在处理数据管道方面。上述所有警告都不适用于 TPU 虚拟机！
这是一个有偏见的文档，所以这里是我们的一点意见：如果可能，尽量避免使用 TPU 节点。 它比 TPU 虚拟机更令人困惑且更难调试。未来 TPU 节点也很可能不再被支持——谷歌最新的 TPU，TPUv4，只能作为 TPU 虚拟机访问，这表明 TPU 节点越来越成为一个“遗留”访问方法。然而，我们理解唯一免费的 TPU 访问是在 Colab 和 Kaggle Kernels 上，这些平台使用 TPU 节点——所以我们将会解释如何处理这个问题！请查看 TPU 示例笔记本 获取更详细的代码示例。
有哪些大小的 TPU 可供使用？¶
单个 TPU（如 v2-8/v3-8/v4-8）运行 8 个副本。TPU 以 吊舱（pod） 形式存在，可以同时运行数百或数千个副本。当使用的 TPU 数量超过单个 TPU 但少于整个吊舱时（例如 v3-32），你的 TPU 集群称为 吊舱切片（pod slice）。
通过 Colab 免费访问 TPU 时，通常会获得一个单个 v2-8 TPU。
我经常听说 XLA 这个东西。什么是 XLA？它与 TPU 有何关系？¶
XLA 是一种优化编译器，由 TensorFlow 和 JAX 使用。在 JAX 中它是唯一的编译器，而在 TensorFlow 中是可选的（但在 TPU 上是强制的）。最简单的方式来启用 XLA 是在调用 model.compile() 时传递参数 jit_compile=True。如果你没有遇到错误且性能良好，那就说明你已经准备好在 TPU 上运行了！
在 TPU 上调试通常比在 CPU/GPU 上更困难，因此我们建议先在 CPU/GPU 上使用 XLA 运行代码，然后再尝试在 TPU 上运行。当然，你不必训练很长时间——只需跑几步以确保模型和数据管道按预期工作即可。
XLA 编译后的代码通常更快——即使你不打算在 TPU 上运行，添加 jit_compile=True 也可以提高性能。不过，请注意以下关于 XLA 兼容性的注意事项！
经验之谈： 尽管使用 jit_compile=True 是提高速度并测试 CPU/GPU 代码是否 XLA 兼容的好方法，但实际在 TPU 上训练时将其保留可能会导致很多问题。XLA 编译将在 TPU 上隐式进行，因此记得在实际运行代码时删除这一行！
如何使我的模型 XLA 兼容？¶
在许多情况下，你的代码可能已经是 XLA 兼容的！但是，有些在普通 TensorFlow 中有效的方法在 XLA 中无效。我们将其归纳为以下三条核心规则：
🤗 Hugging Face 小贴士 🤗： 我们投入了很多努力重写了 TensorFlow 模型和损失函数，使其 XLA 兼容。我们的模型和损失函数默认遵循规则 #1 和 #2，因此如果你使用的是 transformers 模型，可以跳过这两条规则。但在编写自己的模型和损失函数时不要忘记这些规则！
XLA 规则 #1：代码中不能有“数据依赖的条件语句”¶
这意味着任何 if 语句都不能依赖于 tf.Tensor 内部的值。例如，以下代码块无法用 XLA 编译：
if tf.reduce_sum(tensor) > 10:
tensor = tensor / 2.0
这可能一开始看起来非常限制，但大多数神经网络代码并不需要这样做。你可以通过使用 tf.cond（参见文档 这里）或移除条件并找到一个巧妙的数学技巧来绕过这个限制，如下所示：
sum_over_10 = tf.cast(tf.reduce_sum(tensor) > 10, tf.float32)
tensor = tensor / (1.0 + sum_over_10)
这段代码与上面的效果相同，但由于避免了条件语句，确保它可以无问题地用 XLA 编译！
XLA 规则 #2：代码中不能有“数据依赖的形状”¶
这意味着代码中所有 tf.Tensor 对象的形状不能依赖于它们的值。例如，函数 tf.unique 无法用 XLA 编译，因为它返回一个包含输入中每个唯一值的 tensor。这个输出的形状显然会根据输入 Tensor 的重复程度而不同，因此 XLA 拒绝处理它！
通常情况下，大多数神经网络代码默认遵循规则 #2。但有一些常见的情况会导致问题。一个非常常见的情况是使用 标签掩码，将标签设置为负值以表示在计算损失时应忽略这些位置。如果你查看支持标签掩码的 NumPy 或 PyTorch 损失函数，通常会看到如下代码，使用 布尔索引：
label_mask = labels >= 0
masked_outputs = outputs[label_mask]
masked_labels = labels[label_mask]
loss = compute_loss(masked_outputs, masked_labels)
mean_loss = torch.mean(loss)
这段代码在 NumPy 或 PyTorch 中没有问题，但在 XLA 中会出错！为什么？因为 masked_outputs 和 masked_labels 的形状取决于有多少位置被掩码——这是 数据依赖的形状。然而，就像规则 #1 一样，我们可以重写这段代码以在没有数据依赖形状的情况下产生相同的输出。
label_mask = tf.cast(labels >= 0, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask # 将负标签位置设为 0
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)
通过计算每个位置的损失，但在计算平均值时将掩码位置在分子和分母中都设为 0，我们得到了与第一段代码相同的结果，同时保持了 XLA 兼容性。请注意，我们使用了与规则 #1 相同的技巧——将 tf.bool 转换为 tf.float32 并将其用作指示变量。这是一个非常有用的技巧，记住它如果需要将代码转换为 XLA！
XLA 规则 #3：XLA 需要为每个不同的输入形状重新编译模型¶
这是最重要的一个规则。这意味着如果输入形状变化很大，XLA 将不得不反复重新编译模型，从而导致严重的性能问题。这在 NLP 模型中很常见，因为文本经过分词后的长度各不相同。在其他模态中，静态形状更为常见，这个规则的问题要少得多。
如何解决规则 #3？关键是 填充——如果你将所有输入填充到相同的长度，并使用 attention_mask，你就可以得到与可变形状相同的结果，而不会出现 XLA 问题。然而，过度填充也会导致严重的减速——如果将所有样本填充到整个数据集的最大长度，可能会得到大量填充标记的批次，浪费大量的计算和内存！
没有完美的解决方案。但你可以尝试一些技巧。一个非常有用的技巧是 将样本批次填充到某个数字（如 32 或 64）的倍数。这通常只会增加少量标记，但大大减少了独特的输入形状数量，因为每个输入形状现在必须是 32 或 64 的倍数。较少的独特输入形状意味着较少的 XLA 编译！
🤗 Hugging Face 小贴士 🤗： 我们的分词器和数据整理器有一些方法可以帮助你。你可以在调用分词器时使用 padding="max_length" 或 padding="longest" 来获取填充后的数据。我们的分词器和数据整理器还有一个 pad_to_multiple_of 参数，可以用来减少你看到的独特输入形状数量！
如何在我的 TPU 上实际训练模型？¶
一旦你的训练 XLA 兼容并且（如果你使用 TPU 节点/Colab）数据集已适当准备，运行 TPU 实际上非常简单！你只需要在代码中添加几行来初始化 TPU，并确保模型和数据集在 TPUStrategy 范围内创建。查看 我们的 TPU 示例笔记本 了解实际操作！
总结¶
这里涉及的内容很多，让我们总结一个快速检查清单，帮助你为 TPU 训练做好准备：
- 确保你的代码遵循 XLA 的三条规则
- 在 CPU/GPU 上使用
jit_compile=True编译模型并确认可以使用 XLA 进行训练 - 要么将数据集加载到内存中，要么使用 TPU 兼容的数据加载方法（参见 笔记本）
- 将代码迁移到 Colab（将加速器设置为“TPU”）或 Google Cloud 上的 TPU 虚拟机
- 添加 TPU 初始化代码（参见 笔记本）
- 创建
TPUStrategy并确保数据加载和模型创建都在strategy.scope()内（参见 笔记本） - 不要忘记在迁移到 TPU 时移除
jit_compile=True！ - 🙏🙏🙏🥺🥺🥺
- 调用
model.fit() - 你成功了！