ML系統與架構小組：如何在單塊GPU上訓練超大型深度學習模型

機器之心原創

人工智慧研學社

問題：GPU 內存限制

GPU 在深度神經網路訓練之中的強大表現無需我贅言。通過現在流行的深度學習框架將計算分配給 GPU 來執行，要比自己從頭開始便捷很多。然而，有一件事你會避之唯恐不及，即 GPU 的動態隨機存取內存（DRAM（Dynamic Random Access Memory））限制。

在給定模型和批量大小的情況下，事實上你可以計算出訓練所需的 GPU 內存而無需實際運行它。例如，使用 128 的批量訓練 AlexNet 需要 1.1GB 的全局內存，而這僅是 5 個卷積層加上 2 個全連接層。如果我們觀察一個更大的模型，比如 VGG-16，使用 128 的批量將需要大約 14GB 全局內存。目前最先進的 NVIDIA Titan X 內存為 12GB。VGG-16 只有 16 個卷積層和 3 個全連接層，這比大約有 100 個層的 resnet 模型小很多。

圖 1: 當使用基線、全網分配策略時 GPU 內存的使用情況（左軸）。(Minsoo Rhu et al. 2016)

現在，如果你想要訓練一個大於 VGG-16 的模型，而你又沒有一個具有內存的 GPU，你也許有幾個解決內存限制問題的選擇。

• 減小你的批量大小，但這可能會妨礙你的訓練速度和精確度。

• 在多 GPU 環境下做模型並行，這是另一個複雜的事情。

• 縮小你的模型，如果你不情願做出上述兩個選擇，或者已經嘗試但效果不好。

或者你也可以等待下一代更強大的 GPU 出現。然而將來的網路變得更大更深是大勢所趨，我們並不希望物理內存的限制成為演算法開發的一個障礙。

觀察：什麼佔用內存？

我們可以根據功能性把 GPU 內存中的數據分為 4 個部分：

• 模型參數

• 特徵圖

• 梯度圖

• 工作區

前 3 個功能容易理解。模型參數的意義對於所有人來說都很熟悉了。特徵圖是正向過程中生成的中間結果。梯度圖是反向過程中生成的中間結果。而工作區是 cuDNN 庫函數所使用的臨時變數／矩陣的一個緩衝區。對於一些 cuDNN 函數，用戶需要將緩衝區作為函數參數傳給 GPU 內核函數。一旦函數返回，該緩衝區將被釋放。

我們可以看到，一般而言，我們有的層越多，分配給特徵圖的內存比重就越多（由上圖中的三角形表示）。我們也可以看到對於像 VGG-16 這樣的更大模型來說，這一比例基本上要超過 50%。

想法：使用 CPU 內存作為臨時容器

有一個關於特徵圖的事實：它們在正向過程中生成之後立刻被用於下一層，並在稍後的反向過程中僅再用一次。每個 GPU 內核函數僅使用與當前層（通常只有 1 個張量）相關的特徵映射。這將導致絕大多數內存在幾乎所有的時間上出現空置的情況（它們保有數據但不使用）。

這一想法是：如果 GPU 內存中的大部分數據出現空置，為什麼不把它們保存在更便宜的 CPU 內存上呢？下圖更清晰地展現了這一想法。

左側部分所示的間隙表明特徵圖如何在內存之中被空置。右側部分表明這一節省 GPU 內存使用的想法：使用 CPU 內存作為那些特徵圖的臨時容器。

權衡：時間 vs 空間

根據論文，vDNN（虛擬化 DNN，也就是原文的系統設計的名稱）把 AlexNet 的平均 GPU 內存使用成功降低了 91%，GoogLeNet 的降低了 95%。但你很可能已經看到，這樣做的代價是訓練會變慢。例如，vDNN 可以在 12GB 的 GPU 上使用 256 的批量訓練 VGG-16，但是假設我們在一塊擁有足夠內存的 GPU 上訓練同樣的模型而不使用 vDNN 來優化內存使用，我們可以避免 18% 的性能損失。

