ZeROを使用して、DeepSpeedとFairScaleを介してより多くのフィットと高速なトレーニングを実現

使用するZeROを介して、DeepSpeedとFairScaleを使って、より多くのフィットと高速なトレーニングを実現する

Hugging FaceフェローであるStas Bekmanによるゲストブログ投稿

最近の機械学習モデルは、新しくリリースされたカードに追加されるGPUメモリ量よりもはるかに速く成長しているため、多くのユーザーはこれらの巨大なモデルを自分のハードウェアにトレーニングしたり、ロードしたりすることができません。これらの巨大なモデルをより管理しやすいサイズに縮小するための取り組みが進行中ですが、それらの努力は十分に早く小さなモデルを生み出すことはありません。

2019年の秋に、Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase、Yuxiong Heが「ZeRO: Memory Optimizations Toward Training Trillion Parameter Models」という論文を発表しました。この論文には、以前に考えられていたよりもハードウェアを遥かに高い性能で動作させるための、多くの独創的なアイデアが含まれています。その後しばらくして、DeepSpeedがリリースされ、その論文のアイデアのほとんどをオープンソースで実装しました（いくつかのアイデアはまだ進行中です）。同時に、FacebookのチームもZeROの論文のいくつかの核心的なアイデアを実装したFairScaleをリリースしました。

Hugging FaceのTrainerを使用している場合、transformers v4.2.0以降、DeepSpeedとFairScaleのZeRO機能の実験的なサポートが提供されています。新しい--sharded_ddpおよび--deepspeedコマンドラインのTrainer引数は、それぞれFairScaleとDeepSpeedの統合を提供します。こちらが完全なドキュメントです。

このブログ投稿では、単一のGPUを所有している場合でも、複数のGPUを所有している場合でも、ZeROの利点をどのように得るかについて説明します。

• Hugging Face Transformersでより高速なTensorFlowモデル
• PyTorch / XLA TPUsでのHugging Face
• Huggingface TransformersとRayを使用した検索増強生成

翻訳タスクの実験として、t5-largeモデルとtransformers GitHubリポジトリのexamples/seq2seq内にあるfinetune_trainer.pyスクリプトを使用して、小規模な微調整を行います。

テストには2つの24GB（Titan RTX）GPUを使用します。

これは概念実証のベンチマークであり、さらに改善できる余地があるため、比較のためにトレーニング用の2000アイテムと評価用の500アイテムの小さなサンプルでベンチマークを行います。評価はデフォルトでサイズ4のビームサーチを実行するため、同じサンプル数でトレーニングするよりも遅くなります。そのため、これらのテストでは評価アイテムは4分の1になります。

以下はベースラインの主なコマンドライン引数です：

export BS=16
python -m torch.distributed.launch --nproc_per_node=2 ./finetune_trainer.py \
--model_name_or_path t5-large --n_train 2000 --n_val 500 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro [...]

ベースラインのパフォーマンスを向上させるために、単にDistributedDataParallel（DDP）を使用しています。メモリエラー（OOM）が発生する前に、バッチサイズ（BS）16を適合させることができました。

このデモンストレーションのために重要なコマンドライン引数のみを表示して、理解しやすくするために簡略化しています。完全なコマンドラインはこの投稿で見つけることができます。

次に、ベンチマークを実行し、毎回以下のいずれかを追加していきます：

1. --fp16
2. --sharded_ddp（fairscale）
3. --sharded_ddp --fp16（fairscale）
4. --deepspeed（CPUオフロードなし）
5. --deepspeed（CPUオフロードあり）

ここでの主な最適化は、各技術がGPUメモリを効率的に展開するという点です。バッチサイズを継続的に増やし、トレーニングおよび評価がより速く完了することを期待します（メトリクスが安定したままであるか、いくつか改善されることには焦点を当てません）。

トレーニングと評価のステージは非常に異なるため、トレーニングではモデルの重みが変更され、勾配が計算され、オプティマイザの状態が保存されます。評価中はこれらの操作は行われませんが、特に翻訳のタスクでは、モデルが最適な仮説を探し出すために複数の実行を行わなければならないため、実際には複数の実行が必要です。そのため、モデルが大きい場合には高速ではありません。

これらの6つのテストランの結果を見てみましょう:

