只需兩行程式碼，2080Ti 就能當 V100用，這個煉丹神器真牛！

  • 2021 年 6 月 17 日
  • AI

作者 | 陳大鑫、青暮

話說人工智慧江湖門派眾多且繁雜，好似那大腦中的神經網路一般，但繁雜中卻仍然有著一統的跡象……
許久之前，ML派坐落美利堅合眾山中，百年來武學奇才輩出，隱然成江湖第一大名門正派，門內有三套入門武功，曰：圖模型加圈，神經網加層，優化目標加正則。有童謠為證：熟練ML入門功，不會作文也會謅。
不期數年，北方DL神教異軍突起，內修表示學習，外練神經網路，心法眾多，曰門，曰注意，曰記憶，曰對抗，曰增強。經ImageNet一役威震武林，豢Alpha犬一匹無人可近。一時家家築丹爐，人人煉丹忙，門徒雲集，依附者眾，有一統江湖之勢。有童謠為證：左手大數據，右手NVIDIA ，每逢高峰會煉丹忙。
轉至今時，竟知DL神教仍不是源頭，它背靠「大力神教」，Transformer、BERT、GPT-3 、DALL.E等等神丹相繼紛紛被祭出……
丹爐越來越大、爐火越燒越旺，「大力出奇蹟」已被眾多教徒們信奉為至高無上的教義。
……………………….
就這樣時間來到了2021年，訓練大模型已幾乎成為了AI界的共識，近期MLP的涅槃重生與Transformer的一場爭鬥大戲，讓大家不禁悲嘆「Attention is all you need」是假的，別的神馬也都是不確定的，唯有「Money  is all you need 」是真的。
然而，因為虛擬貨幣挖礦等眾多原因，顯示卡的價格一直居高不下，不容樂觀，且一度缺貨。隨手打開某電商平台，可以看到一張2080Ti 漲到一萬五起步 ，一張 V100 更是得四五萬。
可就在這樣的背景下，如果現在有人告訴你2080Ti可以當作 V100來用，那得節省多少Money？
你會說這怎麼可能呢？V100的顯示記憶體可是有16GB ，而2080Ti的顯示記憶體只有11GB，相比於計算能力，顯示記憶體容量可是一票否決制的，雖然計算能力決定了模型訓練的快慢，但是顯示記憶體容量決定了模型能否訓練。即使 2080Ti 的計算能力高於 V100，但是由於顯示記憶體容量不夠，batchsize設置為128 的 ResNet50 根本無法進行訓練，單就這點兒而言對於吃大顯示記憶體的 AI 模型就不可能繞的過去。
然而人定勝天，世間上最令人興奮的事莫過於把看似不可能變成可能，好似魔法一般，其秘訣就在於下面這兩行程式碼：
沒錯，就是上面這兩行看似簡簡單單的程式碼，它就像一個「點石成金」的按鈕，蘊含著巨大的能量，只要一鍵開啟之後，2080Ti 就可以當作 V100來用！
這麼神奇的事情背後究竟是怎麼一回事呢？
在一番探尋之後發現，原來是曠視在開源深度學習框架MegEngine 的最新版本V1.4中，新增了動態圖的顯示記憶體優化技術（DTR）。
「煉丹」時使用這項新功能後可以實現顯著降低顯示記憶體佔用的魔法般效果。
在使用DTR功能後，11G顯示記憶體的2080Ti 可以訓練出原本32G顯示記憶體的V100才能訓練的模型，可以節省很多Money！另外很多學生黨只有 1060 之類的4~6GB的小顯示記憶體卡，而靠 DTR就能在便宜的民用顯示卡上訓練原本需要吃10GB顯示記憶體以上的大模型了。

1

「魔法」從何而來