當使用 cuDNN 核時，工作區大小也會出現權衡的情況。一般而言，你有的工作區越多，你可以使用的演算法就越快。如果有興趣請查閱 cuDNN 庫的參考。在後面的整個討論中我們都將會看到有關時間空間的這一權衡。

優化策略：在前向過程中卸載，在後向過程中預取

你應該已經知道 vDNN 是如何在正向過程中優化內存分配的。基本的策略是在生成特徵圖後將其從 GPU 上卸下傳給 CPU，當它將在反向過程中被重新使用時再從 CPU 預取回 GPU 內存。這個存儲空間可被釋放以作他用。這樣做的一個風險是如果網路拓撲是非線性的，特徵圖的一個張量可能被應用於數個層，從而導致它們不能被立刻卸載。當然，這個問題可以通過簡單的完善優化策略來解決。

在後向過程中，vDNN 採用一種更具侵略性的策略，這是由於相較於特徵圖，梯度圖不存在「稍後重用」的問題。因此一旦生成了相關的權值更新，它們就可以被釋放（相較於那些映射，這些權值更新是很小的）。

優化策略：內存管理器 CUDA 流

vDNN 實現的關鍵是使用一個 cuda 流來管理內存分配／釋放、卸載和預取。以下是一些細節：

傳統的 cuda 內存分配／釋放（cudaMalloc & cudaFree）是同步性 API。由於這兩個操作隨著訓練過程需要不斷地被執行，同步性 API 並不夠高效。

如同分配/釋放操作，卸載 API 也需要是非同步性的。當 vDNN 選擇卸載特徵圖的一個張量時，vDNN 的內存管理器流（memory manager stream of vDNN）將在主機上分配一個頁鎖定的內存區域，並且通過 PCIe 發出一個非阻塞傳輸（a non-blocking transfer）將其傳輸至 CPU。這些特徵圖張量在正向過程中為只讀的，因此這個傳輸過程可以在計算中被安全地覆蓋。當依賴於其傳輸和計算都完成後，相關的張量可以被釋放。只有當前層的卸載完成時，才能進行下一層的處理。

預取操作是在反向處理中從 CPU 返回到 GPU 以得到之前被卸載的特徵圖。和上面的操作類似，預取操作也需要是非同步性的。由於在預取和相同層的計算之間存在數據依賴，vDNN 將同時非同步開始當前層的計算以及前一層的預取。

成本：為了節省內存的性能損失在哪裡？

最顯著的潛在性能損失來自於由卸載／預取引入的隱式依賴（implicit dependency）。我們考慮這樣一種情況：當數據傳輸比正向計算需要花費更長的時間時，卸載／預取操作將會帶來性能損失。下圖表清晰地表明這種情況（圖 9：卸載和預取的性能影響 (Minsoo Rhu et al. 2016)）

相似情形也可能在反向過程中發生。

問題的形式化：在限制內存預算的情況下，如何獲得最佳性能？

如上所述，在時間和空間之間有一個權衡，並且在前的章節中我們已經看到這個權衡是產生作用的。想像一下你正在 12GB 的 GPU 上使用 128 的批量（這需要 14GB 內存如果沒有使用卸載／預取）訓練 VGG-16。假設我們可以使用上面提到的想法，把內存使用壓縮到只有 2GB，那麼僅使用了大約 2GB 的內存也許很浪費，因為你本可以利用空閑的 10GB 空間來減少性能損失（這裡忽略能耗方面的成本）。因此，我們可以這種方式重新形式化這個問題：在限制內存預算的情況下，如何獲得最佳性能？

配置時間-空間權衡：決定一個層是否應該被卸載/預取（offloaded/prefetched），以及應該選擇哪種卷積演算法。

從這個問題的形式化中我們可以看出，我們往往並不需要把內存使用壓縮到極致，有些層完全可以不被卸載/預取。為了取得最佳配置，我們需要為每個層決定兩件事：一個是我們是否需要卸載/預取，一個是我們在前向/反向傳播過程中該選擇哪個演算法（更快收斂的演算法需要更大的存儲空間）。