簡単にわかるように、FairScaleとDeepSpeedの両方がベースラインに比べて、総合的なトレーニング時間や評価時間だけでなく、バッチサイズでも大幅な改善を提供していますが、DeepSpeedはより多くのマジックを実装しており、現時点では短期的な勝者のようですが、Fairscaleはデプロイが容易です。DeepSpeedでは、簡単な設定ファイルを書き、コマンドラインのランチャーを変更する必要がありますが、Fairscaleでは--sharded_ddpのコマンドライン引数を追加するだけで済みますので、最も簡単な方法として最初に試してみることをお勧めします。

80:20の法則に従い、これらのベンチマークにはわずか数時間しか費やしておらず、コマンドライン引数や設定を細かく調整してMBや秒を最適化することはしていません。簡単な表から、次に試すべきことは明らかです。数時間や数日間実行される実プロジェクトに直面する場合には、確実により最適なハイパーパラメータを使用して仕事をより速く、最小限のコストで完了するために、より多くの時間を費やすことをお勧めします。

このベンチマークを自分で試したり、実行に使用されたハードウェアやソフトウェアの詳細を知りたい場合は、この投稿を参照してください。

Fairscaleは複数のGPUでのみパフォーマンスを向上させますが、DeepSpeedは単一のGPUでも利点を提供しています。

では、不可能なことに挑戦してみましょう – 24GBのRTX-3090カードでt5-3bをトレーニングしましょう。

まず、通常の単一GPUセットアップを使用して巨大なt5-3bを微調整してみましょう:

export BS=1
CUDA_VISIBLE_DEVICES=0 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 [...]

残念ながら、BS=1でも以下のエラーが発生します:

RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 23.70 GiB total capacity;
21.37 GiB already allocated; 45.69 MiB free; 22.05 GiB reserved in total by PyTorch)

前述のように、重要な部分のみを表示しています。完全なコマンドライン引数はこちらにあります。

次にtransformersをv4.2.0以上に更新し、DeepSpeedをインストールします:

pip install deepspeed

そして、再度試してみましょう。今回はDeepSpeedをコマンドラインに追加します:

export BS=20
CUDA_VISIBLE_DEVICES=0 deepspeed --num_gpus=1 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 --deepspeed ds_config_1gpu.json [...]

すると、20のバッチサイズで正常にトレーニングできます。おそらくさらに進めることもできるでしょう。プログラムはBS=30でメモリ不足により失敗しました。

以下に関連する結果を示します:

2021-01-12 19:06:31 | INFO | __main__ |   train_n_objs = 60
2021-01-12 19:06:31 | INFO | __main__ |   train_runtime = 8.8511
2021-01-12 19:06:35 | INFO | __main__ |   val_n_objs = 10
2021-01-12 19:06:35 | INFO | __main__ |   val_runtime = 3.5329

ベースラインとは比較できませんが、ベースラインは起動せずにすぐにメモリ不足で失敗します。

本当に素晴らしいですね！

私は主に24GBのGPUに収まらないこの巨大なモデルでトレーニングと評価ができるようになるかどうかに興味があったので、わずかなサンプルのみを使用しました。

このベンチマークを自分で試したり、実行に使用されたハードウェアやソフトウェアの詳細を知りたい場合は、この投稿を参照してください。

transformersはこれらの素晴らしいソリューションを統合しただけであり、その発明の一部ではなかったため、詳細を自分自身で発見できるリソースを共有します。ただし、以下にいくつかのクイックな洞察があります。これらの素晴らしい偉業をどのように達成するかを理解するのに役立つでしょう。

ZeROの主な特徴は、データ並列トレーニングの馴染みのあるコンセプトに分散データストレージを追加することです。

各GPU上の計算は、データ並列トレーニングとまったく同じですが、パラメータ、勾配、およびオプティマイザの状態は、分散/パーティションされた方法ですべてのGPUに格納され、必要なときにのみ取得されます。

次の図は、このブログ記事から取得されたもので、これがどのように機能するかを示しています:

ZeROの巧妙なアプローチは、パラメータ、勾配、およびオプティマイザの状態をすべてのGPUに均等に分割し、各GPUに単一のパーティション（またはシャードとも呼ばれる）を与えることです。これにより、GPU間のデータストレージには重なりがなくなります。実行時には、各GPUは参加するGPUに対して必要な情報を送信することで、各レイヤーのデータを動的に構築します。