曠視MegEngine 團隊在2020年下半年的時候希望由靜態圖徹底向動態圖遷移，所以開始探索動態圖上的 sublinear 類似技術，這時團隊偶然看到了華盛頓大學和卡納基梅隆大學團隊合作的一篇名為DTR論文：《Dynamic Tensor Rematerialization》，目前該論文已被 ICLR 2021接收為 Spotlight。
論文鏈接：//arxiv.org/abs/2006.09616
這篇論文主要講了一種動態圖的顯示記憶體優化技術——在前向計算時釋放保存中間結果的tensor，反向求導時根據計算歷史恢復之前釋放的tensor，從而達到訓練更大模型的目的，亮點是提出了一種動態選取釋放的tensor的策略。
而MegEngine框架本來就有靜態圖上的顯示記憶體優化，功能上與 DTR 類似，經過一番理論到工程的實現之後，MegEngine團隊通過DTR成功實現了動態圖顯示記憶體優化技術。
至於為什麼MegEngine 團隊要針對動態圖做優化，那當然是因為動態圖程式碼易寫易調試、是現在的主流趨勢，且隨著越來越多的深度學習框架支援動態圖模式，能否在動態圖訓練時最大程度地利用有限的顯示記憶體資源，成為了評估深度學習框架性能的重要指標。然而對於目前業界的主流框架而言，如果想使用顯示記憶體優化訓練更大的模型必須要先轉靜態圖之後再用Sublinear來優化才行。
意思就是你動態圖想直接優化顯示記憶體不行，必須先走靜態圖這個橋樑才行，這意味什麼呢？舉個例子哈，假如你要去民事局和心愛的對象辦理結婚，但是民事局卻告訴你們必須先去一趟警察局登記才行，那誰能樂意呢，不知道情況的被誤會了可咋辦，多一趟手續多少會帶來一些麻煩啊。
所以說，直接在簡單易用的動態圖上做顯示記憶體優化就成為了一個迫切的需求，可以給到AI研究員和企業工程師很大的幫助，而未來 MegEngine 也將會全面遷移到動態圖上，並且會致力於優化動態圖的性能，使得動態圖和靜態圖一樣快，甚至更快。
那麼在程式碼里開啟動態圖優化功能會很麻煩嗎？
NO NO NO~就像前文已經說過的一樣，真的只需添加下面兩行程式碼就行：
第一行程式碼是聲明DTR的參數，第二行程式碼是打開DTR開關，只要把這兩行放在開頭，後面的模型程式碼保證不用再做任何改動。
真可謂是一鍵開啟動態Sublinear顯示記憶體優化功能，簡直方便到不行！
那麼優化效果如何呢？下面通過一張動態Sublinear顯示記憶體優化圖來展示一下：
如上圖所示，這是某個AI 模型的訓練圖，藍色線是不加任何動態圖顯示記憶體優化的訓練過程，在峰值條件下需要佔用 16GB 的GPU顯示記憶體；
綠色線是Sublinear優化之後的靜態圖訓練過程，在峰值條件下需要佔用 8GB的GPU顯示記憶體，
紅色線是開啟動態Sublinear顯示記憶體優化之後的訓練過程，在峰值條件下，僅需要佔用4GB的GPU顯示記憶體，也就是藍線的1/4。
對比藍色線和紅色線，很明顯看出原先的16GB顯示記憶體佔用現在優化到來4GB左右！
大家都知道2080Ti 的計算能力比 V100 高，但是卻受限於顯示記憶體容量比 V100 低，而開啟了 DTR 優化之後，相當於用多出來的計算能力彌補了顯示記憶體容量上的不足，使得原來在 2080Ti 上無法訓練的模型可以訓練起來。
也即等於一張 2080 Ti 可以用出 V100 的感覺！
且曠視MegEngine的工程師親自做了測試，發現在 2080Ti上，ResNet50、ShuffleNet等網路模型的最大batchsize可以達到原來的3倍以上。
而且這個黑科技魔法是可疊加到多卡上的，所以說十張2080 Ti 同樣可以用出十張V100的效果。
好傢夥，這簡直就是訓練 AI 大模型的神器啊！
因為並不是每個實驗室都有像Google、Fackbook這樣的土豪，而依靠曠視這項 DTR技術實現，可以大大節省購買顯示卡的資金，能讓更多的大學、實驗室可以參與到這場不分對錯的「大力出奇蹟」的 「軍備競賽」當中來。
這背後具體是什麼原理呢？MegEngine團隊的工程師寫了一篇詳細的文章介紹其DTR的實現與優化路徑，感興趣的讀者可以點擊文末的閱讀原文鏈接一鍵直達。

