PyTorch機械学習モデルのGPUメモリ消費の要因を細かくトラックする。

はじめに

　本記事では今更ながら、この世で一番シンプルな機械学習モデル（全結合1層）を使ってMNISTの分類課題を行い、訓練の各過程でGPU上にどれくらいメモリが割り当てられるかを確認、解釈を行った。
　pytorch_memlab(あるいはtorch.cuda.memory_allocated(device))で表示される、PyTorchのテンソルによって占有されているGPUメモリ消費に関して説明が可能だったので、それを雑多にまとめ、検証している。

実験

条件

PyTorch == 1.13.1
GPU : NVIDIA Tesla A30
CUDA version == 12.1

MNISTデータセットを一層の全結合層で分類。
最適化手法はSGD。
以下のmain.py（これはmodels.pyをimportして動くファイル)をpython main.pyで実行し、分類を行った。
今回は訓練1エポック目の一番最初のバッチで重みの更新が行われる所だけに注目し、この時GPU(1台)上でのメモリの立ち上がり方を追跡した。

main.py

 1 # import library
 2 import torch
 3 import torch.nn as nn
 4 from torchvision.datasets import MNIST
 5 import torchvision.transforms as transforms
 6 from torch.utils.data import random_split
 7 from torch.utils.data import DataLoader
 8 import torch.optim as optim
 9
10 # import model
11 from models import *
12
13 def main():
14     # data prepare
15     mnist_train_data = MNIST('path/to/data', train=True, download=True, transform=transforms.ToTensor())
16     train_size = int(len(mnist_train_data)*0.8)
17     mnist_train_data, mnist_validation_data = random_split(mnist_train_data, [train_size, len(mnist_train_data) - train_size])
18     mnist_test_data = MNIST('path/to/data', train=False, transform=transforms.ToTensor())
19     train_loader = DataLoader(mnist_train_data, batch_size=100, shuffle=True)
20     valid_loader = DataLoader(mnist_validation_data, batch_size=100, shuffle=True)
21     test_loader = DataLoader(mnist_test_data, batch_size=100, shuffle=True)
22
23     # gpu setup
24     if torch.cuda.is_available():
25        gpu_id = 0
26        device = torch.device('cuda:{}'.format(gpu_id))
27
28    # training setup
29    image_size = 28*28
30    model = simple_net_mnist(image_size, 10).to(device)
31    optimizer = optim.SGD(model.parameters(), lr=0.01)
32    criterion = nn.CrossEntropyLoss()
33
34    # train
35    for epoch in range(100):
36        train(model, device, train_loader, optimizer, criterion, epoch, image_size)
37
38 # train function
39 def train(model, device, train_loader, optimizer, criterion, epoch, image_size):
40    model.train()
41    running_loss = 0.0
42    for batch_idx, data in enumerate(train_loader):
43        inputs, labels = data
44        inputs, labels = inputs.to(device), labels.to(device)
45        optimizer.zero_grad()
46        inputs = inputs.view(-1, image_size)
47        output = model(inputs)
48        loss = criterion(output, labels)
49        loss.backward()
50        optimizer.step()
51        running_loss += loss.item()
52        if batch_idx % 480 == 479:
53            print(f'[{epoch+1}, {batch_idx + 1:5d}] loss : {running_loss/2000:.3f}')
54            running_loss = 0.0
55
56 if __name__ == '__main__':
57    main()

models.py

 1 import torch.nn as nn
 2 
 3 class simple_net_mnist(nn.Module):
 4    def __init__(self, input_dim, output_dim):
 5        super(simple_net_mnist, self).__init__()
 6        self.fc = nn.Linear(input_dim, output_dim)
 7   
 8    def forward(self, x):
 9         out = self.fc(x)
10         return out

やったこと

　上記main.py, models.pyの各行の前後でテンソルによって占有されているGPUメモリ消費の変動を確認。pytorch_memlabかtorch.cuda.memory_allocated(device)を使えば良い。どこでどのくらいメモリが割り当てられているかを確認した。

