視訊記憶體不夠，框架來湊：兩行程式碼視訊記憶體翻倍，2080Ti也能當V100來用

機器之心原創

作者：思

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 下一個推出的功能會是什麼？讓我們拭目以待。