見出し画像

DeepSpeedの紹介

株式会社レトリバ

※こちらの記事は、2021年7月26日にRetrieva TECH BLOGにて掲載された記事を再掲載したものとなります。

Chief Research Officerの西鳥羽 (Jiro Nishitoba (@jnishi) | Twitter) です。
前回のブログでBigBirdを触ってみたを予告してましたが、BigBirdのような巨大なモデルを学習するために有用なライブラリがあったので、先にそちらを紹介したいと思います。
皆様は最近のモデルをみて、「お、いいな」と思うものの学習環境で16GPUとか64GPUなどの記述を見つけてしまい、遠い目をしながらそっ閉じした経験などありませんでしょうか。
今回紹介するDeepSpeed というライブラリは、物理メモリや外部SSDなどを活用してより大きなモデルを学習できるようにするものです。
実際GPUメモリ24GBのGeForce TitanRTX 2台でbaseサイズのBERTがほぼ同等の条件で学習できます。


DeepSpeed

DeepSpeedは「巨大なモデルの学習を容易にする」ことをコンセプトにしたライブラリです。
Deep Learningの学習における並列化というと、Data ParallelとModel Parallelの2つがあります。Data Parallelは計算効率は良い一方で必要なメモリ量が多く、計算機資源を増やした時に計算できるモデルの大きさのスケーラビリティがあまり良くないことが知られています。Model Parallelはメモリの使用効率が良く、計算機資源を増やした時に計算できるモデルの大きさは良いのですが、一方で必要な通信量が多くなってしまうため、計算効率はあまり良くありません。そこでData Parallelの計算効率を維持しながら、計算に用いないパラメータを削除することによってメモリの使用効率を上げる方法がZero Redundancy Model(ZeRO)です。また、ZeROでパラメータを削減すると同時に、GPUのメモリを必要とする計算の一部をCPU及びメインメモリで計算する(CPU-Offload)ことに依ってGPUのメモリを節約するZeRO-Offloadがあります。
ZeRO
ZeROではData Parallelのようにforward、backwardを行います。それによりモデルのパラメータは全てのGPUでもつ必要がありますが、学習データの転送以外の通信を必要とせずそれぞれのGPU内だけでforwardとbackwardの計算が行なえます。そしてその後にModel ParallelのようにそれぞれのGPUで分担してパラメータの更新を行います。それにより、後述するようにGPUのメモリを節約できます。
ZeROではoptimizerのstate(Adamにおけるmomentumやvarianceなど)、backward時に計算する勾配、モデルのパラメータを削除します。

図1: ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelsより引用

optimizerのstateはパラメータの更新時にしか用いません。GPUを複数台用いる場合、パラメータ更新は分担して行われます。そのため、図のPosのようにそのGPUが更新を担当するパラメータ以外のoptimizer stateは用いられないので削除できます。これによりGPUのメモリを節約できます。
backwardで計算する勾配はパラメータの更新に使われます。そのため、計算された勾配はそのGPUが担当するパラメータ以外のものは不要となりますので、図のPos+gのように勾配を必要とするGPUに転送してしまえば削除することが出来ます。これによりGPUのメモリを節約できます。optimizerのstateの削除と比較して、追加の通信や計算が発生することはありません。ここまでメモリを削減する機能をDeepSpeedではstage 2としており、ZeRO-2と銘打っているのもここまでのメモリ削減を行っているもののようです。
最後にparameterの削除になります。図のPos+g+pのようにそのGPUが担当するパラメータ以外のパラメータは削除することが出来ます。これにより、更にGPUのメモリが削減できます。ただ、forward及びbackwardの計算には他のパラメータが必要になるので、計算のたびにパラメータをbroadcastして共有します。このため、データの通信量は増加します。計算は省略しますが、上記optimizer state及び勾配のみの削減時に比べて 1.5倍の通信量になります。ここまでメモリを削減する機能をDeepSpeedではstage 3としており、ZeRO-3と銘打っています。

ZeRO-Offload

ZeRO-Offloadはパラメータの更新をCPU及びメインメモリ上で行うことにより、optimizer stateをGPUメモリから削減します。

