教程｜如何在單塊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 內存中的數據分為 4 個部分：

• 模型參數（權重）

• 特徵映射

• 梯度映射

• 工作區

前 3 個功能性容易理解。每個人都知道權重是什麼。特徵映射是正向過程中生成的中間結果。梯度映射是反向過程中生成的中間結果。工作區是 cuDNN 函數臨時變數／矩陣的一個緩衝區。對於一些 cuDNN 函數，用戶需要通過緩衝區到達作為函數參數的內核。一旦返回函數，該緩衝區將被釋放。

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

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

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

這一想法是：如果 GPU 內存中的大部分數據出現空置，為什麼不把它們保存在更便宜的 CPU 內存上呢？這裡是一張在 AlexNet 的實例中表明什麼在進行的圖表。

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

權衡：時間 vs 空間

根據論文，vDNN（可視化DNN）把 AlexNet 的平均 GPU 內存使用成功降低了 91% ， GoogLeNet 的降低了 95%。你很可能已經看到，這樣做的代價是訓練會變慢。例如，vDNN 可以在 12GB 的 GPU 上使用 256 的批量訓練 VGG-16，但是相比於帶有足夠內存的假設性 GPU ，這會產生 18% 的性能損失。

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

優化策略：在正向與反向的過程中

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

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

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

vDNN 的關鍵部分是 cuda 流，它管理內存分配／釋放、卸載和預取。這裡是一些細節：

傳統的 cuda 內存分配／釋放（cudaMalloc & cudaFree）是同步性 API。由於隨著訓練過程的進行連續地發生，同步性 API 並不夠高效。

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

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

成本：如何償還為節省內存而付出的性能代價？

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

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

算的情況下，如何獲得最佳性能？

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

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

為了取得最佳配置，我們需要為每個層決定兩件事：我們是否需要卸載/預取，我們在前向/反向傳播過程中該選擇哪個演算法（更快收斂的演算法需要更大的存儲空間）。

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

為每層決定演算法也並不可行（同樣選擇隨層級數成指數增長）。我們可以通過強制每層使用相同的演算法（gemm 或fft 等演算法）來簡化這一問題。

現在我們有一個小型配置空間，這樣我們可以盡情搜索以確定最佳配置方案。 下面的圖表說明了我們的配置空間：

左上角表示存儲器優化配置（卸載/預取每一層，並使用最慢的演算法），右下角表示性能最優的配置（當然，真實的配置空間應該是網格，可行模型與不可行模型之間的邊界應改成階梯狀，不過這一圖表足以說明這一點）。

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