機器之心原創

作者：思

2080Ti 竟然可以當 V100 來用，這個功能有點兒厲害。

自深度學習大潮興起，模型就朝着越來越大、越來越 「深」 的方向發展。

2012 年，擁有 5 個卷積層的 AlexNet 第一次在視覺任務上展現出強大的能力。在此之後，基礎模型就開始「深」化起來：2014 年的 VGG-Net 達到了 19 層；2015 年的 ResNet、2017 年的 DenseNet 更是将深度提升到了上百層。

模型大小的提升極大地提高了性能。因此，各大視覺任務都将 ResNet、DenseNet 等當做基本的 BackBone。但與此同時，模型的增大也意味着對顯存的需求随之變高。

為什麼 GPU 顯存如此重要？

九年前，Hinton 等人率先用兩張 3GB 顯存的 GTX 580 GPU 高效訓練 AlexNet。在此之後，顯存需求與模型大小就一直同步增長。打比賽想要取到好成績、做實驗想要超越 State of the art 效果、做工程想要拟合龐大的業務數據等等，這些都離不開顯存的加持。

模型加一層，顯存漲一分

在深度學習模型中，占用顯存的總是那些特别大的張量，比如各層的權重矩陣、計算出來的張量（激活值）、反向傳播需要的張量等。在視覺任務中，占據絕大多數的是中間計算出來的張量。随着模型變得更深更大，每一層的激活值張量都需要保留在顯存中。

以 ResNet50 為例，在模型的訓練中，前向傳播中 50 層的計算結果都需要保存在顯存中，以便讓反向傳播利用這些張量計算梯度。如果使用 ResNet108，需要的顯存就會比 ResNet50 多出一倍多。顯存的增加，帶來的當然是模型效果的提升。另一方面，如果顯存不夠，許多工作也必将無法實現。

顯存不夠，寫論文、打比賽屢遭掣肘

在實驗室跑模型、寫論文的過程中，顯存不夠用也是常有的事。一般實驗室的顯卡都是大家共用的，可能分配到每個人的手上已經所剩無幾。甚至于，随着頂尖模型越來越大，所有人都沒有足夠的算力、顯存去複現終極實驗，更不用說超越其 SOTA 結果。

遇到這種情況，學生無非隻有兩種選擇：向導師申請新的 GPU 資源，或者縮減模型做一個 Mini 版的實驗。前者并不總是能夠成功，後者則可能會有種種不完美。如果能用有限的顯存跑頂尖的大模型，做實驗、寫論文都會變得更加簡單。

此外，無論是在學校還是在公司打比賽，算力不夠、顯存不足都是常有的事。頂尖競争者的模型結構可能相差無幾，區别就在于誰的模型更大、更有能力去處理複雜的樣本。更直觀地說，排行榜領先者的模型也許就隻差十幾層，但也正是因為顯存受限少了那十幾層，有些模型才與冠軍失之交臂。

顯存：約束算法工程師的瓶頸

再舉一個常見的例子，企業中的算法工程師擁有足夠的算力，顯存沒那麼重要。然而，隻使用并行策略分擔顯存，還是可能會出現顯存足夠、但每張 GPU 的計算負載又不足的情況。

4 張 V100，顯存占滿，而 GPU 利用率很低。

即使是 V100 這樣強大的算力，訓練大模型時也很容易占滿 16GB 顯存。然而由于批量不夠大，上圖每張 V100 GPU 的利用率隻有 20% 到 30%。隻有繼續增大每次疊代的數據吞吐量，才能增加 GPU 的利用率。

MegEngine：顯存需要優化

其實對于深度學習從業者來說，日常應用中出現的情況遠不止上面三種。做深度學習，不論是研究還是工程，時不時就會遇到顯存問題。但這個問題優化起來又很複雜，需要利用大量的工程實現來緩解。顯然，這樣的優化應該由深度學習框架來完成。不過，在實際應用中不難發現，TensorFlow、PyTorch 似乎都沒有提供完善的官方解決方案。

但如果把目光投向新生勢力，情況可能就不一樣了。在曠視開源深度學習框架 MegEngine 最近發布的 1.4 版本中，該框架首次引入了動态圖顯存優化技術，大大降低了顯存占用問題。

