Hugging FaceとDeepSpeedによる大規模言語モデルのファインチューニング

Original Blog : Fine-Tuning Large Language Models with Hugging Face and DeepSpeed

翻訳： junichi.maruyama 

ChatGPTのセンセーショナルなリリースを受け、大規模言語モデル（LLM）が現在脚光を浴びています。多くの人が、このようなモデルを自分のアプリケーションでどのように活用できるかを考えています。しかし、これは変換器ベースのモデルのいくつかの進歩の一つに過ぎず、他の多くのモデルは、チャットだけでなく、翻訳、分類、要約などのタスクでオープンかつ容易に利用できます。

以前のブログでは、人気のあるHugging Face トランスフォーマーライブラリを通じて、Databricks上でこれらのモデルにアクセスするための基本的な方法を説明しました。T5やBERT のような、あらかじめ訓練された既製のLLMは、追加のデータや訓練なしで、実世界のさまざまな問題に対してうまく機能します。しかし、特定のタスクでより良いパフォーマンスを発揮するために、これらのモデルを「Fine-Tuning」することが価値ある、あるいは不可欠な場合があります。

このブログでは、T5言語モデルファミリーを最小サイズから最大サイズまで簡単にFine-Tuningし、シンプルなユースケース（多くの製品レビューから製品レビューの概要を構築する）向けに特化させる方法を探ります。また、DatabricksのHugging Faceを使用し、MLflow integration.も行います。

また、MicrosoftのDeepSpeed を導入し、非常に大きな言語モデルのFine-Tuningを加速させる予定です。110億のパラメータを持つモデルでも、Databricks上でFine-Tuningするのは難しいことではありません-それが必要なことであれば！

課題: 製品レビューの要約

あなたがカメラ製品を販売するEコマースサイトを運営しているとしましょう。ユーザーは商品に対してレビューを残しています。あなたは、ユーザーが100のレビューを吟味するのではなく、商品に対するすべてのレビューを1つの要約に凝縮できたらいいと考えています。あなたは、何十万ものレビューと、ユーザーが提供した見出し（レビューの「要約」のようなもの）を収集し、LLMを使ってこれらの製品要約を作成したいと考えています。

この例では、そのようなデータセットの代用として、Amazonの顧客による1億3000万件の商品レビューを含む、Amazon Customer Reviewデータセットを使用することにします。ここで注目すべきは、レビューのテキストと見出し、そして当然ながら「カメラ」カテゴリのものだけである。

このデータセットのフリーテキストは、完全にきれいなものではありません。あなたのeコマースサイトが保持する素敵なキュレーションされたデータセットを模倣するために、データをロードし、いくつかの基本的なクリーニング（詳細は付属のノートブックを参照）を適用し、Deltaテーブルとして記述します。ここでは、レビューの長さを（やや任意に）100トークンに制限しています。これは、非常に長いシーケンスがいくつかあるとメモリ不足のエラーが発生するためです。短い配列の方がFine-Tuningが早く、もちろんある程度の精度は犠牲になりますが、レビューテキストの一部は省略されています。

That's better. Start from here, then see what large language models can do with this data.

t5-smallを使ったクイックサマリー

T5 (Text-to-Text Transfer Transformer)は、Googleが提供する汎用LLMのファミリーです。要約、分類、翻訳など多くのタスクに役立ち、「小さい」（～60Mパラメータ）からかなり大きい（～11Bパラメータ）ものまで、いくつかのサイズが用意されている。これらのサイズは、ますます強力になる一方で、扱うにはますます高価になっています。これらのLLMを扱う際の重要なテーマは、「シンプルに保つ」ことです。小さいモデルで十分であれば、それを使用し、可能であれば既製のリソースで始める。大きなモデルを使った推論には時間とコストがかかるので、レイテンシーや予算の制約から、最初から大きなモデルは使えないかもしれません。

前回のブログでは、SparkとHugging Faceを使ってT5を簡単に適用する方法を紹介しましたが、ここでも同じことが可能です。入力テキストと出力テキストがかなり短いので、まもなくわかるように、このデータセットですぐにうまく機能するかどうかは不明です。いくつかの新要素を加えての総括です：

