一文讀懂 PyTorch 顯存管理機制

本文使用的 PyTorch 源碼為 master 分支，commit id 為 a5b848aec10b15b1f903804308eed4140c5263cb（https://github.com/pytorch/pytorch/tree/a5b848aec10b15b1f903804308eed4140c5263cb）。

背景介紹

剖析 PyTorch 顯存管理機制主要是為了減少顯存碎片化帶來的影響。一個簡單示例為：

如上圖所示，假設(shè)當(dāng)前想分配 800MB 顯存，雖然空閑的總顯存有 1000MB，但是上方圖的空閑顯存由地址不連續(xù)的兩個 500MB 的塊組成，不夠分配這 800MB 顯存；而下方的圖中，如果兩個 500MB 的空閑塊地址連續(xù)，就可以通過顯存碎片的整理組成一個 1000MB 的整塊，足夠分配 800MB。上方圖的這種情況就被稱為顯存碎片化。

解決方法： 當(dāng)有多個 Tensor 可以被釋放時，可以優(yōu)先釋放那些在內(nèi)存空間上前后有空閑內(nèi)存的 Tensor。這樣便于 PyTorch 整理這些空閑塊組成的碎片，讓三個塊組成一整個更大的塊。

核心問題/前提： 能否以非常小的代價（ O(1) 或 O(logn) 復(fù)雜度）拿到當(dāng)前想要釋放的塊，以及它前后的空閑空間大小？有了這個，我們可以對 Tensor 排序，優(yōu)先選擇：“前后空閑塊+自己大小” 最大的 Tensor 來釋放。

這一前提可以轉(zhuǎn)化為：找到下面兩個函數(shù)（pseudo code）：

block = get_block(tensor) # 找到存儲該 Tensor 的 Block
size = get_adjacent_free_size(block) # 返回該 Block 前后的空余空間，便于排序

研究 PyTorch 顯存管理機制后可能能回答的問題：

1. 為什么報錯信息里提示顯存夠，但還是遇到了 OOM？
2. 顯存的多級分配機制是怎樣的？為什么要這樣設(shè)計？

源碼解讀

主要看 c10/cuda/CUDACachingAllocator.cpp 里面的 DeviceCachingAllocator?類（L403）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L403

主要的數(shù)據(jù)結(jié)構(gòu)

Block：

• 分配 / 管理內(nèi)存塊的基本單位，(stream_id, size, ptr) 三元組可以特異性定位一個 Block，即 Block 維護一個 ptr 指向大小為 size 的內(nèi)存塊，隸屬于 stream_id 的 CUDA Stream。
• 所有地址連續(xù)的 Block（不論是否為空閑，只要是由 Allocator::malloc 得來的）都被組織在一個雙向鏈表里，便于在釋放某一個 Block 時快速檢查前后是否存在相鄰碎片，若存在可以直接將這三個 Block 合成為一個。

struct?Block?{?
????int?device;?//?gpu

????//?Block?和被分配時調(diào)用者的?stream?是綁定的，即使釋放了也無法給其他?stream?使用
????cudaStream_t?stream;?//?allocation?stream

????stream_set?stream_uses;?//?streams?on?which?the?block?was?used
????size_t?size;?//?block?size?in?bytes
????BlockPool*?pool;?//?owning?memory?pool
????void*?ptr;?//?memory?address
????bool?allocated;?//?in-use?flag
????Block*?prev;?//?prev?block?if?split?from?a?larger?allocation
????Block*?next;?//?next?block?if?split?from?a?larger?allocation
????int?event_count;?//?number?of?outstanding?CUDA?events
};

BlockPool：

• 內(nèi)存池，用 std::set 存儲 Block 的指針，按照 (cuda_stream_id -> block size -> addr) 的優(yōu)先級從小到大排序。所有保存在 BlockPool 中的 Block 都是空閑的。

struct?BlockPool?{
????BlockPool(
????????Comparison?comparator,
????????bool?small,
????????PrivatePool*?private_pool?=?nullptr)
????????:?blocks(comparator),?is_small(small),?owner_PrivatePool(private_pool)?{}
????std::set?blocks;
????const?bool?is_small;
????PrivatePool*?owner_PrivatePool;?//?供?CUDA?Graphs?使用，此處不詳講
};

