GPU捉襟見肘還想訓練大批量模型？誰說不可以

知識 10-17

選自Medium

機器之心編譯

深度學習模型和數據集的規模增長速度已經讓 GPU 算力也開始捉襟見肘，如果你的 GPU 連一個樣本都容不下，你要如何訓練大批量模型？通過本文介紹的方法，我們可以在訓練批量甚至單個訓練樣本大於 GPU 內存時，在單個或多個 GPU 伺服器上訓練模型。

2018 年的大部分時間我都在試圖訓練神經網路時克服 GPU 極限。無論是在含有 1.5 億個參數的語言模型（如 OpenAI 的大型生成預訓練 Transformer或最近類似的BERT模型）還是饋入 3000 萬個元素輸入的元學習神經網路（如我們在一篇 ICLR 論文《Meta-Learning a Dynamical Language Model》中提到的模型），我都只能在 GPU 上處理很少的訓練樣本。

但在多數情況下，隨機梯度下降演算法需要很大批量才能得出不錯的結果。

如果你的 GPU 只能處理很少的樣本，你要如何訓練大批量模型？

有幾個工具、技巧可以幫助你解決上述問題。在本文中，我將自己用過、學過的東西整理出來供大家參考。

在這篇文章中，我將主要討論 PyTorch 框架。有部分工具尚未包括在 PyTorch（1.0 版本）中，因此我也寫了自定義代碼。

我們將著重探討以下問題：

在訓練批量甚至單個訓練樣本大於 GPU 內存，要如何在單個或多個 GPU 伺服器上訓練模型；

如何儘可能高效地利用多 GPU 機器；

在分散式設備上使用多個機器的最簡單訓練方法。

在一個或多個 GPU 上訓練大批量模型

你建的模型不錯，在這個簡潔的任務中可能成為新的 SOTA，但每次嘗試在一個批量處理更多樣本時，你都會得到一個 CUDA RuntimeError：內存不足。

這位網友指出了你的問題！

但你很確定將批量加倍可以優化結果。

你要怎麼做呢？

這個問題有一個簡單的解決方法：梯度累積。

梯度下降優化演算法的五個步驟。

與之對等的 PyTorch 代碼也可以寫成以下五行：

在 loss.backward() 運算期間，為每個參數計算梯度，並將其存儲在與每個參數相關聯的張量——parameter.grad 中。

累積梯度意味著，在調用 optimizer.step() 實施一步梯度下降之前，我們會對 parameter.grad 張量中的幾個反向運算的梯度求和。在 PyTorch 中這一點很容易實現，因為梯度張量在不調用 model.zero_grad() 或 optimizer.zero_grad() 的情況下不會重置。如果損失在訓練樣本上要取平均，我們還需要除以累積步驟的數量。

以下是使用梯度累積訓練模型的要點。在這個例子中，我們可以用一個大於 GPU 最大容量的 accumulation_steps 批量進行訓練：

擴展到極致

你可以在 GPU 上訓練連一個樣本都無法載入的模型嗎？

如果你的架構沒有太多跳過連接，這就是可能的！解決方案是使用梯度檢查點（gradient-checkpointing）來節省計算資源。

基本思路是沿著模型將梯度在小組件中進行反向傳播，以額外的前饋傳遞為代價，節約存儲完整的反向傳播圖的內存。這個方法比較慢，因為我們需要添加額外的計算來減少內存要求，但在某些設置中挺有意思，比如在非常長的序列上訓練 RNN 模型（示例參見 https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a）。

這裡不再贅述，讀者可以查看以下鏈接：

TensorFlow：https://github.com/openai/gradient-checkpointing

PyTorch 文檔：https://pytorch.org/docs/stable/checkpoint.html

「節約內存」（Memory-poor）策略需要 O(1) 的內存（但是要求 O(n2) 的計算步）。

充分利用多 GPU 機器

現在我們具體來看如何在多 GPU 上訓練模型。

在多 GPU 伺服器上訓練 PyTorch 模型的首選策略是使用 torch.nn.DataParallel。該容器可以在多個指定設備上分割輸入，按照批維度（batch dimension）分割，從而實現模塊應用的並行化。

DataParallel 非常容易使用，我們只需添加一行來封裝模型：

但是，DataParallel 有一個問題：GPU 使用不均衡。

在一些設置下，GPU-1 會比其他 GPU 使用率高得多。

這個問題從何而來呢？下圖很好地解釋了 DataParallel 的行為：