数行のコードと数分の実行で GPUを必要としないのは悪くありません。結果はレビューの簡潔な要約としてもっともらしく見えます。しかし純正モデルは、過度に短いレビューを要約することに少し苦戦しており、最初の2つの要約では少し行き過ぎた表現になっています！これは、Fine-tuningの必要性を示唆しています。

上記のリストは、いくつかの有用な使用法のヒントを強調しています：

• TRANSFORMERS_CACHE環境変数を/dbfs上の場所に設定し、モデルを一度だけダウンロードし、多くのジョブやクラスタで再利用できるようにします。
• パイプラインをブロードキャストすることで、Sparkが効率的に転送・共有することができます。ここでは必須ではありませんが、大きなモデルでは重要です
• Spark で pandas UDFs を使用して、一度に全バッチの入力を効率的に処理する。
• t5-smallのような小さなモデルから始めて、必要に応じてアップグレードしていく
• 要約パイプラインにはいくつかの設定がありますが、そのうちのいくつかは知っておくと便利です：
  • min_new_tokens は、出力が短くなりすぎないようにするためのものです。これは、より多くのテキストに対してより長い要約を生成するために以下で使用されます。
  • num_beamsは、より多くの可能性を試すことで出力の質を向上させることができますが、その分、計算量は多くなります。

ここまでは目新しいものではありません。しかし、目標は各製品のすべてのレビューを要約することです。Sparkを使えば簡単で、レビューのテキストを集約して代わりにまとめるだけです：

以下は出力の一例で、簡潔のため集計レビューは切り捨てています：

興味深いですが要約がうまくいっていないのは明らかです。次の論理的なステップとして、より大きなモデルを試すか、これらのモデルの一つをFine-Tuningすることが考えられます。次のセクションでは、その両方を試してみることにします。

既製品のファインチューニング

ファインチューニングとは、あらかじめ訓練されたモデルを新しいデータでさらに訓練し、ある特定のタスクの性能を向上させることを意味します。T5のようなモデルは、ある単語列を別の単語に変換するような多くのことを行うように訓練されています。ここで、T5が得意とするのは、多くの商品レビューをレビュー要約に変換することです。これはT5が得意とする要約作業のようなもので、手元にある実際のレビューデータによりよく合うように要約を調整することを望んでいるだけです。これは、T5アーキテクチャをゼロからトレーニングするのとは違います。時間がかかるだけでなく、事前に訓練されたT5モデルが持っている言語に関する学習がすべて失われてしまうからです。

上記の例は、既存のモデルと、その作成に費やした研究、データ、計算能力を再利用しているため、起こっていることの複雑さに比べれば、非常にシンプルです。しかし、fine-Tuningはモデルのトレーニングであり、経験豊富な実務家であっても、トレーニングプロセスを継続するために必要なPyTorchやTensorflowのコードを書くことは簡単ではありません。

幸いなことに、これらのオープンソースモデルには、トレーニングやファインチューニングのコードが付属していることが多い。ノートブックユーザーにとって残念なのは、それらがノートブックではなく、Pythonスクリプトであることです。Databricksではノートブックでもシェルコマンドやgitリポジトリからのスクリプト、あるいはWebターミナルでの対話的な実行が可能なためこのような問題はありません。

ファインチューニングスクリプトの入手

実は、Hugging Faceは、T5モデルで動作する便利なFine-Tuningスクリプトを、Trainer APIを通じて提供しています。なぜこれらのスクリプトを使うのが有利なのかは、後で明らかになりますが、最初はちょっと考えにくいかもしれませんね。既存のソリューションを再利用することは、すぐに始められる素晴らしい方法です。

まず、Hugging FaceのGithubリポジトリをDatabricksのRepoとしてクローンします。スパースチェックアウトモードを使用して、レポ全体ではなく、要約例のみをクローンできます：

これは run_summarization.py スクリプトのコピーを取得します。 run_summarization.py をコピーして好きなRepoに貼り付けても問題ないです。