• DeviceCachingAllocator 中維護兩種 BlockPool (large_blocks, small_blocks)，分別存放較小的塊和較大的塊（為了分別加速小需求和大需求），簡單地將 <= 1MB 的 Block 歸類為小塊，> 1MB 的為大塊。

直觀理解 Block、BlockPool 見下圖：

BlockPool 示意圖，左邊兩個 500MB 的塊 size 相同，因此上面的比下面的地址更小

總結(jié)： Block 在 Allocator 內(nèi)有兩種組織方式，一種是顯式地組織在 BlockPool（紅黑樹）中，按照大小排列；另一種是具有連續(xù)地址的 Block 隱式地組織在一個雙向鏈表里（通過結(jié)構(gòu)體內(nèi)的 prev, next 指針），可以以 O(1) 時間查找前后 Block 是否空閑，便于在釋放當(dāng)前 Block 時合并碎片。

malloc 函數(shù)解析：返回一個可用的 Block（L466）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L466

一個區(qū)分：（出于簡潔考慮，下面全文都對這兩個表述進行區(qū)分，不多贅述）

• size：請求分配的顯存大小
• alloc_size：需要 cudaMalloc 時的顯存大小

一個 hard-coded： 根據(jù) size 決定實際上的 alloc_size（get_allocation_size 函數(shù)，L1078）：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1078

• 為小于 1MB 的 size 分配 2MB；
• 為 1MB ～ 10MB 的 size 分配 20MB；
• 為 >= 10MB 的 size 分配 { size 向上取整至 2MB 倍數(shù) } MB。

尋找是否有合適的 block 會有五個步驟，如果這五個步驟都沒找到合適的 Block，就會報經(jīng)典的 [CUDA out of memory. Tried to allocate ...] 錯誤（具體見下“分配失敗的情況”）。

步驟一：get_free_block 函數(shù)（L1088）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1088

TLDR：嘗試在 Allocator 自己維護的池子中找一個大小適中的空閑 Block 返回。
* TLDR = Too Long; Didn't Read

• 用當(dāng)前的 (size, stream_id) 這二元組制作 Block Key 在對應(yīng)的 BlockPool 中查找；
• 環(huán)境變量 PYTORCH_CUDA_ALLOC_CONF 中指定了一個閾值 max_split_size_mb，有兩種情況不會在此步驟分配：

1. 需要的 size 小于閾值但查找到的 Block 的比閾值大（避免浪費block）；
2. 兩方都大于閾值但 block size 比需要的 size 大得超過了 buffer（此處是 20MB，這樣最大的碎片不超過 buffer 大小）。

• 這里的這個閾值 max_split_size_mb 涉及一個有趣的特性，后面會講到。
• 若成功分配一個 Block，則將這個 Block 從 BlockPool 中刪除。后續(xù)用完釋放（free）時會再 insert 進去，見free_block : L976：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L976）

步驟二：trigger_free_memory_callbacks 函數(shù)（L1106）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1106

TLDR：手動進行一波垃圾回收，回收掉沒人用的 Block，再執(zhí)行步驟一。

• 若第一步的 get_free_block 失敗，則在第二步中先調(diào)用 trigger_free_memory_callbacks，再調(diào)用一次第一步的 get_free_block；
• trigger_free_memory_callbacks 本質(zhì)上通過注冊調(diào)用了 CudaIPCCollect 函數(shù)，進而調(diào)用 CudaIPCSentDataLimbo::collect 函數(shù)（torch/csrc/CudaIPCTypes.cpp : L64）；https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/torch/csrc/CudaIPCTypes.cpp#L64
• CudaIPCSentDataLimbo 類管理所有 reference 不為 0 的 Block，所以實際上 collect 函數(shù)的調(diào)用算是一種懶更新，直到無 Block 可分配的時候才調(diào)用來清理那些 reference 已經(jīng)為 0 的 Block（值得一提的是，該類的析構(gòu)函數(shù)會首先調(diào)用 collect 函數(shù)，見 torch/csrc/CudaIPCTypes.cpp : L58）https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/torch/csrc/CudaIPCTypes.cpp#L58；
• 相關(guān)源碼可以看 torch/csrc/CudaIPCTypes.h & .cpp。