結果

前後でGPU上での割り当てメモリが増えた行を全て列挙すると以下の通り。数値の単位はByte。

1. main.py 30行目 : model = simple_net_mnist(~).to(device)
   0 --> 32,256 (+32,256)

2. main.py 44行目 : inputs, labels = inputs.to(device), ~
   32,256 --> 347,136 (+314,880)

3. main.py 47行目 : output = model(inputs)
   347,136 --> 351,232 (+4,096)

4. main.py 48行目 : loss = criterion(output, labels)
   351,232 --> 356,352 (+5,120)

5. main.py 49行目 : loss.backward()
   356,352 --> 384,000 (+27,648)

6. models.py 9行目 ： out = self.fc(x)
   347,136 --> 351,232 (+4,096)

これより、最大で384,000 Byteが割り当てられていたことがわかる。

解釈

　初めに、本モデル( simple_net_mnist )のデータの流れを簡単にまとめる。
モデルへの入力にはバッチサイズ(100) × 28 × 28の画像が入り、これをmain.py46行目でreshapeし、バッチサイズ × 784に変換する。その後、全結合層（784 to 10）を経てバッチ内の各画像で10値の予測確率が出力される。この図を踏まえ、実際GPU上で確保されているメモリ量について解釈を行う。

1. main.py 30行目 ( model = simple_net_mnist(~).to(device) ) : +32,256

　ここではモデル( simple_net_mnist )のパラメータ数分のメモリが確保されている。このモデルのパラメータ数は784 to 10の変換を行う全結合層であるため、重み＋バイアスで784 × 10 + 10で、floatは４Byteであるため、

(784 \times 10 + 10)\times 4 = 31,760 

より31,760 Byteが確保されていることがわかる。これは実際の確保分32,256 Byteより496 Byte少ないものの、実際の確保メモリの 98.5 %の大きさ($31,760/32,256\approx0.985$より)で見積もることができている。

2. main.py 44行目 ( inputs, labels = inputs.to(device), ~ ) : +314,880

　ここは単純にデータを転送しているだけ。入力データ = バッチサイズ(100) × 28 × 28 とラベル = バッチサイズ × 10 それぞれfloat, intでともに４Byteであるため、

(100 \times 28 \times 28 + 100) \times 4 = 314,000

より314,000Byteが確保されていることがわかる。これは実際の確保分314,880より880 Byte少ないものの、実際の確保メモリの 99.7 %の精度($314,000/314,880\approx0.997$より)で見積もることができている。

3. main.py 47行目 ( output = model(inputs) ) ＝ 6. models.py 9行目( out = self.fc(x) ) (両者は同じ) : +4,096

　ここはforwardを流している。forwardの中間出力(今回は1層のため、ここでは単純にforwardを経て得られた出力値=バッチサイズ×10)を保存し、

100 \times 10 \times 4 = 4,000

より4,000 Byteが確保されていることがわかる。実際の確保分4,096より96 Byte少ないものの、実際の確保メモリの 97.7 %の精度($4,000/4,096\approx0.977$より)で見積もることができている。

4. main.py 48行目 ( loss = criterion(output, labels) ) : +5,120

　ここはcross entropy誤差を求めている。既にforwardを経て得られたoutputの情報もlabelの情報もあるはずなので、更に追加で必要なメモリとは一体何なのだろうか。一つは1バッチ内の各画像のロスを入れる箱で、これはバッチサイズ × 4 = 400(Byte)だろう。
　また、cross entropyを求める式は、真の分布（今回だとone-hotの10個の要素からなる分布）を$p(x)$, 予測分布（今回だと10個の実数の要素からなる分布）を$q(x)$とするとcross entropy $H(p, q)$は

H(p, q) = -\sum_x p(x) \log (q(x))

なる式で表される。したがって、以下二つの可能性が考えられる。

