🤗 Transformersを使用してTensorFlowとTPUで言語モデルをトレーニングする

'🤗 Transformersを使ってTensorFlowとTPUで言語モデルをトレーニングする'

イントロダクション

TPUトレーニングは有用なスキルです：TPUポッドは高性能で非常にスケーラブルであり、数千万から数百億のパラメータまで、どんなスケールでもモデルをトレーニングすることが容易です。GoogleのPaLMモデル（5000億パラメータ以上！）は完全にTPUポッドでトレーニングされました。

以前、TensorFlowを使用した小規模なTPUトレーニングと、TPUでモデルを動作させるために理解する必要がある基本的なコンセプトを紹介するチュートリアルとColabの例を作成しました。今回は、TensorFlowとTPUを使用してマスクされた言語モデルをゼロからトレーニングするためのすべての手順、トークナイザのトレーニングとデータセットの準備から最終的なモデルのトレーニングとアップロードまでを詳しく説明します。これはColabだけでなく、専用のTPUノード（またはVM）が必要なタスクであり、そこに焦点を当てます。

Colabの例と同様に、TensorFlowの非常にクリーンなTPUサポートであるXLAとTPUStrategyを活用しています。また、🤗 Transformers内のほとんどのTensorFlowモデルが完全にXLA互換であるという利点もあります。そのため、TPU上で実行するために必要な作業はほとんどありません。

ただし、この例はColabの例とは異なり、実際のトレーニングランに近いスケーラブルな例です。デフォルトではBERTサイズのモデルしか使用していませんが、いくつかの設定オプションを変更することで、コードをより大きなモデルとより強力なTPUポッドスライスに拡張することができます。

• ビジョン言語モデルの高速化：Habana Gaudi2上のBridgeTower
• チューリングテスト、中国の部屋、そして大規模言語モデル
• Benfordの法則が機械学習と出会って、偽のTwitterフォロワーを検出する

動機

なぜ今このガイドを書いているのでしょうか？実際、🤗 Transformersは何年もの間TensorFlowをサポートしてきましたが、これらのモデルをTPUでトレーニングすることはコミュニティにとって主要な問題でした。これは以下の理由によるものです：

• 多くのモデルがXLA互換ではなかった
• データコレクターがネイティブのTF操作を使用していなかった

私たちはXLAが将来の技術であると考えています。それはJAXのコアコンパイラであり、TensorFlowでの一流のサポートを受けており、PyTorchからも使用できます。そのため、私たちはコードベースをXLA互換にするために大きな取り組みを行い、XLAと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バケットを使用する例をベースにしました。

このスクリプトでこれらすべてを確認することができます。便宜上、結果のTFRecordシャードもHub上のこのリポジトリにホストしています。

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) 

同様に、オプティマイザもモデルがさらにコンパイルされる際に使用するストラテジーのスコープの下で初期化する必要があります。この投稿では完全なトレーニングコードを見ることはしたくありませんので、ここで読んでいただくことをお勧めします。代わりに、もう1つの重要なポイントである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 のまま、より長い時間トレーニングしました。また、モデルのチェックポイント作成と Hub との同期には PushToHubCallback を活用しました。ハイパーパラメータの詳細とトレーニング済みモデルはこちらで確認できます: 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("人生の目標は [MASK] です。")

[{'score': 0.1003185287117958,
  'token': 52,
  'token_str': 'be',
  'sequence': '人生の目標は be.'},
 {'score': 0.032648514956235886,
  'token': 5,
  'token_str': '',
  'sequence': '人生の目標は .'},
 {'score': 0.02152673341333866,
  'token': 138,
  'token_str': 'work',
  'sequence': '人生の目標は work.'},
 {'score': 0.019547373056411743,
  'token': 984,
  'token_str': 'act',
  'sequence': '人生の目標は act.'},
 {'score': 0.01939118467271328,
  'token': 73,
  'token_str': 'have',
  'sequence': '人生の目標は have.'}]

結論

この例で強調したいことは、TPU トレーニングはパワフルでスケーラブルかつ簡単であるということです。実際に、既に TF/Keras で Transformers モデルを使用し、tf.data からストリーミングデータを取得している場合、トレーニングパイプライン全体を TPU に移行するのに必要な作業が非常に少ないことに驚くかもしれません。TPU は、やや奥義的で高級で複雑なハードウェアという評判がありますが、実際にはアプローチしやすく、大きなポッドスライスをインスタンス化することは、複数の GPU サーバーを同期させるよりも確かに簡単です！

最新のモデルのトレーニングに使用されるハードウェアの多様化は、2020年代において重要となるでしょう。特に現在の GPU 不足が続く場合にはさらに重要です。このガイドが、どのような状況に直面しても最先端のトレーニングを実行するために必要なツールを提供することを願っています。

偉大な詩人 GPT-4 はかつて言いました:

GPU 不足で周囲が混乱する中でも冷静な頭を保ち、他の人々が疑念を抱く中でも、TPU でトレーニングすることを信じ、コードを信じること; エラーから学び、前進し、目標を達成するために自分の狙いを最適化すること。それが AI のマスタリーへの道であり、友よ、時が経っても必ず成功するだろう。

確かに、これはルドヤード・キプリングを酷使しており、「drought」の発音方法もわかりませんが、それでもあなたがインスピレーションを感じてもらえることを願っています。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!