PyTorch顯存機(jī)制分析

作者最近兩年在研究分布式并行，經(jīng)常使用PyTorch框架。一開(kāi)始用的時(shí)候?qū)τ赑yTorch的顯存機(jī)制也是一知半解，連蒙帶猜的，經(jīng)常來(lái)知乎上來(lái)找答案，那么我就吸收大家的看法，為PyTorch的顯存機(jī)制做個(gè)小的總結(jié)吧。

1.1 深度學(xué)習(xí)訓(xùn)練過(guò)程

1. 模型定義：定義了模型的網(wǎng)絡(luò)結(jié)構(gòu)，產(chǎn)生模型參數(shù)；

2. 前向傳播：執(zhí)行模型的前向傳播，產(chǎn)生中間激活值；
3. 后向傳播：執(zhí)行模型的后向傳播，產(chǎn)生梯度；
4. 梯度更新：執(zhí)行模型參數(shù)的更新，第一次執(zhí)行的時(shí)候產(chǎn)生優(yōu)化器狀態(tài)。

1.2 前向傳播

1.3 后向傳播（反向傳播）

1.4 梯度更新

2.1 分析方法

這也是令很多人在使用PyTorch時(shí)對(duì)顯存占用感到困惑的罪魁禍?zhǔn)住?/span>(2) torch.cuda is all you need在分析PyTorch的顯存時(shí)候，一定要使用torch.cuda里的顯存分析函數(shù)，我用的最多的是torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()，前者可以精準(zhǔn)地反饋當(dāng)前進(jìn)程中Torch.Tensor所占用的GPU顯存，后者則可以告訴我們到調(diào)用函數(shù)為止所達(dá)到的最大的顯存占用字節(jié)數(shù)。還有像torch.cuda.memory_reserved()這樣的函數(shù)則是查看當(dāng)前進(jìn)程所分配的顯存緩沖區(qū)是多少的。memory_allocated+memory_reserved就等于nvidia-smi中的值啦。非~常~好~用chao dashengTorch 官方文檔2.2 PyTorch context開(kāi)銷(xiāo)-----之前沒(méi)有提到PyTorch context的開(kāi)銷(xiāo)，做個(gè)補(bǔ)充...我注意到有很多同學(xué)在做顯存分析的時(shí)候是為了在訓(xùn)練的時(shí)候可以把卡的顯存用滿，這個(gè)之前沒(méi)有考慮到呢。其實(shí)PyTorch context是我們?cè)谑褂胻orch的時(shí)候的一個(gè)大頭開(kāi)銷(xiāo)。主要參考的是論壇里的這篇討論：How do I create Torch Tensor without any wasted storage space/baggage?https://discuss.pytorch.org/t/how-do-i-create-torch-tensor-without-any-wasted-storage-space-baggage/131134什么是PyTorch context? 其實(shí)官方給他的稱(chēng)呼是CUDA context，就是在第一次執(zhí)行CUDA操作，也就是使用GPU的時(shí)候所需要?jiǎng)?chuàng)建的維護(hù)設(shè)備間工作的一些相關(guān)信息。如下圖所示這個(gè)值跟CUDA的版本，pytorch的版本以及所使用的設(shè)備都是有關(guān)系的。目前我在ubuntu的torch1.9上測(cè)過(guò)RTX 3090和V100的context 開(kāi)銷(xiāo)。其中3090用的CUDA 11.4，開(kāi)銷(xiāo)為1639MB；V100用的CUDA 10.2，開(kāi)銷(xiāo)為1351MB。感興趣的同學(xué)可以在shell中執(zhí)行下面這兩行代碼，然后用nvidia-smi去看看自己的環(huán)境里context的大小。然后用總大小減去context的大小再做顯存分析。

importtorch temp=torch.tensor([1.0]).cuda()我估計(jì)會(huì)有人問(wèn)怎么去減小這個(gè)開(kāi)銷(xiāo)...官方也給了一個(gè)辦法，看看自己有哪些cuda依賴是不需要的，比如cuDNN，然后自己重新編譯一遍PyTorch。編譯的時(shí)候把對(duì)應(yīng)的包的flag給設(shè)為false就好了。我是還沒(méi)有試過(guò)，要搭編譯的環(huán)境太難受了，而且還要經(jīng)常和庫(kù)做更新。

2.3Torch顯存分配機(jī)制

在PyTorch中，顯存是按頁(yè)為單位進(jìn)行分配的，這可能是CUDA設(shè)備的限制。就算我們只想申請(qǐng)4字節(jié)的顯存，CUDA也會(huì)為我們分配512字節(jié)或者1024字節(jié)的空間。

2.4Torch顯存釋放機(jī)制

