Hugging Face Transformers、TensorFlow、TPU を使用した言語モデルのトレーニング

Hugging Face は、TensorFlow と Tensor Processing Units（TPU）を使用して、マスク付き言語モデル（MLM）をスクラッチからトレーニングするためのスケーラブルなワークフローを詳述しています。XLA（Accelerated Linear Algebra）と TPUStrategy を活用することで、開発者は数百万パラメータのモデルから、Google の PaLM モデルのように 5000 億パラメータを TPU ポッドで利用した大規模モデルまでトレーニングできます。

XLA 互換性と TPU のアクセシビリティ

TPU 上で TensorFlow モデルをトレーニングすることは、これまで XLA 互換性の欠如や、TensorFlow のネイティブでない操作に依存するデータコラレータが原因で困難でした。Hugging Face はコードベースを更新し、ほとんどの TensorFlow モデルが XLA 互換になるようにしたことで、これらの障壁を取り除き、TPU トレーニングをより利用しやすくしました。

この変化は、高性能 GPU の不足が続く中で重要です。TPU は超高性能コンピューティングハードウェアへのアクセス手段として高性能な代替手段を提供し、GPU のみへの依存なしに大規模生成 AI や LLM のトレーニングをスケーラブルに行う道筋を示します。

エンドツーエンドのトレーニングワークフロー

このアプローチのスケーラビリティを示すために、Hugging Face は WikiText（v1）データセットを使用して RoBERTa-base モデルをスクラッチからトレーニングしました。プロセスは TPU ハードウェア上での効率性を確保するために、特定のパイプラインに従います：

1. トークナイザのトレーニングとデータ準備

モデルをスクラッチからトレーニングするため、カスタムトークナイザが必要です。ワークフローは以下を含みます：

  • WikiText データセットの train 分割を ‚datasets を介してロードします。
  • tokenizers を使用して Unigram モデルをトレーニングします。
  • 生成されたトークナイザを Hugging Face Hub にアップロードします。

2. TFRecord シャードの作成

大規模な並列処理を可能にするため、データは単一ファイルではなく TFRecord シャードに変換されます。トークナイゼーション戦略は、サンプルを連結し、固定サイズのチャンク（128 トークン）に分割することで、切り捨てによるテキスト内容の過度な損失を防ぎます。

これらのシャードは Google Cloud Storage（GCS）バケットにアップロードされます。TPU ノードはホストメモリが限られており、データを GCS から直接ストリームする必要があるためです。（注：TPU VM はローカルデータセットや永続ストレージを使用できます）

3. モデルの初期化と分散トレーニング

データ並列性を用いて TPU ワーカー間でトレーニングを分散させるには、モデルとオプティマイザを 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-Native Data Collators: DataCollatorForLanguageModeling は return_tensor="tf" で設定する必要があります。これにより、コラレータは NumPy 配列ではなく TensorFlow テンソルを返すようになり、TPU 互換性に不可欠です。
  • GCS Integration: TensorFlow の tf.io.gfile.glob を使用すると、gs:// 識別子で GCS バケットから TFRecord シャードをシームレスに読み取れます。
  • Model Checkpointing: PushToHubCallback は、トレーニング中にモデルのチェックポイントを直接 Hugging Face Hub に同期するために使用されます。

推論と結果

トレーニングが完了すると、標準の Hugging Face pipeline API に framework="tf" 引数を指定して推論にデプロイできます。トレーニングされた RoBERTa-base モデルは学習率 1e-4 で長時間トレーニングされ、得られた重みは Hugging Face Hub で入手可能です。

Sources