使用 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.ndarray
或tf.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_function
或 Dataset.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_outputs
和masked_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()
- 你做到了!