一つ小さな変更が必要です。このスクリプトは、トランスフォーマーライブラリのバージョンが、期待するものと一致するかどうかをチェックします。このソースチェックアウトは、transformersのソースインストールを期待しています。このファイルを変更する代わりに %pip install git+https://github.com/huggingface/transformers を追加することもできますが、このチェックを削除しても問題ないでしょう。あなたのコピーから check_min_version("...dev0") を読んでいる行を削除してください。

環境セットアップ

どこでもそうですが、これらのスクリプトを使うには少し設定が必要です。Databricks Runtime 12.2 ML (GPU - これは間違いなくGPUが必要です！)を利用します：

• ランタイムにまだ入っていない必要なライブラリのインストール:
• DatabricksがホストするMLflowトラッキングサーバーとHugging FaceのMLflow integrationを接続するための環境変数を設定します:

t5-smallのFine-Tuningは、最近のGPU1つで簡単にこなせます。これらのモデルには、NVIDIAのA10やA100のような最近のAmpereアーキテクチャのGPUを使用するのが有利です。例えば、AWSではg5インスタンスタイプになります。A10はA100よりも入手しやすいかもしれません。

このようなスクリプトは、通常ローカルの入力ファイルを必要とします。ここでは、チューニングスクリプトは、(text,summary)ペアのCSVファイルを必要としています。問題ありません。分散ストレージは /dbfsでローカルファイルのように見えます。Deltaデータセットをトレーニングファイルとバリデーションファイルのペアとして書き出せばよいのです：

Fine-Tuningのチューニング

実際のFine-Tuningは、その後スクリプトを実行することになります。拍子抜けするほどです。

このスクリプトには多くのパラメータがあり、ここで取り上げることはできませんが、いくつかの項目は注目に値するでしょう：

• 環境変数のエクスポートにより、MLflow統合が動作するようになる
• source_prefix： "summarize： " は、データがテキストの要約の例であることをT5が理解するのに役立つ
• optim： Adafactor オプティマイザは、ここでは厳密には必要ありませんが、メモリ使用量の削減は、後で大きなモデルをチューニングする際に重要になります。
• bf16: より多くの数値範囲を保持する bfloat16タイプを使用して、より高速な 16 ビット浮動小数点演算を可能にします。これは最新のAmpere GPUを必要とします
• num_train_epochs: ファインチューニングでは、ゼロからのトレーニングほど多くのエポックを必要としないため、一握りのエポックでも問題ない。
• per_device_train_batch_size: この重要なパラメータは、バッチサイズを制御します。どのようなトレーニングでもそうですが、バッチサイズが大きいとGPUメモリを使い果たす可能性があります。これは、より大きなモデルをFine-Tuningする際に重要な問題となりますが、64のような大きな値は、短いシーケンスを持つデータセットや、24GB以上のメモリを持つ最新のGPUでは問題ないでしょう。あなたの走行距離は大きく変わるでしょう。

これは、例えば数ドルするA10 GPUで1時間程度かかります。トレーニング中（とトレーニング後）、MLflowはステップごとの損失などのトレーニングメトリクスを記録します：

トレーニングはもう少し長く続けられたようですが、8エポックはすでに大まかな再コンバージョンを行うのに十分でした。MLflowがどのようにモデルを追跡し、REST APIとしてデプロイすることができるかについての詳細は、付属のノートブックを参照してください。

サマリーの品質を評価するROUGE metrics のように、モデルとともに自動的に記録される追加のメトリクスを見つけることができます。このメトリクスは、損失よりも結果の品質についてやや意味のあるイメージを与えるため、Fine-Tuningの期間を決定する際に役立ちます。

以上です！オープンなモデルやツールを使って、DatabricksでT5モデルのFine-Tuningを行いました。結果はどうだったでしょうか？t5-small」をあなたのモデルの出力パスに置き換えて、上述の要約パイプラインを再実行するだけです：

文章は、予想通り、確かにレビューの見出しのような感じです。深みに欠けるかもしれませんね！これは、数クリックでREST APIにすることもできます。

REST APIを考えるなら、レイテンシーについても考える価値があります。要約を返すのにどれくらいの時間がかかるか？レイテンシーが重要なら、GPUで実行することが重要です。結論から言うと、1つのサマリーを実行するだけで、約480msかかります。