具體來說，通過複現并優化 ICLR 2021 Spotlight 論文《Dynamic Tensor Rematerialization》（以下簡稱 DTR），MegEngine 實現了「用計算換取更多顯存」。有了這項技術的加持，模型的顯存占用大大降低，同樣的硬件可以訓練更大的模型、承載更大的 BatchSize。如此一來，學生的小顯卡也能開始訓練大模型，而工程師們的服務器也經得起更充分的應用。

原本需要 16GB 顯存的模型，優化後使用的顯存峰值就降到了 4GB。

MegEngine 這種顯存優化技術，讓 1060 這樣的入門級顯卡也能訓練原本 2080Ti 才能加載得上的模型；而 11GB 顯存的 2080Ti，更能挑戰原本 32GB V100 才能訓練的模型。要知道，V100 的價格可是 2080Ti 的 9 倍還多。

兩行代碼，顯存「翻倍」

如要需要自己去優化顯存， 可能 99% 的算法工程師都會放棄。最好的辦法是告訴深度學習框架，這次訓練就分配多少顯存，剩下的就交給框架自己去優化。MegEngine 的動态圖顯存優化就是基于這一邏輯。

通過兩行代碼，框架可以全自動地完成顯存優化，将所有優化邏輯與複雜的工程實現都隐藏在 MegEngine 内部。

如上圖所示，在動态計算圖中導入 DTR 顯存優化模塊，并配置顯存釋放阈值為 5GB。訓練時，因為顯存已經「翻倍」了，Batch Size 翻四倍也能裝到 GPU 中。

顯存擴增帶來的收益

很多時候，提高顯存的利用率，最顯著的作用就是能訓練更大的模型。從一定程度上來說，參數量越大就意味着效果越好；而批大小越大，梯度更新方向就越準确，模型性能也就越優異。MegEngine 開發團隊做了很多實驗，以确保提高顯存利用率的同時訓練是優質的。

最簡單的驗證方法就是不斷增加批大小，看看顯卡到底能堅持到什麼程度。下面兩張表分别展示了在 PyTorch 及 MegEngine 上加載或不加載動态圖顯存優化（DTR）技術的效果。

如果不使用動态圖顯存優化技術，PyTorch 上的模型一次訓練疊代最多隻能處理 64 個樣本，MegEngine 能處理 100 個樣本。隻要加上 DTR，PyTorch 模型一次疊代就能處理 140 個樣本，MegEngine 能嘗試處理 300 個樣本。

如果換算成模型大小，加上動态圖顯存優化技術的 MegEngine，在相同的 GPU 及批大小情況下，能高效訓練增大近乎 5 倍的模型。

MegEngine 動态圖顯存優化技術

深度學習模型的顯存占用一般分為權重矩陣、前向傳播的中間張量、反向傳播的梯度矩陣（Adam 優化器）三部分。

權重矩陣和梯度矩陣占的内存很難優化，各個模型基本上都有一個定值。前向傳播的中間計算結果則不然：随着 Batch Size 的增加以及模型層和數量的增加，顯存必然跟着增加。如果模型比較大，中間計算結果将占據最主要的顯存。

如上圖所示，在前向傳播中（第一行從左到右），藍色圓圈表示模型的中間計算結果開始占用顯存。一直到前向傳播完成，第一行完全變為藍色圓圈，前面計算所占用的顯存都不能釋放。

等到反向傳播開始（第二行從右到左），随着梯度的計算與完成應用，前向傳播保留在顯存中的張量才可以釋放。

很明顯，如果要降低顯存占用，就要拿前向傳播保存的中間計算結果開刀，這也正是 MegEngine 動态圖顯存優化的主要方向。

用計算換顯存

對于動态計算圖，最直接的方法就是用計算或内存換顯存。因此，MegEngine 首先要決定到底使用哪種技術。

MegEngine 團隊通過實驗發現，用計算耗時遠比交換耗時少。例如從顯存中節省 612.5MB 空間，用帶寬換顯存要比用計算換顯存慢了幾十上百倍。

因此很明确，動态計算圖中也應該使用梯度檢查點技術，用計算換顯存。

如下為梯度檢查點技術原理示意，前向傳播中第三個點為檢查點，它會一直保存在顯存中。第四個點在完成計算後即可釋放顯存，在反向傳播中如果需要第四個點的值，可以從第三個點重新計算出第四個點的值。

