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