可能性1：元のラベルはintで入っていたが、cross entropyを求めるために内部でone-hotに変換され、その分のメモリが確保されている可能性がある。

可能性2：モデルから出力されるoutputsの各予測値に対してlog (q(x))を算出しなければならない。この分のメモリが中間出力とは別に確保されている可能性がある。

いずれにせよ、モデルから出力されるoutputsと同じメモリを最初の400（Byte）と足し合わせて確保する必要があり、

400 + 100 \times 10 \times 4 = 4,400

より4,400Byteが確保されていることがわかる。実際の確保分5,120より720 Byte少ない。これより今回の見積もりは、実際の確保メモリの 85.9 %程度($4,400/5,120\approx0.859$より)であることがわかる。（見積もりの精度が低いので、見落としや解釈間違いがありそう..）

5. main.py 49行目 ( loss.backward() ) : +27,648

　ここはbackwardを計算している。今回最適化アルゴリズムとしてSGDを使ったので、勾配の情報さえもとまれば重みの更新が可能となる。勾配の情報は元の重みとは別のコピー配列を用意して格納する必要があり、（コピーしないと勾配か元の重みの値のどちらかの情報が消えてしまう）モデルのパラメータ数分を新たに別で確保する。また、誤差が一度求まってしまえば中間出力の値の保持は不要であるため、その削除が行われる。これらを総合すると、

(784 \times 10 + 10) \times 4 - 100 \times 10 \times 4 = 27,400

より27,400 Byteが確保されたことになる。この見積もりは、実際の確保分27,648 Byteよりも248 Byte少ないものの、 実際の確保メモリの99.1 %の大きさ($27,400/27,648\approx0.991$より)で見積もれていることがわかる。

解釈まとめ

　ここまで、典型的な機械学習モデルで段階ごとにGPU上でどれ程のメモリが確保されているかを見積もった。その見積りの精度はほとんどの段階で 97 % ~ 99 % だったことから、少なくともSGDを使った学習では、今回行った解釈でそれ程大きく外してはいないはずである。今回の解釈はまとめると以下の4段階になる。

段階0. 初期状態 : 0
段階1. モデルをロード : + モデルのパラメータ数 (main.py 30行目)
段階2. フォワード : + バッチサイズ × ( 入力データ + 中間出力(=中間層から出力された値, 最終層から出力された値outputs) + 誤差取得用のoutputs(?) )　　(main.py 44, 47, 48行目)
段階3. バックワード : + モデルのパラメータ数(段階1.とは別)を確保 & 中間出力を削除 (main.py 49行目)

以上をまとめると、MNISTのSGD学習で、今回の確認方法で表示された最大GPU割り当てメモリ/Byteは段階3時点での割り当てメモリ

maxparam = ( モデルのパラメータ数 $\times$ 2 + バッチサイズ $\times$ (入力データ + outputs) ) $\times$ 4

で近似できる事になる。ただしこの式は普遍的なものではなく、中間層の次元やモデルのサイズのバランス次第ではフォワードで最大GPU割り当てメモリとなる可能性もあるため、その場合には式の形も変わってくるだろう。

　今回の例だと、モデルのパラメータ数=$784\times 10 + 10$、バッチサイズ=$100$、入力データ=$784$, outputs=$10$より、

{\rm max param} = ((784\times 10 + 10) \times 2 + 100 \times (784 + 10) ) * 4 = 380,400

となり、実際の割り当ては最大で$384,000$ Byteであったため$380,400/384,000 = 0.991$より 99.1 % の精度で確保メモリの見積もりができていることがわかる。

検証

1. nvidia-smiコマンドで表示されるメモリ使用量に洞察を与えられないか？

