文 | 立交橋跳水冠軍@知乎
本文已獲作者授權，禁止二次轉載

魚與熊掌不可兼得，深度學習領域中的幾個指標也相同。

主要的指標有如下四個：

（1）精度：自然精度是一個模型最根本的衡量指標，如果一個模型精度不高，再快，再綠色環保也無濟于事。基本上所有刷榜的工作都是用其他所有指標換精度：比如用更深的網絡就是用memory和computation換精度。然而到了實際應用中，尤其是部署側，工程師越來越多的用一些方法適當的減少精度從而換取更小的內存占用或者運算時間 。

（2）內存：Out Of Memory Error恐怕是煉丹師最常見的情況了。內存（或者說可以高效訪問的存儲空間）的尺寸是有限的，如果網絡訓練需要的內存太大了，可能程序直接就報錯了，即使不報錯，也需要把內存中的數據做個取舍，一部分存到相對較慢的存儲介質中（比如host memory)。

(3) 通信：隨著網絡規模越來越大，分布式訓練已經是state-of-the-art的網絡模型必不可少的部分（你見過誰用單卡在ImageNet訓練ResNet50？），在大規模分布式系統，通信帶寬比較低，相比于computation或者memory load/sotre，network communication會慢很多，如果可以降低通信量，那么整個網絡的訓練時間就會有大幅減少：這樣研究員就不會借口調參，實際上把模型往服務器上一扔自己就跑出去浪了。(資本家狂喜)

（4）計算：雖然我們用的是計算機，但實際上恐怕只有很少的時間用于計算(computation)了，因為大多數時間都在等待數據的讀取或者網絡通信，不過即便如此，對于一些計算密集型的神經網絡結構（比如BERT，幾乎都是矩陣乘法），制約我們的往往是設備的計算能力（FLOPS），即每秒鐘可以處理多少浮點計算。

計算能力是重要指標，即使多數情況用不滿無疑同時將以上四點做好是我們追求的，但現實往往很殘酷，需要我們做很多取舍。今天就在這兒介紹一些業內常見的trade-off。

（1）計算換內存

很多時候內存是最重要的：計算慢了我們多等一等就行了，內存爆了就徹底訓練不了了。在神經網絡訓練中，內存的占用大頭往往是activation，即神經網絡每層的輸出。我們在訓練的時候需要把這些輸出（activation）記錄下來，因為我們反向傳播的時候需要用這些activation計算梯度。

一個很直觀的想法就是：我們干脆把一堆activation扔掉，到時候需要他們的時候再算一遍。這就是checkpoint機制的想法。雖然這個想法很簡單，但是對于一個特定的神經網絡，究竟扔掉/保留哪些activation一直沒有定論，有興趣的同學可以看一下我之前寫的另一篇文章了解這個專題：

立交橋跳水冠軍：DNN顯存優化的終點？Checkmate論文總結：

https://zhuanlan.zhihu.com/p/299861314

（2）通信換內存

隨著BERT，GPT的發展，研究員發現一件更尷尬的事情：內存不夠了，但這次不僅僅是裝不下activation，甚至光是參數（parame）和參數對應的optimizer都裝不下了。那之前說的checkpoint就不管用了（人家只負責省activation）。

這時候有些讀者會有想法：那如果我一張卡裝不下，就兩張卡來唄。恭喜你！你的想法和世界上最頂尖的程序員一樣！這種做法可以被稱為Model Parallel，即每個分布式節點存儲不同的參數，feed一樣的數據。目前Model Parallel有兩種，粗略來說可以分成intra-layer拆分和inter-layer拆分：

▲intra-layer拆分，多見于NLP模型▲inter-layer拆分，多見于CV模型

上面的例子可能比較抽象，我們來結合下面的兩個具體工作說一下這兩種model-parallel：

首先是intra-layer的拆分：

我們知道神經網絡是一層一層的，每層可能是一個卷積，一個Pool，或者BN什么的。如果我們對一層進行了拆分，那么就是intra-layer的。下面這張圖摘自英偉達的Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（https://arxiv.org/abs/1909.08053）,描述了如何把兩個很大的矩陣乘法拆開到兩個節點上來算

