介绍
TPU训练是一项有用的技能：TPU Pod具有高性能和极高的可扩展性，可以轻松训练具有任意规模的模型，从几千万个参数到数百亿个参数：Google的PaLM模型（超过5000亿个参数！）完全是在TPU Pod上进行训练的。
我们之前写过一个教程和一个Colab示例，展示了如何使用TensorFlow进行小规模的TPU训练，并介绍了需要了解的核心概念，以便在TPU上运行您的模型。这次，我们将进一步提升水平，使用TensorFlow和TPU从头开始训练一个遮蔽语言模型，包括从训练分词器和准备数据集到最终模型训练和上传的每个步骤。这是一项任务，您可能需要专用的TPU节点（或VM），而不仅仅是Colab，因此我们将重点放在这里。
与我们的Colab示例一样，我们利用了TensorFlow非常干净的TPU支持，通过XLA和TPUStrategy。我们还将受益于🤗 Transformers中大多数TensorFlow模型完全兼容XLA。令人惊讶的是，只需做一点点工作就可以让它们运行在TPU上。
然而，与我们的Colab示例不同，这个示例旨在可扩展，并且更接近实际的训练过程——尽管我们默认只使用了一个BERT大小的模型，但只需更改几个配置选项，代码就可以扩展到更大的模型和更强大的TPU Pod片段。
动机
为什么我们现在要写这个指南呢？毕竟，🤗 Transformers多年来一直支持TensorFlow。但是，让这些模型在TPU上训练一直是社区的一个主要痛点。这是因为：
- 许多模型不兼容XLA
- 数据收集器没有使用原生的TF操作
我们认为XLA是未来的发展方向：它是JAX的核心编译器，TensorFlow有一流的支持，甚至还可以从PyTorch中使用它。因此，我们大力推动使我们的代码库与XLA兼容，并消除XLA和TPU兼容性之间的任何其他障碍。这意味着用户应该能够在TPU上训练我们的大多数TensorFlow模型而无需麻烦。
现在还有另一个重要的理由关心TPU训练：最近在LLMs和生成式AI方面取得的重大进展引起了公众对模型训练的巨大兴趣，因此对于大多数人来说，获得最先进的GPU变得非常困难。知道如何在TPU上进行训练为您提供了另一种获得超高性能计算硬件的途径，这比在eBay上竞标最后一台H100然后在办公桌前丑陋地哭泣要好得多。您值得更好的待遇。根据经验来说：一旦您熟悉了在TPU上训练，您可能就不想回到过去了。
预期结果
我们将从头开始在WikiText数据集（v1）上训练一个RoBERTa（基础模型）。除了训练模型之外，我们还将训练分词器，将数据进行分词并以TFRecord格式上传到Google Cloud Storage，以供TPU训练使用。您可以在此目录中找到所有的代码。如果您是某种类型的人，您可以跳过本文的其余部分，直接转到代码。如果您留下来，我们将深入研究代码库中的一些关键思想。
这里的许多思想也在我们的Colab示例中提到过，但我们希望向用户展示一个完整的端到端示例，将所有内容整合在一起并展示其实际效果，而不仅仅是在高层次上介绍概念。以下图表以图形方式概述了使用TensorFlow和TPU训练🤗 Transformers语言模型的步骤：
获取数据和训练分词器
如前所述，我们使用了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虚拟机，我们可以在本地使用数据集，甚至将持久存储附加到这些虚拟机上。由于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数组。然而，许多收集器还支持将这些值作为TensorFlow张量返回，如果我们指定return_tensor="tf"。这对于我们的数据管道与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和从tf.data流式传输数据的Transformers模型，您可能会对将整个训练流程迁移到TPU上所需的工作量感到惊讶。它们在某种程度上被认为是晦涩、高端、复杂的硬件，但它们实际上是非常容易上手的，而且相对于保持多个GPU服务器同步而言，实例化一个大型Pod切片肯定更容易!
在2020年代，为最先进的模型提供多样化的训练硬件将至关重要，特别是如果持续存在GPU短缺的情况。我们希望这个指南能够为您提供所需的工具，以便无论面临何种情况，都能进行最先进的训练。
正如伟大的诗人GPT-4曾经说过:
如果你能在周围的人都为GPU短缺而失去头脑的时候保持冷静，如果你相信自己的代码，而其他人对你的能力表示怀疑，如果你能够在TPU上进行训练而没有丝毫犹豫;
如果你能够从错误中学习，并不断前进，不断优化你的目标，以达到更高的高度，那么你就走上了AI的大道，我的朋友，随着时间的推移，你将获得胜利。
当然，这是在无耻地抄袭鲁德亚德·吉卜林的作品，并且它不知道如何发音“drought”，但无论如何，我们希望您感到受到了鼓舞。