步驟三：alloc_block 函數(shù)（L1115）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1115

TLDR：Allocator 在已有的 Block 中找不出可分配的了，就調(diào)用 cudaMalloc 創(chuàng)建新的 Block。

• 步驟一、二中重用 block 失敗，于是用 cudaMalloc 分配內(nèi)存，大小為 alloc_size；
• 注意有一個參數(shù) set_fraction 會限制可分配的顯存為當(dāng)前剩余的顯存 * fraction（若需要分配的超過這一限制則失敗），但還沒搞清楚哪里會指定這個（TODO）；
• 新分配的內(nèi)存指針會被用于創(chuàng)建一個新 Block，新 Block 的 device 與 cuda_stream_id 與 caller 保持一致。

上面幾個步驟都是試圖找到一些空閑顯存，下面是兩個步驟是嘗試進行碎片整理，湊出一個大塊顯存

步驟四：release_available_cached_blocks 函數(shù)（L1175）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1175

TLDR：先在自己的池子里釋放一些比較大的 Block，再用 cudaMalloc 分配看看

• 如果上面的 alloc_block 失敗了，就會嘗試先調(diào)用這一函數(shù)，找到比 size 小的 Block 中最大的，由大至小依次釋放 Block，直到釋放的 Block 大小總和 >= size（需要注意，這一步驟只會釋放那些大小大于閾值 max_split_size_mb 的 Block，可以理解為先釋放一些比較大的）；
• 釋放 block 的函數(shù)見 release_block（L1241），主要就是 cudaFree 掉指針，再處理一些 CUDA graphs 的依賴，更新其他數(shù)據(jù)結(jié)構(gòu)等等，最后把這個 Block 從 BlockPool 中移除；
• 當(dāng)前整個 BlockPool（作為提醒，Allocator 有兩個 BlockPool，這里指的是最初根據(jù) size 指定的大 pool 或小 pool）中，可以釋放的 Block 需要滿足兩個條件：cuda_stream_id 相同的，且大小要大于閾值 max_split_size_mb。如果將這樣的 Block 全部釋放的空間仍比 size 小，那么這一步就會失敗。
• 釋放掉了一批 Block 之后，再次執(zhí)行步驟三中的 alloc_block 函數(shù)，創(chuàng)建新 Block。

步驟五：release_cached_blocks 函數(shù)（L1214）

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1214

• 如果釋放一些 Block 還不夠分配，則把整個 Allocator 中的 large / small pool 全部釋放掉（同樣調(diào)用 release_block：L1241：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1241），再次調(diào)用 alloc_block 函數(shù)。

malloc 分配失敗的情況

https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L519-L560

會報經(jīng)典的 CUDA out of memory. Tried to allocate ... 錯誤，例如：

CUDA out of memory. Tried to allocate 1.24 GiB (GPU 0; 15.78 GiB total capacity; 10.34 GiB already allocated; 435.50 MiB free; 14.21 GiB reserved in total by PyTorch)

• Tried to allocate：指本次 malloc 時預(yù)計分配的 alloc_size；
• total capacity：由 cudaMemGetInfo 返回的 device 顯存總量；
• already allocated：由統(tǒng)計數(shù)據(jù)記錄，當(dāng)前為止請求分配的 size 的總和；
• free：由 cudaMemGetInfo 返回的 device 顯存剩余量；
• reserved：BlockPool 中所有 Block 的大小，與已經(jīng)分配的 Block 大小的總和。
  即 [reserved] = [already allocated] + [sum size of 2 BlockPools]

注意，reserved + free 并不等同于 total capacity，因為 reserved 只記錄了通過 PyTorch 分配的顯存，如果用戶手動調(diào)用 cudaMalloc 或通過其他手段分配到了顯存，是沒法在這個報錯信息中追蹤到的（又因為一般 PyTorch 分配的顯存占大部分，分配失敗的報錯信息一般也是由 PyTorch 反饋的）。

在這個例子里，device 只剩 435.5MB，不夠 1.24GB，而 PyTorch 自己保留了 14.21GB（儲存在 Block 里），其中分配了 10.3GB，剩 3.9GB。那為何不能從這 3.9GB 剩余當(dāng)中分配 1.2GB 呢？原因肯定是碎片化了，而且是做了整理也不行的情況。