このアイデアは理解するのが難しいかもしれませんが、こちらで説明を試みています。

作成時点では、FairScaleとDeepSpeedはオプティマイザの状態と勾配のパーティショニング（シャーディング）のみを実行しています。DeepSpeedとFairScaleでは、モデルパラメータのパーティショニングも近々行われる予定です。

もう1つの強力な機能は、ZeRO-Offloadです。この機能は、一部の処理とメモリの必要性をホストのCPUにオフロードすることで、より多くをGPUに適合させることができます。24GBのGPUでのt5-3bの成功にその劇的な影響を見ることができました。

PyTorchのフォーラムでは、多くの人がGPUメモリの断片化について苦情を言っています。以下のようなOOMエラーが頻繁に発生します:

RuntimeError: CUDA out of memory. Tried to allocate 1.48 GiB (GPU 0; 23.65 GiB total capacity;
16.22 GiB already allocated; 111.12 MiB free; 22.52 GiB reserved in total by PyTorch)

プログラムは約1.5GBを割り当てたいとしていますが、GPUにはまだ6-7GBの未使用メモリがありますが、連続した空きメモリは約100MBしかなく、OOMエラーが発生します。これは、異なるサイズのチャンクが繰り返し割り当てられ、解放されるため、時間の経過とともに所望のサイズの連続したチャンクが作成されず、メモリの断片化が起こるためです。上記の例では、プログラムはおそらく100MBの連続したメモリを割り当てることができますが、明らかに1.5GBの単一のチャンクを取得することはできません。

DeepSpeedは、GPUメモリを自ら管理し、長期的なメモリ割り当てが短期的なものと混ざらないようにし、断片化が少なくなるようにします。論文では詳細には触れていませんが、ソースコードは利用可能なので、DeepSpeedがどのようにそれを達成しているかを確認することができます。

ZeROはZero Redundancy Optimizerの略であるため、その名前にふさわしいものであることは簡単に理解できます。

DeepSpeedでは、モデルの変更は必要ありません。必要な変更は、トレーニングコードにのみあります。

これらのプロジェクトの統合部分に問題が発生した場合は、transformersのIssueを開いてください。

ただし、DeepSpeedとFairScaleのインストール、設定、展開に問題がある場合は、各ドメインの専門家に質問する必要がありますので、DeepSpeed IssueまたはFairScale Issueを使用してください。

これらのプロジェクトのいずれかがどのように動作するかを理解する必要はありません。単にtransformers Trainerを使用して展開することができますが、なぜそのようになるのかを理解したい場合は、以下のリソースを参照してください。

• FairScale GitHub

• DeepSpeed GitHub

• 論文：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models。この論文は非常に興味深いですが、非常に簡潔です。

• 以下のビジュアルを使用した論文の良いビデオディスカッションがあります。

• 論文：ZeRO-Offload: Democratizing Billion-Scale Model Training。最近発表されました。これはZeRO Offload機能の詳細について説明しています。

• DeepSpeedの設定とチュートリアル

• 論文に加えて、次の詳細なブログ記事とダイアグラムを読むことを強くお勧めします。

  • DeepSpeed: Extreme-scale model training for everyone
  • ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters
  • Turing-NLG: A 17-billion-parameter language model by Microsoft

• DeepSpeedのGitHub上の例

私たちは、FairScaleおよびDeepSpeedの開発チームから受けた素晴らしいサポートのレベルに非常に驚いています。これらのプロジェクトをtransformersに統合する作業中に私たちが遭遇した問題の迅速な解決と寛大なサポートを提供してくれた、FairScaleチームのBenjamin Lefaudeux @blefaudeuxとMandeep Baines @msbaines、そしてDeepSpeedチームのJeff Rasley @jeffra、Olatunji Ruwase @tjruwase、Samyam Rajbhandari @samyamに特に感謝したいと思います。

また、ベンチマークを実行するためのハードウェアへのアクセスを提供してくれたHuggingFaceにも感謝しています。

Sylvain Gugger @sguggerとStas Bekman @stas00は、これらのプロジェクトの統合に取り組みました。

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