　メモリ使用量を確認する際に手軽でよく登場するコマンドとしてnvidia-smi(あるいはnvitop)が挙げられる。ただしこのコマンドは非テンソル占有メモリもまとめて表示され、今回の解釈をストレートには適用できない。
　そこでここでは、テンソル占有メモリに影響が及ぶハイパラだけを変えたときに占有メモリがどう変化するかを確認した。具体的には先ほどと全く同じ全結合一層のアーキテクチャ、データはMNIST、最適化手法はSGDで揃え、バッチサイズだけを様々に変えてnvidia-smiで表示される使用メモリ量/MiBをプロットした。その結果、以下のような図が得られた。

このプロットは凡そ線形になっており、傾きは3,212 Byte / 1 Batchだった。この傾きを解釈まとめで考察したmaxparamの式と紐付けると、傾き(Byte / 1 Batch)は "(入力データ + outputs）× 4" に相当することが分かる。この値を算出すると、

(784 + 10) \times 4 = 3,176

となることから、限りなくnvidia-smiコマンドで表示されるメモリ使用量のバッチサイズに対する変化率3,212 Byte / 1 Batchを近似できており、解釈まとめで示したmaxparamの式の一定の妥当性が示唆された。

　なお、線形とは言いつつ切片の値が異常に高くなっていそうだという点については、CUDAの初期化時点で色々なカーネルをロードするCUDAコンテキストが作成されることによるものらしい。（ここら辺の非テンソル占有メモリについてはよくわからなかったため撤退した。他にも原因はありそうだが、一因と思しきリンクのdiscussionをご参照頂ければ幸いである。）

2-1. Adamではどうか？

　SGD以外の最適化アルゴリズムでも各行の前後でtorch.cuda.memory_allocated(device) or pytorch_memlabによりメモリを確認した。コードmain.pyの31行目を

main.py

~  ...
31 optimizer = optim.Adam(model.parameters(), lr=0.01)
~  ...

に置き換えて同じ実験を行った結果、

・GPU上で確保されているメモリはSGDの時とmain.py49行目までは同じ結果だった。

・ただし、main.py50行目に記述されているoptimizer.step()で384,000 --> 448,512へと+64,512 Byte分のメモリ上昇が確認された。

2-2. Adam実験の解釈は？

　この上昇分はAdamで必要となる勾配情報の1次のモーメントと2次のモーメントのメモリ確保によるものだと考えられ、モデルのパラメータ数を2倍した式

(784 \times 10 + 10) \times 4 \times 2 = 62,800

により、62,800 Byteと見積もられた。この見積もりは実際に確保されたメモリ64,512 Byteの 97.3 % という精度であった。このことから、loss.backward時点まではSGDと同じ挙動をしているが、更新optimizer.step()の際に勾配情報に加えて新しく1次・2次のモーメント用のメモリが割り当てられ、これらを使って重みの更新が行われるということが示唆された。

結論

　今回確認、検証をしてみて、機械学習モデルのGPUへのメモリ割り当ては大きく以下の4段階で行われると解釈された。これらの解釈で見積もられたGPU割り当てメモリサイズは実際の割り当てメモリサイズの97 % ~ 99 %であった。

・段階1：モデルをロード。
・段階2：フォワードでデータを流す。この時中間層からの出力の分もメモリを確保。
・段階3：バックワードで勾配を取得。そして中間出力の分のメモリは削除。
・段階4：重み更新。SGDでは新しいメモリ確保は不要。Adamなどの他の最適化アルゴリズムではモデルのパラメータ数×n個分のメモリを新しく確保。
　
　今回扱った1層の全結合モデルでSGDを用いてMNISTを分類する例では、GPU割り当てメモリは段階３で最大となり、既出のmaxparamの式で近似ができた。いつ割り当てメモリが最大となるかはモデルと中間出力の次元の両方により決まるため、それ次第では他の段階時点での割り当てメモリの近似式を用いるべきだと考えられる。

参考文献

・https://pytorch.org/docs/stable/cuda.html
・https://runebook.dev/ja/docs/pytorch/cuda?page=2
・https://medium.com/deep-learning-for-protein-design/a-comprehensive-guide-to-memory-usage-in-pytorch-b9b7c78031d3