またGPUメモリ足りないや...
深層学習初学者の自分は
上司「とりあえずこれ使ってください(16GBメモリくらいの弱マシン)」
と言われて、いざ自分なりにgithubクローンして深層学習を始めてみると
RuntimeError: CUDA out of memory. Tried to allocate 220.00 MiB (GPU 0; 15.78 GiB total capacity; 13.16 GiB already allocated; 201.44 MiB free; 14.44 GiB reserved in total by PyTorch
みたいなエラーが出てきて、
僕「あれ、上司がこれ使えって言ったくらいだから動くはずなのに...」
僕「cuda、docker周りの設定ミスったか？」
僕「いや、でも初学者だから何か知らないメモリ周りの設定間違えてるのか...」
などと悶々としていて、最終的に調べてみると
上司「ごめんごめんメモリ足りんわ。もっといいマシン使っていいよ」
となることが多かったのはいい思い出です。が、実際のところ ある深層学習モデルをあるタスクについて動かすのには最低どれくらいメモリが必要なのか？ を見積もるのは最初の一番大事なタスクになるかと思います。マシンスペックが足りないだけの場合、時間が無駄になりますしね。1
必要なメモリの見積もりかた
僕は学習に最低限必要なメモリの概算は典型的な設定において次のように見積もっています：
必要なGPUサイズ ≒ 4bytes × (モデルのパラメータ数)
+ 8bytes × (optimizerの必要な記憶)
+ 4bytes × (モデルのパラメータにおける勾配)
+ 4bytes × (入力データをforwardする時に必要なメモリ) × (バッチサイズ)
+ α(pythonカーネル、他ソフト系、計1〜2GB)
≒ 4bytes × 4 × (モデルのパラメータ数)
+ 4bytes × (入力データサイズ) × (バッチサイズ)
+ 2GB
注１）一時的にこれを超えるメモリが必要になる可能性はあります。
注2）自分の経験則です！forward計算時に必要なメモリはモデルのアーキテクチャによって異なるので（例えば「分岐構造があるか？」「CNN使ってるか？」等によってデータサイズと必要な演算結果保持に必要なメモリサイズの関係は変わります。）
注3)実際にはこれより少ないメモリでも動かせる可能があります。
どうして上の式で見積もれるのかを以下で解説していきます。
必要なメモリの内訳
深層学習で必要なメモリの内訳は以下となります：
- モデルの重み
- Optimizerの状態
- 勾配
- フォワード計算結果の保持
- tempバッファ
- 関数用領域
それぞれについて必要なデータはを以下で見ていきます。
1. モデルの重み
≒ネットワークの重みパラメータのことです。数式だとよくWで表されるものですね。
- 4 bytes × パラメータ数
2. Optimizer計算に必要な状態のメモリ
Optimizerのアルゴリズムにより、保持しなければいけないデータ数は変わります。
- 4 bytes × パラメータ数 × 2(AdamWの場合更新前後の2状態を記録する必要がある)
3. 勾配
バックプロパゲーションに必要な数値の保持分です。
- 4 bytes × パラメータ数
4. フォワード計算結果の保持
あるバッチ数モデルに計算させる際の入力データです。batchサイズ、シーケンス長、hidden stateの大きさによって変わる部分です。一番ボトルネックになる部分ですね。
5. tempバッファ
複雑な計算などに必要な一時的なメモリ領域です。利用する側は特に気にする必要はないです。学習中などに一時的に使用率が上がってしまう原因でもあります。従って、使わなくなった変数などはメモリを解放することが望まれる部分です。
6. 関数用領域
最後にbeamサーチなどのライブラリ・ソフトウェアが使う領域です。こちらも利用する側は気にする必要がないでしょう。(pythonカーネルとかもここですかね？)
従って、GPUメモリの内訳としては
- データやモデルアーキテクチャによらず必要なモデルパラメータのデータと、そこで保持する必要のある勾配データ部分
- あるbatchサイズだけ、データをモデルにforwardさせる時に保存するデータ＆行った演算(一般に「計算グラフ」と言ってる部分)のデータを保持する部分
が必要となります。
ただし、後者については利用したいモデルによって違ったりする&学習のボトルネックになる可能性が高いので(≒batchサイズはでかい方が大体学習がうまくいく)、必要なGPUを見積もる際には
- batchサイズを1、データサイズも可能な限り小さくして動かしてみる。
- 動くのであればbatch・データサイズをN倍にしてメモリ使用率を見ることで、maxどれくらいのbatch・データサイズが試せそうか調べる。
- 動かないならそもそもそのマシンでは動かせないから他手段を考える。
というのが一番手っ取り早いと思います。
この見積もりより小さいGPUメモリでも動かせる可能性がある
以上、GPUメモリの内訳を説明しましたが、 現在pytorchやdeepspeedなどのフレームワークではGPUを節約して実際のGPUサイズより大きいサイズのモデルを扱ったり、計算時間を短くするための工夫が様々提案・実装されています。なので、理解が深まって「GPUたりんな〜」というときはhugging faceやpytorchのドキュメントを読むことをお勧めします（、というか自分が読まなきゃいけないと感じてる）。
キーワード：Gradiant Accumulation、Gradient Checkpointing、不動小数点型の変更、ZeRO、マルチGPU、マルチノードにおける並列化手法
-
(P.S. 本当は小さいデータセット作ってから検証進めてれば、メモリ不足であることはすぐ検証できたんですけどね笑。初学者の自分にはその発想すら思いついてませんでした。同じ悩みを持つ人が少しでも減ることを祈ります！) ↩