2

實驗數據對比

對這個訓練 AI 大模型神器的感覺還是不夠直觀？
那就不如看一下更直觀的實驗數據對比。
首先是訓練耗時對比。
下圖是MegEngine的DTR實現與原論文在PyTorch中的實現在ResNet1202上的訓練情況對比，注意到實驗用的顯示卡不同，所以從數據上看MegEngine稍快一些。
不難看出，在顯示記憶體管理上MegEngine要更好一些，因為在11G的顯示卡上它仍然能跑batchsize=100 的模型訓練。除了論文中嘗試的最大batchsize=140之外，MegEngine團隊還嘗試了更大的batchsize，也都是可以運行的。
下面是在MegEngine框架上開啟不同顯示記憶體優化的訓練耗時對比，baseline是在動態圖模式下不加任何顯示記憶體優化運行的結果。
首先是兩個常見的模型——ResNet50和ShuffleNet的對比，可以發現開啟DTR優化後極限batchsize超過了靜態圖Sublinear和baseline，且在batchsize相同時耗時和Sublinear持平。
需要說明的是這裡拿視覺模型只是為了方便做實驗和舉例，實際上這項優化功能適用於CV、NLP等不同領域的任意 AI 模型訓練。
上面的兩個模型都是比較偏靜態的，所以可以用靜態圖的Sublinear顯示記憶體優化來做對比，而下面這個SPOS網路就比較特殊，它是一個從輸入到輸出有多條路徑可以更新的大網路。
在訓練過程中，每一輪會隨機取樣去更新某一條路徑，這就導致每輪執行的語句可能不相同。對於這種網路，在動態圖裡實現會比較自然。因此，這裡只取了動態圖DTR優化的結果與Baseline比較。不論是單卡還是八卡，動態圖的極限batchsize都在100，如果打開DTR可以跑到250甚至更大。
硬幣都是兩面的，在MegEngine框架上實現的這項DTR 技術非常之香——2080Ti 可以當作V100來用，但固然也有代價——犧牲了計算算力和計算時長：

  • 犧牲的計算時長由DTR的參數決定，最壞情況下所有 Tensor 在不被用到的時候都立即釋放，恢復每個 Tensor 的時間都是 O(N)的，總的時間就會達到 O(N^2) 級別。但一般情況下只是原來的時間的常數倍。

  • 例如訓練 batchsize=200 的 ResNet50 需要 16G 左右的顯示記憶體，每輪的訓練耗時是800ms。設置DTR的閾值為7G時，只需要 11G 的顯示記憶體，訓練耗時為 898ms；設置DTR的閾值為3G時，只需要 7.5G 的顯示記憶體，訓練耗時為 1239ms。

但是要清楚地認識到 DTR 的優勢不是體現在耗時上，而是體現在可以訓練 batchsize 更大的模型。用了DTR耗時一定是會增加的，但是這個是可以容忍的，因為原本不能訓練的大模型現在可以訓練了！
最後，還要介紹一下這項DTR技術的最神奇之處——顯示記憶體佔用越大越吃香！
也即是顯示卡容量越大，優化的效果越好，都能把任意大的記憶體優化到原來層數對應的 1/倍，這裡的 O(1/)是對於n層的前饋神經網路的理論下界。

3

總結

1、MegEngine是首個實現DTR技術的開源深度學習框架。
2、一張 2080 Ti 可以用出 V100 的感覺。
3、在2080Ti上，ResNet50、ShuffleNet等網路的最大batchsize可以達到原來的3倍以上！
4、單卡上的優化可以疊加到多卡上一起加倍節省顯示記憶體。
5、顯示記憶體佔用越大越吃香。
6、對於用戶的體驗：運行速度更快、使用體驗更好：API簡潔，用戶只用打開開關，不用改動訓練程式碼。
7、對於框架開發者的體驗：得益於MegEngine框架簡潔的底層實現，在 MegEngine 中實現DTR更加簡單自然，並且便於進行擴展。
（註：本文開頭二、三兩段為清華大學劉知遠老師原創，本文已經過劉老師本人授權使用。）
雷鋒網雷鋒網雷鋒網

VirMach 便宜 VPS

QNews

QNews