より高度なサマリーを得るには、より大きなT5モデルを試す価値があります。

DeepSpeedでt5-largeにスケールアップする。

t5-largeのような大きなモデルにスケールアップしても、パラメータが60Mに対して770Mと、桁違いの大きさであるにもかかわらず、ほとんど変化はありません。実際、上記のFine-Tuningスクリプトを実行する際に変わるべきことは、以下の通りです：

• 複数のGPUに切り替える - たとえば、1つのGPUではなく、4つのA10 GPUに切り替える
• t5-largeをご指定ください。
• バッチサイズを12個程度に減らす

より大きなモデルは、より多くのハードウェアを必要とし、一度にGPUに搭載できる量が少なくなる。バッチサイズを正しく設定するのは、シーケンスの長さがまちまちで、長い場合もあるため、難しい場合があります。そのため、データ準備ではレビューや要約の長さを制限しています。また、スクリプトにはmax_source_lengthのようなオプションがあり、手動で入力を切り捨てることができます。入力が小さいとスケールアップに役立ちますが、問題によっては、任意に入力を切り捨てることでモデルの品質を損なう可能性があります。

これらの変更を加えて再実行した場合、うまくいくことがわかるでしょう。また、1エポックあたり2時間かかる。コストは、もはや数ドルではなく、8エポックで90ドル程度になる。このため、エポック数をよく考えることが重要になる。例えば、上のプロットでは、4エポックだけ実行しても損失が大きくなることはないようなので、少なくとも最初のうちは、4エポックだけ実行するのが妥当かもしれません。必要であれば、チェックポイントを再開し、さらにトレーニングを行うことができます。

DeepSpeedに参入

また、Hugging Faceのようなツールがすぐに提供してくれるものよりも、より高度な並列化を利用することが重要になります。幸いなことに、今回もオープンソースにいくつかの答えがあります。MicrosoftのDeepSpeedは、洗練された最適化を多数実装することで、既存の深層学習トレーニングや推論ジョブをほとんど変更することなく高速化することができます。特に興味深いのは、メモリ使用量を減らそうとする最適化セットであるZeROです。詳細や論文については、DeepSpeedのサイトを参照してください。

DeepSpeedは、Hugging FaceのTrainer APIを使用するFine-Tuningジョブを自動的に最適化することができ、既存のFine-Tuningスクリプトを実行するためのドロップイン代替スクリプトを提供しています。これは、既製のトレーニングスクリプトを再利用することが有利である理由の1つです。

DeepSpeedを使用するには、accelerateと一緒にそのパッケージをインストールします。ソースからインストールすることをお勧めしますが、リリースされたパッケージをインストールすることもできます。

DeepSpeedで同じように実行しても、ほんの少ししか変わりません：

python は deepspeed runner スクリプトに置き換えられます。また、設定ファイルへのパスが必要です。このファイルのオプションについては、ここでは触れませんが、デフォルトの「ZeRO stage 2」構成から始めてマイナーチェンジするのが合理的です。たとえば、現在の設定では、次の2つの編集が必要です：

• bf16 を明示的に有効にする（通常の float16 / fp16 サポートは無効にする）
• オプティマイザーの設定を削除し、AdamWの代わりにAdafactor を使用するようにしました。Adafactorはメモリ消費量が非常に少なく、大規模なモデルの微調整を行う際に良い代替となります。

DeepSpeedでは、さらにいくつかの改良を加えることが可能です：

• gradient_checkpointing：前方パスで大きな中間結果を解放し、後方パスで再計算することで、少し計算量が増える代わりにメモリを節約する。
• デバイスごとのバッチサイズは20まで可能です（必要であればもっと大きくなります）。4GPUの場合、有効なバッチサイズは4 x 20 = 80になることに注意してください。

バッチサイズは重要なチューニングの問題になります。バッチサイズはデバイスごとに調整されることが多いのですが、それは1つのGPUが一度に処理できる量を制約するのは、個々のGPUメモリだからです。バッチサイズが大きいとGPUメモリを使い切らない限りスループットが向上します！最大バッチサイズは、GPUメモリ、入力配列のサイズ、モデル内の最大レイヤーのサイズ、オプティマイザーの設定など、いくつかの要因に依存します。一般的には試行錯誤しながら調整する必要があります。

