こんにちは、CCCMKホールディングス TECH LABの三浦です。
9月になりました。夏から秋に変わる時期で気候もコロコロ変わるので、体調管理をちゃんとしないと・・・と思います。
さて、LLMに限らずDeep LearningモデルをGPUを使って学習していると、「なんでこんなにGPUメモリを消費しているんだろう・・・？」と不思議に思うことがあります。特にモデルそのもののファイルサイズ上は十分にGPUメモリに読み込めるはずなのに、学習を始めるとGPUメモリの容量オーバー時に発生する"Out-Of-Memory"エラーが出てしまったりします。この理由について、最近調べていた分散学習に関する論文にまとまっていてとても勉強になりました。
それはMicrosoftが公開しているDeep Learningモデルの学習ライブラリ"DeepSpeed"の中に含まれる"ZeRO"というテクニックに関する論文です。
- Title
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - Author
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He - Submit
4 Oct 2019 (v1), last revised 13 May 2020 - arXiv
https://arxiv.org/abs/1910.02054
"Zero Redundancy Optimizer (ZeRO)"は従来の分散処理を実行する際に発生するGPUメモリ消費の最適化を行い、巨大なLLMの学習を可能にする技術です。この論文の中の"3 Where Did All the Memory Go?"の箇所に学習時に使用されるGPUメモリの内容について詳しく記述されています。
今回のブログでは、その内容について自身の理解のためにまとめてみたいと思います。
GPUメモリはどこで消費されているのか？
たとえば1.5Bパラメータで構成されるGPT-2のモデル自身のサイズはパラメータをfp16で保存した場合、およそ3GBになります。これは16bitが2byteであることから計算することが出来ます。ところがこのモデルは32GBのメモリを持つGPUでもそのままでは学習することが出来ません。モデル学習時にはモデルそのもの以外にもメモリに読み込まないといけないものが存在し、実はそちらの方にGPUメモリの大部分が消費されていることが分かります。では一体何の用途でメモリが使われているのでしょうか？その多くが"Optimizer States", "Gradients", "Parameters"の3つで構成される"Model States"に使われています。
Model States
Model Statesは"Optimizer States", "Gradients", "Parameters"で構成される学習中に前後の学習ステップで保持しておかなければならないモデルの状態に関するデータです。Optimizerとしてよく使用される"Adam"の場合、各パラメータに対する"momentum"と勾配(gradients)の"variance"を保持する必要があり、そのためのGPUメモリが必要になります。
実は巨大なモデルの学習には"Mixed-Precision Training"という学習アプローチが採用されていて、パラメータと勾配はfp16で保持する一方、効果的なパラメータ更新処理のため、Optimizer Statesはfp32(fp16の倍の容量)で保持する必要があります。
つまり"Mixed-Precision Training"を採用する場合、Model Statesとしてはfp16(2byte)でモデルのパラメータと勾配を、fp32(4byte)でOptimizer States(パラメータごとにmomentumとvariance of gradients)を保持するためのメモリが必要になるのです。これは総合すると元のモデルをfp16で保存したものの8倍のサイズにもなります。
Residual Memory Consumption
Model States以外にも学習中にGPUメモリが使われる用途があります。それらは"Residual Memory Consumption"と論文では言及されていて、"Activations", "Temporary buffers", "Memory Fragmentation"に分けることが出来ます。
Activationsは学習の順伝播の時に計算される各層の出力で、逆伝播の勾配計算時に使用するために保存されます。学習時のバッチサイズやトークンの長さにも影響され、たとえば1.5BのGPT-2の場合、1Kのトークン列、32のバッチサイズでなんと60GBものメモリが必要になるそうです。バッチサイズを減らすとOut-Of-Memoryが解消される経験をこれまで何回もしてきましたが、このActivationsが関係していたんだな・・・と勉強になりました。
Activationsに使われるメモリを減らすテクニックでよく取られるもので"Activation checkpointing"というものがあるそうです。このテクニックは別の論文で提案されたものですが、一部のActivationsをメモリ上に保存せずに逆伝播の際に再計算をする、という内容のようです。もちろん再計算にかかる時間は増加するものの、その影響を最小限にしながら効率的にメモリの使用が抑えられる工夫がされており、先ほどの例だと60GB必要だったGPUメモリがActivation checkpointingによって8GB程度まで削減できるそうです。しかし、Activation checkpointingを用いても100B程度のモデルになると先ほどと同じ条件で60GBのGPUメモリをActivationsで消費してしまうため、巨大なモデルを学習させる場合は別の何らかのメモリ対策が必要になります。
Temporary buffersは学習時に一時的に使用されるGPUメモリを指しますが、それなりに大きなサイズになるそうです。たとえばデータ並列の分散処理で複数のGPUで学習をする際に、それぞれのプロセスで計算した勾配を集約する"all-reduce"の処理を効率的に行うためなどに、Temporary buffersが発生するようです。
Memory Fragmentationは、大きなモデルの学習中にしばしば発生することがあるメモリのfragmentation(断片化)によって、メモリ容量はあるのにメモリにアクセス出来なくなってOut-Of-Memoryが発生してしまう現象です。場合によっては総メモリの30%程利用可能であってもこの現象によってOut-Of-Memoryが発生することもあるようです。
まとめ
今回はZeROの論文の中で、DLモデルの学習時にどんな用途でGPUメモリが使われているのかについて述べられている箇所について読んで理解した内容をまとめてみました。順伝播や逆伝播など、Deep Learningモデルの学習の基礎的な内容が結構頭から抜け落ちていて、少し苦労しながら調べました。私自身モデルを学習する機会はそこそこあるのですが、内部で行われる処理はライブラリ任せにしてしまっていたので、理解せずに使っていたな、と少し反省しました。
今回はまとめられなかったのですが、これらのGPUメモリの問題に対してZeROではどのようなアプローチで解決するのかについても調べていますので、そちらの内容は次回以降の記事でご紹介したいと思います。
