Transformers 文档

使用TensorFlow在TPU上进行训练

使用 TensorFlow 在 TPU 上进行训练

如果你不需要冗长的解释,只想获取TPU代码示例来开始,请查看我们的TPU示例笔记本!

什么是TPU?

TPU 是 张量处理单元。 它们是由谷歌设计的硬件,用于大幅加速神经网络中的张量计算,类似于 GPU。它们可以用于网络训练和推理。通常通过谷歌的云服务访问,但小型 TPU 也可以通过 Google Colab 和 Kaggle Kernels 直接免费访问。

因为🤗 Transformers中的所有TensorFlow模型都是Keras模型,所以本文档中的大多数方法通常适用于任何Keras模型的TPU训练!然而,有一些点是特定于HuggingFace生态系统(hug-o-system?)的Transformers和Datasets的,我们会在遇到它们时确保标记出来。

有哪些类型的TPU可用?

新用户通常对TPU的种类及其不同的访问方式感到非常困惑。首先需要理解的关键区别是TPU节点TPU虚拟机之间的差异。

当你使用TPU节点时,你实际上是间接访问远程TPU。你需要一个单独的虚拟机,它将初始化你的网络和数据管道,然后将它们转发到远程节点。当你在Google Colab上使用TPU时,你是在以TPU节点的方式访问它。

使用TPU节点可能会让不习惯它们的人感到非常意外!特别是,由于TPU位于与运行Python代码的机器物理上不同的系统上,你的数据不能存储在本地机器上——任何从机器内部存储加载的数据管道都会完全失败!相反,数据必须存储在Google Cloud Storage中,这样即使管道在远程TPU节点上运行,你的数据管道仍然可以访问它。

如果你可以将所有数据以np.ndarraytf.Tensor的形式放入内存中,那么即使在使用Colab或TPU节点时,你也可以在该数据上fit(),而无需将其上传到Google Cloud Storage。

🤗特定的 Hugging Face 提示🤗: 方法 Dataset.to_tf_dataset() 及其更高级别的封装 model.prepare_tf_dataset(),在我们的 TF 代码示例中随处可见,但在 TPU 节点上都会失败。原因是尽管它们创建了一个 tf.data.Dataset,但它不是一个“纯”的 tf.data 管道,并且使用 tf.numpy_functionDataset.from_generator() 从底层的 HuggingFace Dataset 流式传输数据。这个 HuggingFace Dataset 由本地磁盘上的数据支持,远程 TPU 节点将无法读取这些数据。

访问TPU的第二种方式是通过TPU虚拟机。 使用TPU虚拟机时,您直接连接到TPU所连接的机器,就像在GPU虚拟机上训练一样。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但不到整个pod时(例如v3-32),您的TPU集群被称为pod slice。

当您通过Colab访问免费TPU时,通常会获得一个v2-8 TPU。

我经常听到关于XLA的事情。什么是XLA,它与TPU有什么关系?

XLA 是一个优化编译器,被 TensorFlow 和 JAX 使用。在 JAX 中,它是唯一的编译器,而在 TensorFlow 中它是可选的(但在 TPU 上是强制使用的!)。在训练 Keras 模型时启用它的最简单方法是将参数 jit_compile=True 传递给 model.compile()。如果你没有遇到任何错误并且性能良好,那是一个很好的迹象,表明你已经准备好迁移到 TPU 了!

在TPU上调试通常比在CPU/GPU上稍微困难一些,因此我们建议在尝试在TPU上运行之前,先在CPU/GPU上使用XLA运行你的代码。当然,你不需要长时间训练,只需进行几步以确保你的模型和数据管道按预期工作。

XLA编译的代码通常更快 - 所以即使你不打算在TPU上运行,添加jit_compile=True也可以提高性能。不过,请务必注意下面关于XLA兼容性的注意事项!

痛苦经验带来的提示: 虽然使用 jit_compile=True 是提升速度和测试你的CPU/GPU代码是否与XLA兼容的好方法,但如果你在实际在TPU上训练时保留它,实际上可能会引起很多问题。XLA编译在TPU上会隐式发生,所以记得在实际在TPU上运行代码之前移除那一行!

如何使我的模型兼容XLA?

在许多情况下,您的代码可能已经与XLA兼容!然而,有一些在普通TensorFlow中有效的东西在XLA中无效。我们已将它们提炼为以下三个核心规则:

🤗特定的HuggingFace提示🤗: 我们投入了大量精力重写我们的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_outputsmasked_labels的形状取决于有多少位置被屏蔽——这使得它成为一个数据依赖的形状。然而,就像规则#1一样,我们通常可以重写这段代码,以产生完全相同的输出,而没有任何数据依赖的形状。

label_mask = tf.cast(labels >= 0, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask  # Set negative label positions to 0
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)

在这里,我们通过计算每个位置的损失来避免数据依赖的形状,但在计算平均值时,将掩码位置在分子和分母中都置为零,这样在保持XLA兼容性的同时,得到了与第一个块完全相同的结果。请注意,我们使用了与规则#1中相同的技巧——将tf.bool转换为tf.float32并将其用作指示变量。这是一个非常有用的技巧,所以如果你需要将自己的代码转换为XLA,请记住它!

XLA 规则 #3: XLA 需要为每个不同的输入形状重新编译您的模型

这是最重要的一点。这意味着如果你的输入形状非常多变,XLA将不得不反复重新编译你的模型,这将导致巨大的性能问题。这在NLP模型中很常见,其中输入文本在分词后具有可变长度。在其他模式中,静态形状更为常见,因此这条规则的问题要小得多。

如何绕过规则#3?关键是填充 - 如果你将所有输入填充到相同的长度,然后使用attention_mask,你可以得到与可变形状相同的结果,但不会遇到任何XLA问题。然而,过多的填充也会导致严重的减速 - 如果你将所有样本填充到整个数据集中的最大长度,你可能会得到由无尽的填充标记组成的批次,这将浪费大量的计算和内存!

这个问题没有一个完美的解决方案。然而,你可以尝试一些技巧。一个非常有用的技巧是将样本批次填充到32或64个标记的倍数。这通常只会稍微增加标记的数量,但它大大减少了唯一输入形状的数量,因为现在每个输入形状都必须是32或64的倍数。更少的唯一输入形状意味着更少的XLA编译!

🤗特定的HuggingFace提示🤗: 我们的分词器和数据整理器有一些方法可以在这里帮助你。你可以在调用分词器时使用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兼容的数据集加载方法(参见notebook
  • 将您的代码迁移到Colab(加速器设置为“TPU”)或Google Cloud上的TPU VM
  • 添加TPU初始化代码(参见notebook
  • 创建你的 TPUStrategy 并确保数据集加载和模型创建在 strategy.scope() 内部(参见 notebook
  • 别忘了在迁移到TPU时再次移除jit_compile=True
  • 🙏🙏🙏🥺🥺🥺
  • 调用 model.fit()
  • 你做到了!
< > Update on GitHub