トレーニングにおいても、過度に大きなバッチサイズが問題になることもあります。しかし非常に大きな言語モデルの場合、一般的には数個、あるいは1個のバッチを各デバイスのメモリに収める方法を見つけることが問題です。少なくともバッチサイズは結果を再現するために重要なので、有効なバッチサイズは（デバイス数×デバイスごとのバッチサイズ）であることは覚えておく必要があります。

これにより、実行時間は1エポックあたり40分程度に低下しています。3倍速いということは、3倍安いということでもあり、このチューニングは30ドル程度で済むかもしれません。これは、時間が日単位、コストが数百ドル、数千ドル単位になったときに、ますます価値が高まるでしょう。

この改善は、DatabricksクラスタのGangliaメトリクスのように、GPUメトリクスで確認することができます：

DeepSpeedを使用しない場合：

With DeepSpeed:

このモデルでサマリーを見るとどうでしょうか？

簡潔で、レビューの文章から正確な詳細を提供するようになったので、おそらくさらに良くなっています。シングルGPUでのモデルレイテンシは約3秒となり、より大きなモデルへのスケールアウトを検討する際には、すでに躊躇してしまうかもしれません。ここで止まってしまうかもしれませんが、最大のT5モデルまでスケールアウトすることも可能です。

Super-Size It: Fine-Tuning t5-11b

T5の最大モデルはt5-11bで、ご想像の通り110億個のパラメータを持ち、t5-largeの14倍以上です。それをFine-Tuningすることは、クラウドで利用できる最大のタイプとはいえ、1台のマシンで同じ手法で可能です。

これは大変な作業で、現実的に考えて、これを合理的な時間で完成させるには、さらにいくつかの変更が必要です：

• 400GB以上のRAMを搭載した8GPUマシン
• 1トレーニングエポック（またはそれ以下）
• 1台あたりのバッチサイズは8（有効バッチサイズは8×8＝64）
• パラメーターのオフロードを可能にした、パラメーターの完全なパーティショニングを含む「ZeROステージ3」最適化

しかし、このデータセットと入力サイズでは、1回のエポックに約1.8日かかってしまいます。コストは1エポックあたり数百ドルに近づきます。これは確かにもっとチューニングできますが、予想される計算の規模を知ることができます。この例では、数百ステップ（1エポックの1％）のFine-Tuningで次のようなサマリーを得ることができました：

質的にも、より良いまとめ方だと思います。しかしこの特定の問題では、たとえそれが完全に可能であったとしても、時間とコストに見合うとは思えません。Fine-Tuningのコストと時間が許容できる場合でも、推論のコストと時間は許容できないかもしれません。例えば t5-11bを使った推論では、GPUで数十秒かかる可能性があり、それは遅すぎるかもしれません。ほとんどの問題では、この規模かそれ以下で十分ですが、非常に大規模なチューニングは容易に可能です。

Conclusion and Next Directions

大規模言語モデルは、様々なビジネス問題に対する強力な新ツールであり、オープンソースのものは、Databricks上でオープンソースツールを使って、そのまま簡単に適用することができます。このような大規模な言語モデルのFine-Tuningも、オープンソースのツールを使えば同様に簡単です。スクリプトもDatabricksのノートブックでは問題ありません。これらの簡単なアプローチは、ほとんどすべての実世界の問題に対して十分なサイズにスケールアップすることができます。

T5の様々なサイズをFine-Tuningするための迅速な実験の結果、Fine-Tuningと推論の実行に必要なリソースの大きさが桁違いであることがわかります：

しかし、最大級のマシンが提供する以上のリソースを必要とするケースは常に存在する。さらに作業を進めれば、これらのツールはDatabricks上のマシンのクラスタに適応させることができ、これは将来のブログのテーマとなります。

Databricksで試してみてください！このノートブックアーカイブをリポジトリにインポートして利用できます。