通常，靠近輸入的層有較長重用度（reuse）距離。因此，首先卸載/預取這些層比較好。於是，我們並不需要為每一層決定是否使用（這個選擇總量隨層級數成指數增長）。我們只需要選擇一層作為基準，與之相比更加靠近輸入層的都進行卸載/預取，其餘層將其張量保留在 GPU 上。

為每層決定演算法也並不可行（同樣選擇總量隨層級數成指數增長）。我們可以通過例如強制多層使用相同的演算法（gemm 或 fft 等演算法）來縮小我們的可選備用空間。這將原本指數級大小的選擇數降為線性大小。

經過上面的處理，現在我們的配置選擇空間已經足夠小，這樣我們可以盡情搜索以確定最佳配置方案。下面的圖表說明了我們的配置空間：

左上角表示最節省內存的配置（卸載/預取每一層，並使用最慢的演算法），右下角表示性能最優的配置（當然，真實的配置空間應該是網格，可行模型與不可行模型之間的邊界應改成階梯狀，不過這一圖表足以表明整個搜索空間的大致分布特徵）。

下一步即找到具有最佳性能的可行配置。我們可以沿可行性邊界搜索配置。如果你對該搜索過程的實現細節感興趣，請參閱原論文（搜索過程在第 3.C 節中進行了描述）。

論文連接：https://arxiv.org/pdf/1602.08124.pdf

推薦者介紹：

Hongyu 畢業於上海交通大學 09 級 ACM 班，多倫多大學在讀博士，現從事關於深度神經網路訓練過程的一些 profiling 的工作。

加入機器之心 ML 系統與架構小組：

近些年人工智慧的突破不僅僅是機器學習演算法的努力，更有包含了其所依賴的系統和架構領域的進步。如果你想更加全面的了解機器學習這個領域，對其依賴的上下游領域的了解會很有幫助。系統和架構就是這樣一個下游領域。這個方向的工作極大方便了其上游演算法的開發，或許你也是現在諸多 ML 系統產品的用戶之一。但面對一個這樣重要的跨領域的方向，你可能會感到這樣一些困境：

1. 找不到合適的學習資料

2. 有學習動力，但無法堅持

3. 學習效果無法評估

4. 遇到問題缺乏討論和解答的途徑

不論你是想要獲得相關跨領域的更全面大局觀，還是你只是想對手中的 ML 系統工具更加了解，你都是機器之心想要幫助和挖掘的對象。基於這個方向現在越來越成為眾多研究人員關注的焦點，機器之心發起了一個互助式學習小組——「人工智慧研學社· ML 系統與架構小組」。本小組將通過優質資料分享、教材研習、論文閱讀、群組討論、專家答疑、講座與分享等形式加強參與者對強化學習和深度學習的理解和認知。

• 面向人群：有一定的機器學習演算法基礎，並且對分布式計算、並行計算等也有所了解，同時想掌握此方向最新成果的人

• 學習形式：學習資料推薦、統一進度學習（教材或論文）、群組討論、專家答疑、講座等。

• 加入方式：

1）添加機器之心小助手微信，並註明：加入系統與架構學習組

2）完成小助手發送的入群測試（題目會根據每期內容變化），並提交答案，以及其他相關資料（教育背景 、從事行業和職務 、人工智慧學習經歷等）

3）小助手將邀請成功通過測試的朋友進入「人工智慧研學社· ML 系統與架構學習組」

入群測試 QUIZ

1）教育背景 2）從事行業和職務 3）人工智慧經歷 4）人工智慧系統或架構經歷

1. 數據並行和模型並行的區別在於哪裡？

2. 在訓練過程中，數據並行和模型並行很難達到線性的 scale-up，你認為瓶頸會在哪？有沒有辦法突破這樣的瓶頸？這些突破的方法又會造成什麼樣的副作用？