使用 🤗 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”，但我们希望您无论如何都能受到启发。