使用 🤗 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”，但我們希望您無論如何都能受到啟發。