實現(xiàn)自動碎片整理的關(guān)鍵特性：split

前面提到，通過環(huán)境變量會指定一個閾值 max_split_size_mb，實際上從變量名可以看出，指定的是最大的可以被 “split” 的 Block 的大小。

在本節(jié)最開頭解釋過，步驟三 alloc_block 函數(shù)通過 cuda 新申請的 Block（這是 Allocator 唯一一個創(chuàng)建新 Block 的途徑）大小是計算出的 alloc_size 而不是用戶通過 malloc 請求的大小 size。因此，如果 malloc 在上述五個步驟中成功返回了一個 Block，Allocator 會通過函數(shù) should_split (L1068) 檢查：

• 由于過量分配內(nèi)存，Block 內(nèi)部產(chǎn)生的大小為 alloc_size - size 的碎片大小稱為 remaining；
• 如果這個 Block 屬于 small_blocks 且 remaining >= 512 Bytes；或者 remaining >= 1MB 且該 Block 大小沒有超過上述的閾值，則這個 Block 需要被 split。

被 split 的操作很簡單，當(dāng)前的 Block 會被拆分成兩個 Block，第一個大小正好為請求分配的 size，第二個則大小為 remaining，被掛到當(dāng)前 Block 的 next 指針上（這一過程見源碼 L570～L584：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L570-L584）。這樣一來，這兩個 Block 的地址自然而然成為連續(xù)的了。隨著程序運行，較大的 Block（只要仍小于閾值 max_split_size_mb）會不斷被分成小的 Block。值得注意的是，由于新 Block 的產(chǎn)生途徑只有一條，即通過步驟三中的 alloc_block 函數(shù)經(jīng)由 cudaMalloc 申請，無法保證新 Block 與其他 Block 地址連續(xù)，因此所有被維護在雙向鏈表內(nèi)的有連續(xù)地址空間的 Block 都是由一個最初申請來的 Block 拆分而來的。

一段連續(xù)空間內(nèi)部（由雙向鏈表組織的 Block 們）如下圖所示：

當(dāng) Block 被釋放時，會檢查其 prev、next 指針是否為空，及若非空是否正在被使用。若沒有在被使用，則會使用 try_merge_blocks (L1000) 合并相鄰的 Block。由于每次釋放 Block 都會檢查，因此不會出現(xiàn)兩個相鄰的空閑塊，于是只須檢查相鄰的塊是否空閑即可。這一檢查過程見 free_block 函數(shù)（L952）。又因為只有在釋放某個 Block 時才有可能使多個空閑塊地址連續(xù)，所以只需要在釋放 Block 時整理碎片即可。

try_merge_blocks (L1000)：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L1000

free_block 函數(shù)：https://github.com/pytorch/pytorch/blob/a5b848aec10b15b1f903804308eed4140c5263cb/c10/cuda/CUDACachingAllocator.cpp#L952

關(guān)于閾值 max_split_size_mb ，直覺來說應(yīng)該是大于某個閾值的 Block 比較大，適合拆分成稍小的幾個 Block，但這里卻設(shè)置為小于這一閾值的 Block 才進行拆分。個人理解是，PyTorch 認為，從統(tǒng)計上來說大部分內(nèi)存申請都是小于某個閾值的，這些大小的 Block 按照常規(guī)處理，進行拆分與碎片管理；但對大于閾值的 Block 而言，PyTorch 認為這些大的 Block 申請時開銷大（時間，失敗風(fēng)險），可以留待分配給下次較大的請求，于是不適合拆分。默認情況下閾值變量 max_split_size_mb 為 INT_MAX，即全部 Block 都可以拆分。

總結(jié)

剖析 PyTorch 顯存管理機制可以幫助我們更好地尋找顯存優(yōu)化的思路。目前我們正在聯(lián)合 Dynamic Tensor Rematerialization (DTR)（https://arxiv.org/abs/2006.09616） 的作者@圓角騎士魔理沙 推進 PyTorch 上 DTR 技術(shù)的工業(yè)界落地，其中一項就是解決上述文章中提到的訓(xùn)練過程內(nèi)存碎片化問題。近期會有一篇文章介紹 DTR 落地遇到的問題 & 解決方案 & benchmark，敬請期待。