使用 torch.nn.DataParallel 的前向和後向傳播。

在前向傳播的第四步（右上），所有並行計算的結果都聚集在 GPU-1 上。這對很多分類問題來說是件好事，但如果你在大批量上訓練語言模型時，這就會成為問題。

我們可以快速計算語言模型輸出的大小：

語言模型輸出中的元素數量。

假設我們的數據集有 4 萬辭彙，每一條序列有 250 個 token、每個 batch 中有 32 條序列，那麼序列中的每一個元素需要 4 個位元組的內存空間，模型的輸出大概為 1.2GB。要儲存相關的梯度張量，我們就需要把這個內存翻倍，因此我們的模型輸出需要 2.4GB 的內存。

這是典型 10GB GPU 內存的主要部分，意味著相對於其它 GPU，GPU - 1 會被過度使用，從而限制了並行化的效果。

如果不調整模型和／或優化方案，我們就無法輕易減少輸出中的元素數量。但我們可以確保內存負載在 GPU 中更均勻地分布。

多 GPU 機器上的均衡負載

解決辦法是把每部分輸出保留在其 GPU 上，而不是將它們聚集到 GPU-1 上。我們也需要分配損失標準計算，計算損失並進行反向傳播。

幸而，張航開源了一個名為 PyTorch-Encoding 的 PyTorch 包，它包含了這些定製的並行化功能。

我提取並稍稍改動了這個模塊，你可以從以下地址下載 gist（parallel.py）來納入並調用你的代碼。它主要包括兩個模塊：DataParallelModel 和 DataParallelCriterion，它們的用途如下：

下載地址：https://gist.github.com/thomwolf/7e2407fbd5945f07821adae3d9fd1312

DataParallelModel 和 torch.nn.DataParallel 的區別在於，前向傳播的輸出（predictions）沒有聚集在 GPU-1 上，而是作為 n_gpu 張量的元組，每個張量分布在相應的 GPU 上。

DataParallelCriterion 容器封裝了損失函數，並把 n_gpu 張量元組和目標標籤張量作為輸入。它在每個 GPU 上並行計算損失函數，像 DataParallel 分割模型輸入一樣分割目標標籤張量。

下圖說明了 DataParallelModel/DataParallelCriterion 的內部情況：

使用 DataParallelModel 和 DataParallelCriterion。

以下是你可能會遇到的兩個特定案例的解決辦法：

你的模型輸出幾個張量：你可能想分解它們：output_1, output_2 = zip(*predictions)

有時候你並不想使用並行損失函數：收集 CPU 上的所有張量：gathered_predictions = parallel.gather(predictions)

分散式訓練：在多台機器上訓練

在更大的批量上訓練時，我們要如何控制多個伺服器的算力呢？

最簡單的選擇是使用 PyTorch 的 DistributedDataParallel，它幾乎可以說是以上討論的 DataParallel 的直接替代元件。

但要注意：儘管代碼看起來很相似，但在分散式設定中訓練模型要改變工作流程，因為你必須在每個節點上啟動一個獨立的 Python 訓練腳本。正如我們將看到的，一旦啟動，這些訓練腳本可以通過使用 PyTorch 分散式後端一起同步化。

在實踐中，這意味著每個訓練腳本將擁有：

它自己的優化器，並在每次迭代中執行一個完整的優化步驟，不需要進行參數傳播（DataParallel 中的步驟 2）；

一個獨立的 Python 解釋器：這也將避免 GIL-freeze，這是在單個 Python 解釋器上驅動多個並行執行線程時會出現的問題。

當多個並行前向調用由單個解釋器驅動時，在前向傳播中大量使用 Python 循環/調用的模型可能會被 Python 解釋器的 GIL 放慢速度。通過這種設置，DistributedDataParallel 甚至在單台機器設置中也能很方便地替代 DataParallel。

現在我們直接討論代碼和用途。

DistributedDataParallel 是建立在 torch.distributed 包之上的，這個包可以為同步分散式運算提供低級原語，並能以不同的性能使用多種後端（tcp、gloo、mpi、nccl）。在這篇文章中，我將選擇一種簡單的開箱即用的方式來使用它，但你應該閱讀文檔和 Séb Arnold 寫的教程來深入理解這個模塊。

文檔：https://pytorch.org/docs/stable/distributed.html

教程：https://pytorch.org/tutorials/intermediate/dist_tuto.html

我們將考慮使用具有兩個 4 - GPU 伺服器（節點）的簡單但通用的設置：