使用 🤗 Transformers、TensorFlow 和 TPU 训练语言模型
引言
TPU 训练是一项很有用的技能:TPU pods 具有高性能和极强的可扩展性,可以轻松地以任何规模训练模型,从数千万参数到真正庞大的模型:Google 的 PaLM 模型(超过 5000 亿参数!)完全在 TPU pods 上训练。
我们之前曾撰写过一篇教程和一份Colab 示例,展示了使用 TensorFlow 进行小规模 TPU 训练,并介绍了使模型在 TPU 上运行所需理解的核心概念。这次,我们将更进一步,使用 TensorFlow 和 TPU 从头开始训练一个掩码语言模型,包括从训练分词器、准备数据集到最终模型训练和上传的每一个步骤。这种任务可能需要专用的 TPU 节点(或 VM),而不仅仅是 Colab,因此我们将重点放在这方面。
与我们的 Colab 示例一样,我们利用了 TensorFlow 通过 XLA 和 TPUStrategy
提供的非常简洁的 TPU 支持。我们还将受益于 🤗 Transformers 中大多数 TensorFlow 模型完全XLA 兼容的事实。因此,令人惊讶的是,让它们在 TPU 上运行所需的工作量很少。
然而,与我们的 Colab 示例不同,此示例旨在具有**可扩展性**,并且更接近真实的训练运行——尽管我们默认只使用 BERT 大小的模型,但通过更改一些配置选项,代码可以扩展到更大的模型和更强大的 TPU pod 切片。
动机
为什么我们现在要写这份指南?毕竟,🤗 Transformers 对 TensorFlow 的支持已经有好几年了。但是让这些模型在 TPU 上训练一直是社区的一个主要痛点。这是因为
- 许多模型不兼容 XLA
- 数据校对器不使用原生 TF 操作
我们认为 XLA 是未来:它是 JAX 的核心编译器,它在 TensorFlow 中具有一流的支持,您甚至可以从PyTorch中使用它。因此,我们已经大力推动使我们的代码库与 XLA 兼容,并消除阻碍 XLA 和 TPU 兼容性的任何其他障碍。这意味着用户应该能够轻松地在 TPU 上训练我们大多数的 TensorFlow 模型。
现在关注 TPU 训练还有一个重要原因:LLM 和生成式 AI 的最新重大进展引起了公众对模型训练的巨大兴趣,因此大多数人很难获得最先进的 GPU。了解如何在 TPU 上训练为您提供了访问超高性能计算硬件的另一条途径,这比在 eBay 上竞标最后一个 H100 失败然后坐在办公桌前丑陋地哭泣要体面得多。您值得拥有更好的。并且凭经验而言:一旦您习惯了在 TPU 上训练,您可能就不想再回去了。
预期内容
我们将从头开始在 WikiText 数据集 (v1) 上训练一个 RoBERTa(基础模型)。除了训练模型,我们还将训练分词器,对数据进行分词并将其以 TFRecord 格式上传到 Google Cloud Storage,以便 TPU 训练访问。您可以在此目录中找到所有代码。如果您是那种人,您可以跳过本博客文章的其余部分,直接跳转到代码。但是,如果您留下来,我们将更深入地探讨代码库中的一些关键思想。
本文中的许多想法也曾出现在我们的 Colab 示例中,但我们希望向用户展示一个完整的端到端示例,将其整合在一起并实际运行,而不仅仅是停留在概念层面。下图从图片上概述了使用 🤗 Transformers、TensorFlow 和 TPU 训练语言模型的步骤:
获取数据并训练分词器
如前所述,我们使用了WikiText 数据集 (v1)。您可以访问Hugging Face Hub 上的数据集页面来探索该数据集。
由于该数据集已以兼容格式在 Hub 上可用,我们可以轻松地使用 🤗 datasets 加载并与之交互。然而,对于本示例,由于我们还要从头开始训练分词器,所以我们采取了以下步骤:
- 使用 🤗 datasets 加载 WikiText 的
train
分割。 - 利用 🤗 tokenizers 训练了一个Unigram 模型。
- 将训练好的分词器上传到 Hub。
您可以在这里找到分词器训练代码,并在这里找到分词器。此脚本还允许您使用 Hub 上的任何兼容数据集运行它。
💡 使用 🤗 datasets 托管您的文本数据集非常简单。请参阅此指南了解更多信息。
对数据进行分词并创建 TFRecords
分词器训练完成后,我们可以在所有数据集划分(本例中为 train
、validation
和 test
)上使用它,并从中创建 TFRecord 分片。将数据划分分布在多个 TFRecord 分片中有助于大规模并行处理,而不是将每个划分放在单个 TFRecord 文件中。
我们单独对样本进行分词。然后,我们取一批样本,将它们连接起来,并将其分成几个固定大小的块(在我们的例子中是 128)。我们遵循这种策略而不是以固定长度对一批样本进行分词,以避免过度丢弃文本内容(由于截断)。
然后,我们批量获取这些分词后的样本,并将这些批量序列化为多个 TFRecord 分片,其中总数据集长度和单个分片大小决定了分片的数量。最后,这些分片被推送到Google Cloud Storage (GCS) 存储桶。
如果您使用 TPU 节点进行训练,那么数据需要从 GCS 存储桶流式传输,因为节点主机内存非常小。但是对于 TPU VM,我们可以本地使用数据集,甚至可以将持久存储连接到这些 VM。由于 TPU 节点仍然大量使用,我们的示例基于使用 GCS 存储桶进行数据存储。
您可以在此脚本中查看所有这些代码。为了方便起见,我们还在 Hub 上的此存储库中托管了生成的 TFRecord 分片。
在 GCS 中的数据上训练模型
如果您熟悉使用 🤗 Transformers,那么您已经知道建模代码
from transformers import AutoConfig, AutoTokenizer, TFAutoModelForMaskedLM
tokenizer = AutoTokenizer.from_pretrained("tf-tpu/unigram-tokenizer-wikitext")
config = AutoConfig.from_pretrained("roberta-base")
config.vocab_size = tokenizer.vocab_size
model = TFAutoModelForMaskedLM.from_config(config)
但由于我们是在 TPU 领域,我们需要在策略范围内执行此初始化,以便可以通过数据并行训练将其分布到 TPU 工作器中。
import tensorflow as tf
tpu = tf.distribute.cluster_resolver.TPUClusterResolver(...)
strategy = tf.distribute.TPUStrategy(tpu)
with strategy.scope():
tokenizer = AutoTokenizer.from_pretrained("tf-tpu/unigram-tokenizer-wikitext")
config = AutoConfig.from_pretrained("roberta-base")
config.vocab_size = tokenizer.vocab_size
model = TFAutoModelForMaskedLM.from_config(config)
同样,优化器也需要在与模型将进一步编译的相同策略范围内初始化。在本篇博文中,我们不想详细介绍完整的训练代码,所以我们欢迎您在此处阅读它。相反,让我们讨论另一个关键点——一个 TensorFlow 原生数据校对器——DataCollatorForLanguageModeling
。
DataCollatorForLanguageModeling
负责掩码输入序列中随机选择的标记并准备标签。默认情况下,我们以 NumPy 数组的形式返回这些 collator 的结果。但是,如果我们将 return_tensor="tf"
,许多 collator 也支持将这些值作为 TensorFlow 张量返回。这对于我们的数据管道与 TPU 训练兼容至关重要。
谢天谢地,TensorFlow 为从 GCS 存储桶读取文件提供了无缝支持:
training_records = tf.io.gfile.glob(os.path.join(args.train_dataset, "*.tfrecord"))
如果 args.dataset
包含 gs://
标识符,TensorFlow 将理解它需要查看 GCS 存储桶。本地加载就像删除 gs://
标识符一样简单。对于其余与数据管道相关的代码,您可以参考训练脚本中的此部分。
一旦数据集准备好,模型和优化器初始化完毕,模型也已编译,我们就可以进行社区最喜欢的操作——model.fit()
。为了训练,我们没有进行大量的超参数调优。我们只是以 1e-4 的学习率训练了更长时间。我们还利用了 PushToHubCallback
来进行模型检查点和与 Hub 同步。您可以在这里找到超参数细节和训练好的模型:https://huggingface.co/tf-tpu/roberta-base-epochs-500-no-wd。
模型训练完成后,使用它进行推理就像这样简单:
from transformers import pipeline
model_id = "tf-tpu/roberta-base-epochs-500-no-wd"
unmasker = pipeline("fill-mask", model=model_id, framework="tf")
unmasker("Goal of my life is to [MASK].")
[{'score': 0.1003185287117958,
'token': 52,
'token_str': 'be',
'sequence': 'Goal of my life is to be.'},
{'score': 0.032648514956235886,
'token': 5,
'token_str': '',
'sequence': 'Goal of my life is to .'},
{'score': 0.02152673341333866,
'token': 138,
'token_str': 'work',
'sequence': 'Goal of my life is to work.'},
{'score': 0.019547373056411743,
'token': 984,
'token_str': 'act',
'sequence': 'Goal of my life is to act.'},
{'score': 0.01939118467271328,
'token': 73,
'token_str': 'have',
'sequence': 'Goal of my life is to have.'}]
结论
如果说我们希望通过这个例子强调一件事,那就是 TPU 训练是**强大、可扩展且简单的。**事实上,如果您已经在使用带有 TF/Keras 的 Transformers 模型并从 tf.data
流式传输数据,您可能会对将整个训练管道迁移到 TPU 所需的工作量感到震惊。它们以有点神秘、高端、复杂硬件而闻名,但它们相当平易近人,实例化一个大型 pod 切片肯定比保持多个 GPU 服务器同步更容易!
在 2020 年代,使最先进模型训练所使用的硬件多样化将至关重要,特别是如果持续的 GPU 短缺继续下去。我们希望本指南能为您提供所需的工具,无论您面临何种情况,都能为尖端训练运行提供动力。
正如伟大的诗人 GPT-4 曾经说过:
当周围的人都因 GPU 短缺而失去理智时,
如果你能保持清醒,
并且相信你的代码,而其他人怀疑你,
毫无犹豫地在 TPU 上训练;
如果你能从错误中学习,并继续前进,
并优化你的目标以达到顶峰,
人工智能掌握之路属于你,
我的朋友,随着时间的推移,你将取得胜利。
当然,这无耻地抄袭了拉迪亚德·吉卜林,而且它不知道如何发音“drought”,但我们希望您无论如何都能受到启发。