PyTorch顯存機制分析——顯存碎片問題

參考前文：

https://www.cnblogs.com/devilmaycry812839668/p/15578068.html

 

 

====================================

 

 

 

從前文我們知道pytorch中是分層進行管理顯存的，最小的管理單位是512B，然后上一層是2MB，那么如果我們按照這個原理寫一個碎片化的顯存分配，就可以實現2GB數據占4GB的顯存空間的操作。

 

現有顯存：

 

 

 

 

 

運行代碼：

import torch
import time

device = "cuda:0"

tensor_ = torch.randn(5*256*1024*1024, device=device)
print(torch.cuda.memory_summary())
time.sleep(60000)

 

申請5G顯存，報錯：

 

 

 

 

 

 

更改代碼：

import torch
import time

device = "cuda:0"

tensor_ = torch.randn(4*256*1024*1024, device=device)
print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

 

可以成功運行：

 

 

 

 

說明當前顯卡可以成功分配4G顯存，5G顯存則不夠分配。

 

 

 

執行大量的小顯存分配：

import torch
import time

device = "cuda:0"
data = []


for _ in range(1024*1024*4):
    data.append(torch.randn(128+1, device=device))
    

# tensor_ = torch.randn(4*256*1024*1024, device=device)
print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

 

可以看到實際給tensor分配的顯存空間為4GB，

那么我們的tensor的實際大小為：4*1024*1024*129*4=2164260864=2.015625GB，

但是實際分配的顯存為4*1024*1024*512*2=4294967296=4GB，

其中的原因就是pytorch中顯存的最小分配單位是512B，pytorch分配顯存時如果存在有以前分配512B的空間沒有填滿的情況，這時如果又一次分配顯存時不能在現有所有的未填滿的512B顯存中全部裝下就會重新分配新的512B顯存，這是分配顯存的對齊方式。

上面代碼中128個float32大小為512B，剛好填滿一個最小顯存分配單元，這時再分配一個float32大小為4B，則需要重新分配一個512B大小的顯存，然后再下一次分配512B大小顯存，但是上次分配顯存單元中雖然還剩512-4=508B大小，但是不足以裝下需要分配的512B大小顯存，於是就需要重新分配一個512B大小的顯存，那么上一塊512大小顯存空間中508B大小的空間就空下來了。

 

 

 

采用上面的分配方式如果再申請一個1G連續空間的顯存，那么情況如何呢？

代碼：

import torch
import time

device = "cuda:0"
data = []


for _ in range(1024*1024*4):
    data.append(torch.randn(128+1, device=device))
    

tensor_ = torch.randn(1*256*1024*1024, device=device)
print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

顯示顯存不夠無法分配。

 

如果把1GB連續大顯存空間變成256*1024*1024個4B小顯存分配呢？

代碼：

import torch
import time

device = "cuda:0"
data = []


for _ in range(1024*1024*4):
    data.append(torch.randn(128+1, device=device))
    

# tensor_ = torch.randn(1*256*1024*1024, device=device)
for _ in range(1*256*1024*1024):
    data.append(torch.randn(1, device=device))

print(torch.cuda.memory_summary())
time.sleep(60000)

 

依然報錯：

 

 

 

 

 

由此我們可以大膽猜測pytorch在分配顯存時最小分配單元512B大小空間中只有在連續分配時才會對未填滿空間進行填充，為此補充一次測試：

import torch
import time

device = "cuda:0"
data = []


for _ in range(1024*1024*4*129):
    data.append(torch.randn(1, device=device))
    

# tensor_ = torch.randn(1*256*1024*1024, device=device)
# for _ in range(1*256*1024*1024):
#     data.append(torch.randn(1, device=device))

print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

結果顯示我們上面的猜測是不對的，也就是說最小分配單元512B大小的空間如果一次沒有填滿那么以后則不再對這部分空間進行填充。

 

 

再次進行補充測試;

import torch
import time

device = "cuda:0"
data = []


for _ in range(4*1024*1024*2):
    data.append(torch.randn(1, device=device))
    

# tensor_ = torch.randn(1*256*1024*1024, device=device)
# for _ in range(1*256*1024*1024):
#     data.append(torch.randn(1, device=device))

print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

成功運行，由此我們基本可以得出結論，512B最小分配單元必須一次性進行填充，即使是沒有填充滿在下一次申請空間時也不可以使用為填滿的512B最小分配單元中的空余空間，而上一層的分配空間2MB里面有2*1024=2048個最小分配單元，2MB空間內的這2048個空間是可以進行連續填充的，比如第一次分配顯存使用掉了2048個最小分配單元中的100個，那么下次分配顯存可以從第101個512B大小的最小分配單元開始。那么這2MB空間是否會存在對齊問題呢？

 

 

如果每次分配2MB空間，共分配4GB空間，運行如下：

import torch
import time

device = "cuda:0"
data = []


for _ in range(2*1024):
    data.append(torch.randn(128*2*1024*2, device=device))
    # data.append(torch.randn(128*2*1024+1, device=device))
    


print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

 

 

 

 

如果2MB空間中存在未填滿的情況，而下一次的申請空間又大於未填滿的空間，那么會不會對未填滿的顯存空間進行填充呢？

再次測試：

代碼：

 

import torch
import time

device = "cuda:0"
data = []


for _ in range(2*1024):
    # data.append(torch.randn(128*2*1024*2, device=device))
    data.append(torch.randn(128*2*1024+1, device=device))
    


print(torch.cuda.memory_summary())
time.sleep(60000)

 

 

 

 

 可以看到只有最小分配單元512B空間是不能再次填充的，2MB空間內的4096個最小空間是不僅可以再次填充而且如果剩余空間不夠是可以一段內存跨多個2MB空間的。

 

 

 

 

總結的說就是512B空間是只進行一次分配，不允許多個變量使用這512B空間，即使一個變量是1個float32，4B大小那么也是分配給512B空間的，這時再為另一個變量事情1個float32,4B大小也是不能利用上個空閑的508B大小而是需要重新申請一個512B空間的。同時，2MB空間可以為多個變量進行分配。

2MB大小的顯存空間是pytorch向系統一次申請的最小顯存空間，512B大小顯存時pytorch為變量分配的最小顯存空間。

 

 

 

 

 

 

 

==============================================