一文讀懂 PyTorch 顯存管理機制

最近在做 RLHF 場景下的訓練和推理複用機器的優化，需要及時釋放掉不用的 GPU 資源。翻出了之前和同事一起寫的這篇文章，雖然過去 2 年了，但依然好看，因此轉載到這裏。

本文使用的 PyTorch 源碼爲 master 分支，commit id 爲 a5b848aec10b15b1f903804308eed4140c5263cb。

背景介紹

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

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

解決方法： 當有多個 Tensor 可以被釋放時，可以優先釋放那些在內存空間上前後有空閒內存的 Tensor。這樣便於 PyTorch 整理這些空閒塊組成的碎片，讓三個塊組成一整個更大的塊。

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

這一前提可以轉化爲：找到下面兩個函數（pseudo code）：

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

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

  1. 爲什麼報錯信息裏提示顯存夠，但還是遇到了 OOM？

  2. 顯存的多級分配機制是怎樣的？爲什麼要這樣設計？

源碼解讀

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

主要的數據結構

Block：

struct Block { 
    int device; // gpu

    // Block 和被分配時調用者的 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：

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

直觀理解 Block、BlockPool 見下圖：

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

總結：Block 在 Allocator 內有兩種組織方式，一種是顯式地組織在 BlockPool（紅黑樹）中，按照大小排列；另一種是具有連續地址的 Block 隱式地組織在一個雙向鏈表裏（通過結構體內的 prev, next 指針），可以以 O(1) 時間查找前後 Block 是否空閒，便於在釋放當前 Block 時合併碎片。

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

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

一個 hard-coded： 根據 size 決定實際上的 alloc_size（get_allocation_size 函數，L1078）：

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

步驟一：get_free_block 函數（L1088）

TLDR：嘗試在 Allocator 自己維護的池子中找一個大小適中的空閒 Block 返回。

  1. 需要的 size 小於閾值但查找到的 Block 的比閾值大（避免浪費 block）；

  2. 兩方都大於閾值但 block size 比需要的 size 大得超過了 buffer（此處是 20MB，這樣最大的碎片不超過 buffer 大小）。

步驟二：trigger_free_memory_callbacks 函數（L1106）

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

步驟三：alloc_block 函數（L1115）

TLDR：Allocator 在已有的 Block 中找不出可分配的了，就調用 cudaMalloc 創建新的 Block。

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

步驟四：release_available_cached_blocks 函數（L1175）

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

步驟五：release_cached_blocks 函數（L1214）

malloc 分配失敗的情況

會報經典的 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)

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

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

實現自動碎片整理的關鍵特性：split

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

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

被 split 的操作很簡單，當前的 Block 會被拆分成兩個 Block，第一個大小正好爲請求分配的 size，第二個則大小爲 remaining，被掛到當前 Block 的 next 指針上（這一過程見源碼 L570～L584）。這樣一來，這兩個 Block 的地址自然而然成爲連續的了。隨着程序運行，較大的 Block（只要仍小於閾值 max_split_size_mb）會不斷被分成小的 Block。值得注意的是，由於新 Block 的產生途徑只有一條，即通過步驟三中的 alloc_block 函數經由 cudaMalloc 申請，無法保證新 Block 與其他 Block 地址連續，因此所有被維護在雙向鏈表內的有連續地址空間的 Block 都是由一個最初申請來的 Block 拆分而來的。

一段連續空間內部（由雙向鏈表組織的 Block 們）如下圖所示：

連續空間內部的雙向鏈表（鏈表箭頭已省略）

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

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

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/2Mx-6Tf95Pku3QqqzjAo7A

猜你喜歡