在PyTorch中，只要一個(gè)Tensor對(duì)象在后續(xù)不會(huì)再被使用，那么PyTorch就會(huì)自動(dòng)回收該Tensor所占用的顯存，并以緩沖區(qū)的形式繼續(xù)占用顯存。要是實(shí)在看緩沖區(qū)不爽的話，也可以用torch.cuda.empty_cache()把它歸零，但是程序速度會(huì)變慢哦03 訓(xùn)練過(guò)程顯存分析為了讓大家方便理解，我這里用torch.nn.Linear(1024, 1024, bias=False) 來(lái)做例子。為了省事，loss函數(shù)則直接對(duì)輸出的樣本進(jìn)行求和得到。沒(méi)辦法，想直接執(zhí)行l(wèi)oss.backward()的話，loss得是標(biāo)量才行呢。示例代碼：

import torch model = torch.nn.Linear(1024,1024, bias=False).cuda() optimizer = torch.optim.AdamW(model.parameters()) inputs = torch.tensor([1.0]*1024).cuda() # shape = (1024) outputs = model(inputs) # shape = (1024) loss = sum(outputs) # shape = (1) loss.backward() optimizer.step()

3.1 模型的定義

結(jié)論：顯存占用量約為參數(shù)量乘以4

import torch model = torch.nn.Linear(1024,1024, bias=False).cuda() print(torch.cuda.memory_allocated())打印出來(lái)的數(shù)值為4194304，剛好等于1024×1024×4。

3.2 前向傳播過(guò)程

結(jié)論：顯存增加等于每一層模型產(chǎn)生的結(jié)果的顯存之和，且跟batch_size成正比。

inputs = torch.tensor([1.0]*1024).cuda() # shape = (1024) memory + 4096 outputs = model(inputs) # memory + 4096代碼中，outputs為產(chǎn)生的中間激活值，同時(shí)它也恰好是該模型的輸出結(jié)果。在執(zhí)行完這一步之后，顯存增加了4096字節(jié)。(不算inputs的顯存的話)。

3.3 后向傳播過(guò)程

后向傳播會(huì)將模型的中間激活值給消耗并釋放掉掉，并為每一個(gè)模型中的參數(shù)計(jì)算其對(duì)應(yīng)的梯度。在第一次執(zhí)行的時(shí)候，會(huì)為模型參數(shù)分配對(duì)應(yīng)的用來(lái)存儲(chǔ)梯度的空間。

loss = sum(outputs) # memory + 512(torch cuda分配最小單位) temp = torch.cuda.memory_allocated() loss.backward() print(torch.cuda.memory_allocated() - temp) # 第一次增加4194304第一次執(zhí)行時(shí)顯存增加：4194304字節(jié) - 激活值大小；第二次以后執(zhí)行顯存減少：激活值大??；Note：由于這個(gè)中間激活值被賦給了outputs，所以后面在后向傳播的時(shí)候會(huì)發(fā)現(xiàn)，這個(gè)outputs的顯存沒(méi)有被釋放掉。但是當(dāng)層數(shù)變深的時(shí)候，就能明顯看到變化了。為了讓大家看到變化，再寫(xiě)一段代碼~

import torch # 模型初始化 linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304 print(torch.cuda.memory_allocated()) linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096 print(torch.cuda.memory_allocated()) # 輸入定義 inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304 print(torch.cuda.memory_allocated()) # 前向傳播 loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512 print(torch.cuda.memory_allocated()) # 后向傳播 loss.backward() # memory - 4194304 + 4194304 + 4096 print(torch.cuda.memory_allocated()) # 再來(lái)一次~ loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512沒(méi)了，因?yàn)閘oss的ref還在) print(torch.cuda.memory_allocated()) loss.backward() # memory - 4194304 print(torch.cuda.memory_allocated())

3.4 參數(shù)更新

optimizer.step()#第一次增加8388608，第二次就不增不減了哦第一次執(zhí)行時(shí)，會(huì)為每一個(gè)參數(shù)初始化其優(yōu)化器狀態(tài)，對(duì)于這里的AdamW而言，每一個(gè)參數(shù)需要4*2=8個(gè)字節(jié)。第二次開(kāi)始，不會(huì)再額外分配顯存。顯存開(kāi)銷(xiāo):第一次: 增加8388608字節(jié)第二次及以后: 無(wú)增減3.5 Note由于計(jì)算機(jī)計(jì)算的特性，有一些計(jì)算操作在計(jì)算過(guò)程中是會(huì)帶來(lái)額外的顯存開(kāi)銷(xiāo)的。但是這種開(kāi)銷(xiāo)在torch.memory_allocated中是不能被察覺(jué)的。比如在AdamW在進(jìn)行某一層的更新的時(shí)候，會(huì)帶來(lái)2倍該層參數(shù)量大小的臨時(shí)額外開(kāi)銷(xiāo)。這個(gè)在max_memory_allocated中可以看到。在本例中就是8388608字節(jié)。

審核編輯 ：李倩