使用 Hugging Face Transformers、TensorFlow 與 TPU 訓練語言模型

Hugging Face 詳細說明了一套可擴展的工作流程，使用 TensorFlow 與 Tensor Processing Units（TPU）從頭訓練遮蔽語言模型（MLM）。透過 XLA（Accelerated Linear Algebra）與 TPUStrategy，開發者可以訓練從數百萬參數到大規模模型（如 Google 的 PaLM，使用 TPU pod 來支援 5000 億參數）的模型。

XLA 相容性與 TPU 可取得性

由於 XLA 不相容以及資料整理器依賴非原生 TensorFlow 操作，過去在 TPU 上訓練 TensorFlow 模型相當困難。Hugging Face 已更新其程式碼庫，使大多數 TensorFlow 模型皆相容 XLA，移除這些障礙，讓 TPU 訓練更易取得。

在高階 GPU 短缺的情況下，此轉變尤為關鍵。TPU 提供高效能的替代方案，讓使用者能取得超高效能運算硬體，為大規模生成式 AI 與 LLM 的訓練提供可擴展的路徑，而不必完全依賴 GPU 的可用性。

端對端訓練工作流程

為了展示此方法的可擴展性，Hugging Face 使用 WikiText（v1）資料集，從頭訓練了一個 RoBERTa-base 模型。此流程遵循特定管線，以確保在 TPU 硬體上效率最佳化：

1. Tokenizer 訓練與資料準備

由於模型是從頭訓練，必須自行建立 tokenizer。工作流程包括：

  • 透過 ‚datasets 載入 WikiText 資料集的 train 切分。
  • 使用 ‚tokenizers 訓練 Unigram 模型。
  • 將產生的 tokenizer 上傳至 Hugging Face Hub。

2. 建立 TFRecord Shards

為了支援大規模平行處理，資料會被轉換成 TFRecord shards，而非單一檔案。tokenization 策略會將樣本串接後切割成固定大小的區塊（128 個 token），以避免因截斷而過度遺失文字內容。

這些 shards 之後會上傳至 Google Cloud Storage（GCS）桶。這對 TPU 節點而言是必要的，因為它們的主機記憶體有限，必須直接從 GCS 串流資料。（註：TPU VM 可使用本機資料集或永久儲存空間）

3. 模型初始化與分散式訓練

為了使用資料平行方式在 TPU 工作者之間分散訓練，模型與 optimizer 必須在 TPUStrategy 範圍內初始化：

import tensorflow as tf

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(...)
strategy = tf.distribute.TPUStrategy(tpu)

with strategy.scope():
    # Model and tokenizer initialization happens here
    model = TFAutoModelForMaskedLM.from_config(config)

TPU 訓練的關鍵技術需求

成功整合 TPU 需要在資料管線與模型設定上做特定配置：

  • TensorFlow 原生資料整理器：DataCollatorForLanguageModeling 必須設定 return_tensor="tf"。此設定確保整理器回傳 TensorFlow 張量而非 NumPy 陣列，對 TPU 相容性至關重要。
  • GCS 整合：TensorFlow 的 tf.io.gfile.glob 可無縫讀取 GCS 桶中的 TFRecord shards，使用 gs:// 標識符。
  • 模型檢查點：使用 PushToHubCallback 在訓練期間直接將模型檢查點同步至 Hugging Face Hub。

推論與結果

模型訓練完成後，可使用標準的 Hugging Face pipeline API 並傳入 framework="tf" 參數進行推論。訓練出的 RoBERTa-base 模型以較長的訓練時間、學習率 1e-4 進行訓練，最終權重已上傳至 Hugging Face Hub。

Sources