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

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

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

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

為什么GPU顯存如此重要？

九年前，Hinton等人率先用兩張3GB顯存的GTX580GPU高效訓練AlexNet。在此之后，顯存需求與模型大小就一直同步增長。打比賽想要取到好成績、做實驗想要超越Stateoftheart效果、做工程想要擬合龐大的業(yè)務(wù)數(shù)據(jù)等等，這些都離不開顯存的加持。

模型加一層，顯存漲一分

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

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

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

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

遇到這種情況，學生無非只有兩種選擇：向?qū)熒暾埿碌腉PU資源，或者縮減模型做一個Mini版的實驗。前者并不總是能夠成功，后者則可能會有種種不完美。如果能用有限的顯存跑頂尖的大模型，做實驗、寫論文都會變得更加簡單。

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

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

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

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

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

MegEngine：顯存需要優(yōu)化

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

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

具體來說，通過復現(xiàn)并優(yōu)化ICLR2021Spotlight論文《DynamicTensorRematerialization》（以下簡稱DTR），MegEngine實現(xiàn)了「用計算換取更多顯存」。有了這項技術(shù)的加持，模型的顯存占用大大降低，同樣的硬件可以訓練更大的模型、承載更大的BatchSize。如此一來，學生的小顯卡也能開始訓練大模型，而工程師們的服務(wù)器也經(jīng)得起更充分的應用。

圖：原本需要16GB顯存的模型，優(yōu)化后使用的顯存峰值就降到了4GB

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

兩行代碼，顯存「翻倍」

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

通過兩行代碼，框架可以全自動地完成顯存優(yōu)化，將所有優(yōu)化邏輯與復雜的工程實現(xiàn)都隱藏在MegEngine內(nèi)部。

如上圖所示，在動態(tài)計算圖中導入DTR顯存優(yōu)化模塊，并配置顯存釋放閾值為5GB。訓練時，因為顯存已經(jīng)「翻倍」了，BatchSize翻四倍也能裝到GPU中。

顯存擴增帶來的收益

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

最簡單的驗證方法就是不斷增加批大小，看看顯卡到底能堅持到什么程度。下面兩張表分別展示了在PyTorch及MegEngine上加載或不加載動態(tài)圖顯存優(yōu)化（DTR）技術(shù)的效果。

如果不使用動態(tài)圖顯存優(yōu)化技術(shù)，PyTorch上的模型一次訓練迭代最多只能處理64個樣本，MegEngine能處理100個樣本。只要加上DTR，PyTorch模型一次迭代就能處理140個樣本，MegEngine能嘗試處理300個樣本。

如果換算成模型大小，加上動態(tài)圖顯存優(yōu)化技術(shù)的MegEngine，在相同的GPU及批大小情況下，能高效訓練增大近乎5倍的模型。

MegEngine動態(tài)圖顯存優(yōu)化技術(shù)

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

權(quán)重矩陣和梯度矩陣占的內(nèi)存很難優(yōu)化，各個模型基本上都有一個定值。前向傳播的中間計算結(jié)果則不然：隨著BatchSize的增加以及模型層和數(shù)量的增加，顯存必然跟著增加。如果模型比較大，中間計算結(jié)果將占據(jù)最主要的顯存。

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

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

很明顯，如果要降低顯存占用，就要拿前向傳播保存的中間計算結(jié)果開刀，這也正是MegEngine動態(tài)圖顯存優(yōu)化的主要方向。

用計算換顯存

對于動態(tài)計算圖，最直接的方法就是用計算或內(nèi)存換顯存。因此，MegEngine首先要決定到底使用哪種技術(shù)。

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

圖：因此很明確，動態(tài)計算圖中也應該使用梯度檢查點技術(shù)，用計算換顯存

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

雖然大致原理不難理解，但具體怎么做還是比較復雜的，MegEngine團隊借鑒了論文《DynamicTensorRematerialization》，將其優(yōu)化并實現(xiàn)到MegEngine中。

DTR，最前沿的顯存優(yōu)化技術(shù)

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

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

不考慮小的張量，它們不加入候選集

每次在需要釋放張量的時候，隨機采樣并遍歷少部分張量，以節(jié)省計算開銷

最難的是工程實現(xiàn)

雖然DTR看上去原理也不復雜，但真正的難題在于提高易用性，即將所有細節(jié)都隱藏到框架的底層，只為開發(fā)者提供最簡單的接口。

在此就用一個最簡單的計算例子，跟著框架演算一遍，看看MegEngine是如何利用動態(tài)圖的計算歷史恢復與釋放張量的。

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

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

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

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

1     structComputePath{

2    std：：shared_ptr<OpDef>op；

3    SmallVector<TensorInfo*>inputs；

4    SmallVector<TensorInfo*>outputs；

5    doublecompute_time=0；

6     }*producer；

7    SmallVector<ComputePath*>users；

8     size_tref_cnt=0；

以上為MegEngine底層用于追蹤計算路徑信息的結(jié)構(gòu)體。其中op表示產(chǎn)生該張量的算子；inputs和outputs分別表示這個算子需要的輸入與輸出張量；compute_time表示該算子實際的運行時間。

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

圖：Python代碼會翻譯成C++底層實現(xiàn)，C++代碼會通過指針管理顯卡內(nèi)存中真正的張量（右圖綠色部分）

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

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