図2: ZeRO-Offloadのコンセプト(ZeRO-Offload: Democratizing Billion-Scale Model Trainingより引用)

通常は図2の左のようにパラメータ、勾配、optimizer stateが全てGPUのメモリ上にあることを必要とします。しかし、パラメータの更新をCPU上で行うようにすると図2の右のようにGPUのメモリに必要となるのはパラメータと勾配の一部となります。勾配は計算が終わったら順次メモリに転送してしまえば良いので、全部を保持する必要はありません。これによりGPUのメモリを削減できます。

また、パラメータの更新をCPUで行うことになり、計算速度の劣化が気になるところかもしれませんが、CPUでの処理の高速化のための工夫もしています。例えばSIMD命令による並列化、loop unrolling、OMPのmultithredingによるコア並列化などです。これにより、学習の1イテレーションあたりではほぼ同等の速度で実行できるようになっています。

その他の機能

NVMe接続したSSDを用いてより巨大なモデルを学習できるようにしたZeRO-Infinityや量子化による通信量の削減、attentionのsparse化による高速化 などがあり、それら全てDeepSpeedから使用することが出来ます。

Huggingface Transformersとのインテグレーション

Huggingface TransfomersではDeepSpeedとの連携が進んでおり、手軽に用いることが可能です。詳細はこちらを参照してもらえればと思いますが、基本的には下記の3点の変更で使うことが出来ます。

  • pip install deepspeed でDeepSpeedをインストールする

  • DeepSpeedの設定ファイルを `ds_config.json` というファイル名で用意する、こちらに掲載されているサンプルをほぼそのまま用いることが出来ます

  • 実行時のコマンドを python run_mlm.py などから deepspeed run_mlm.py など、deepspeedコマンドに変更し、最後に --deepspeed ds_config.json というオプションを付ける

実験

BERTのbaseサイズを学習してみました。元論文ではTraining of BERTBASE was performed on 4 Cloud TPUs in Pod configuration (16 TPU chips total).とあり、おそらく2TB程度のメモリ(*1)を必要としたと思われます。

Huggingface/TransformersのDeepSpeedインテグレーションにて下記の機能を用い、BERTのbaseサイズの事前学習を行いました。これにより、GeForce Titan RTX(GPU メモリ 24GB) 2枚の計GPUメモリ48GBで学習できています。

  • ZeRO-2 (optimizer state及び勾配の削減)

  • CPU-Offload

  • fp16(半精度浮動小数点)による計算

  • 系列長を512から128に短縮

完全に同じというわけにはいきませんでしたが、小さな構成のサーバーでも学習ができるようになりました。
なお、系列長に関してはGoogle ResearchのBERTのレポジトリにも、始めは系列長128で学習し、最後に少しだけ系列長512で学習すると良い（*2)とあり、今回は系列長128、バッチサイズ256で500000 step学習した後、系列長512バッチサイズ32で5000 stepの学習を行いました。
学習したBERTモデルで livedoorニュースコーパス の分類を試してみたところ精度93.3%となりました。

まとめ

GPUのメモリ使用を効率化し、より巨大なモデルを学習できるようになるDeepSpeedを紹介しました。今までBERTの学習というとbaseサイズですらGPUクラスタが必要でしたが、DeepSpeedを使うと小規模なサーバーでも学習できるようになります。今回はZeRO-2までしか試していませんが、ZeRO-3及びZeRO-infinityなども用いるとより大きなモデルが学習できることが期待できます。

*1:TPU v3 が一台あたり128GiBのメモリを搭載しているため

*2:Longer sequences are disproportionately expensive because attention is quadratic to the sequence length. In other words, a batch of 64 sequences of length 512 is much more expensive than a batch of 256 sequences of length 128. The fully-connected/convolutional cost is the same, but the attention cost is far greater for the 512-length sequences. Therefore, one good recipe is to pre-train for, say, 90,000 steps with a sequence length of 128 and then for 10,000 additional steps with a sequence length of 512. The very long sequences are mostly needed to learn positional embeddings, which can be learned fairly quickly. Note that this does require generating the data twice with different values of max_seq_length Google ResearchのBERTのレポジトリ より