雖然大緻原理不難理解，但具體怎麼做還是比較複雜的，MegEngine 團隊借鑒了論文《Dynamic Tensor Rematerialization》，将其優化并實現到 MegEngine 中。

DTR，最前沿的顯存優化技術

DTR 是一種完全動态的啟發式策略，核心思想是當顯存超過某個阈值時，動态地釋放一些合适的張量，直到顯存低于阈值。一般而言，釋放張量的标準有三個：重新計算出該張量的開銷越小越好；占用的顯存越大越好；在顯存中停留的時間越長越好。

除去從檢查點恢複前向傳播結果張量帶來的主要開銷，DTR 的額外開銷在于尋找應該被釋放的最優張量，即計算上圖張量 t 的 f(t)值。為了降低這一部分的計算量，MegEngine 還采用了兩種運行時優化：

• 不考慮小的張量，它們不加入候選集
• 每次在需要釋放張量的時候，随機采樣并遍曆少部分張量，以節省計算開銷

最難的是工程實現

雖然 DTR 看上去原理也不複雜，但真正的難題在于提高易用性，即将所有細節都隐藏到框架的底層，隻為開發者提供最簡單的接口。

在此就用一個最簡單的計算例子，跟着框架演算一遍，看看 MegEngine 是如何利用動态圖的計算曆史恢複與釋放張量的。

現在假設輸入有 a 和 b 兩個張量，并希望計算 a*b 與 a b，但是顯存最大隻能保存三個張量。在黃框計算 c=a b 時，顯存還能保留張量 c，然而在下一步綠框計算 d=a*b 時隻能先釋放 c 才能保存 d。

不巧的是，下一步灰框需要獲取黃框的計算結果，然而為了節省顯存，c 已經被釋放了。所以，MegEngine 現在需要做的是重新運行灰框的計算圖，計算 c=a b，并加載到顯存中。顯然，這樣做必然需要釋放 d 的顯存。

這樣一來，鑒于顯存的限制，MegEngine 就會自動選擇合适的張量釋放，并在需要時重新計算。如果需要重新計算某個張量的結果，例如上圖的 d，就需要具體的曆史計算信息（在這裡就是 a b 這樣的計算路徑），與此同時還需要知道 a 和 b 這兩個輸入張量。

所有這樣的曆史計算信息都由 MegEngine 自動獲取與保存，MegEngine 的工程師已經在底層用 C 處理完畢，用戶完全不需要考慮。

struct ComputePath { std::shared_ptr<OpDef> op; SmallVector<TensorInfo*> inputs; SmallVector<TensorInfo*> outputs; double compute_time = 0; } *producer; SmallVector<ComputePath*> users; size_t ref_cnt = 0;

以上為 MegEngine 底層用于追蹤計算路徑信息的結構體。其中 op 表示産生該張量的算子；inputs 和 outputs 分别表示這個算子需要的輸入與輸出張量；compute_time 表示該算子實際的運行時間。

實際上，在使用 MegEngine 的過程中，全都是用 Python 接口創建張量，隻不過框架會對應追蹤每個張量的具體信息。每當需要訪問張量，不用考慮張量是否在顯存中時，沒有也能立刻恢複出來。所有這些複雜的工程化的操作與運算邏輯都隐藏在了 MegEngine C 底層。

Python 代碼會翻譯成 C 底層實現，C 代碼會通過指針管理顯卡内存中真正的張量（右圖綠色部分）。

幸好這樣的複雜操作不需要算法工程師完成，都交給 MegEngine 好了。

MegEngine 能做的事情遠不止于此，隻不過大多是像動态圖顯存優化這種技術一樣，潤物細無聲地把用戶的實際問題解決于無形。2020 年 3 月開源的 MegEngine 在以肉眼可見的速度快速成長，從靜态計算圖到動态計算圖，再到持續提升的訓練能力、移動端推理性能優化、動态顯存優化…… 這也許就是開源的魅力。隻有不斷優化和創新，才能吸引和滿足「挑剔」的開發者。MegEngine 下一個推出的功能會是什麼？讓我們拭目以待。

更多精彩资讯请关注tft每日頭條，我们将持续为您更新最新资讯!