本來矩陣乘法是和? (忽略激活層什么的)。我們知道矩陣乘法有很好的性質：我們可以把矩陣乘法變成分塊矩陣乘法。因此我們可以把上面的矩陣乘法變成

因此如果我們讓第一個節點計算,第二個節點計算?,最后再累加兩個人的結果，不就好了？這樣的好處就是第一個節點只需要存儲A1和B1，第二個只需要A2和B2，相當于節省了一半的空間（和一個節點存儲A和B相比）

inter-layer拆分：和intra-layer不同，我們只做網絡層之間的切分。這種切分方式更符合直覺。上面的例子來自PipeDream(https://arxiv.org/abs/1806.03377、PipeDream: Fast and Efficient Pipeline Parallel DNN Training)。假設我們有5層的神經網絡，有四個節點可以用，那我們可以讓第一個節點算第一層，第二個節點算第二層，第三個節點算第三層，第四個節點算最后兩層。這樣我們很直接的就把網絡拆開了，不管第1-5層的具體操作是卷積，BN還是什么，都可以這么搞。

但這么拆有一個問題，就是每臺機器之間都存在數據依賴:當你發現第1,2,3節點在悠閑地打王者，第四節點在苦逼的干活，你就上去質問他們你們為啥在劃水？他們表示很無辜：我在等第四節點把結果算完，然后把梯度傳給我啊。

把每個節點工作的時間系統的記錄下來，發現所有時刻都只有一個人在干活

很自然的想法就是，如果每個人都處理不同的batch不就好了？但這樣做可能會引發精度的問題。有興趣的讀者可以去看PipeDream的論文。

（3）計算換通信

雖然剛才介紹了模型并行，但目前主流的還是數據并行，即每張卡分到同樣的parameter，每次接收不同的input，算完之后每個人把自己的local gradient做一次同步，得到global gradient來更新本地的參數（如下圖所示）

在這種情況下，我們的通信只發生在gradient allReduce的時刻（即圖中最下面灰色的框）。雖然它只是訓練過程的一部分，但因為隨著分布式系統的增大，通信速度和計算、訪存時間相比會越來越慢，因此這個allReduce操作逐漸成為了性能瓶頸。

為了打破這個瓶頸，有些研究員嘗試壓縮梯度：每次我們并不通信梯度本身，而是先把梯度做一個壓縮，讓他們的size變小，然后把壓縮后的數據做一次傳輸，最后在本地解壓縮這些數據，從而完成一次梯度的allReduce。其中做的比較好的就是TernGrad（TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning（https://arxiv.org/abs/1705.07878））。這個算法將原來一個N*float32這么大的梯度tensor壓縮成了N*3這么大的tensor，再加一些可以忽略不計的meta-data。即用(-1,0,1)來表示原本float32的數值。

（4）顯存換計算

這方面的例子我沒想到太多，就想到諸如用3個3*3的卷積代替一個7*7的卷積：感受野不變，計算量減少，但是原本一個activation變成了三個，顯存變大了。

（5）精度換計算/內存/通信

這種方法很“流氓”：深度學習模型最重要的就是精度，如果為了計算、內存和通信放棄了精度就很沒道理。

不過得益于神經網絡的超強魯棒性，很多看似大膽的做法可以在顯著降低計算，內存或通信的情況下只掉一點點精度。

這里簡單介紹幾種常見的做法：

• 量化/用低精度計算：顯而易見，如果你用Float16代替Float32，那么運行速度，需要的內存，需要的帶寬基本上都可以直接砍一半

• 稀疏通信：精度換通信的一種做法：我們每次對梯度做all reduce的時候并不需要傳所有梯度，只需要選擇一部分（比如數值比較大）的梯度傳輸就好了

• 神經網絡的各種剪枝：比如把很小的weight直接刪掉，畢竟對最終結果沒啥影響

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

總結

以上是生活随笔為你收集整理的深度学习如何均衡精度、内存、计算和通信开销？的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。