作者 | 陳大鑫、青暮
話說人工智能江湖門派眾多且繁雜，好似那大腦中的神經(jīng)網(wǎng)絡(luò)一般，但繁雜中卻仍然有著一統(tǒng)的跡象......許久之前，ML派坐落美利堅合眾山中，百年來武學(xué)奇才輩出，隱然成江湖第一大名門正派，門內(nèi)有三套入門武功，曰：圖模型加圈，神經(jīng)網(wǎng)加層，優(yōu)化目標加正則。有童謠為證：熟練ML入門功，不會作文也會謅。不期數(shù)年，北方DL神教異軍突起，內(nèi)修表示學(xué)習(xí)，外練神經(jīng)網(wǎng)絡(luò)，心法眾多，曰門，曰注意，曰記憶，曰對抗，曰增強。經(jīng)ImageNet一役威震武林，豢Alpha犬一匹無人可近。一時家家筑丹爐，人人煉丹忙，門徒云集，依附者眾，有一統(tǒng)江湖之勢。有童謠為證：左手大數(shù)據(jù)，右手英偉達，每逢頂會煉丹忙。轉(zhuǎn)至今時，竟知DL神教仍不是源頭，它背靠「大力神教」，Transformer、BERT、GPT-3 、DALL.E等等神丹相繼紛紛被祭出......丹爐越來越大、爐火越燒越旺，「大力出奇跡」已被眾多教徒們信奉為至高無上的教義。............................就這樣時間來到了2021年，訓(xùn)練大模型已幾乎成為了AI界的共識，近期MLP的涅槃重生與Transformer的一場爭斗大戲，讓大家不禁悲嘆「Attention is all you need」是假的，別的神馬也都是不確定的，唯有「Money is all you need 」是真的
。然而，因為虛擬貨幣挖礦等眾多原因，顯卡的價格一直居高不下，不容樂觀，且一度缺貨。隨手打開某電商平臺，可以看到一張2080Ti 漲到一萬五起步 ，一張 V100 更是得四五萬。可就在這樣的背景下，如果現(xiàn)在有人告訴你2080Ti可以當(dāng)作 V100來用，那得節(jié)省多少Money？你會說這怎么可能呢？V100的顯存可是有16GB ，而2080Ti的顯存只有11GB，相比于計算能力，顯存容量可是一票否決制的，雖然計算能力決定了模型訓(xùn)練的快慢，但是顯存容量決定了模型能否訓(xùn)練。即使 2080Ti 的計算能力高于 V100，但是由于顯存容量不夠，batchsize設(shè)置為128 的 ResNet50 根本無法進行訓(xùn)練，單就這點兒而言對于吃大顯存的 AI 模型就不可能繞的過去。然而人定勝天，世間上最令人興奮的事莫過于把看似不可能變成可能，好似魔法一般，其秘訣就在于下面這兩行代碼：沒錯，就是上面這兩行看似簡簡單單的代碼，它就像一個“點石成金”的按鈕，蘊含著巨大的能量，只要一鍵開啟之后，2080Ti 就可以當(dāng)作 V100來用！在一番探尋之后發(fā)現(xiàn)，原來是曠視在開源深度學(xué)習(xí)框架MegEngine 的最新版本V1.4中，新增了動態(tài)圖的顯存優(yōu)化技術(shù)（DTR）。“煉丹”時使用這項新功能后可以實現(xiàn)顯著降低顯存占用的魔法般效果。在使用DTR功能后，11G顯存的2080Ti 可以訓(xùn)練出原本32G顯存的V100才能訓(xùn)練的模型，可以節(jié)省很多Money！另外很多學(xué)生黨只有 1060 之類的4~6GB的小顯存卡，而靠 DTR就能在便宜的民用顯卡上訓(xùn)練原本需要吃10GB顯存以上的大模型了。“魔法”從何而來
曠視MegEngine 團隊在2020年下半年的時候希望由靜態(tài)圖徹底向動態(tài)圖遷移，所以開始探索動態(tài)圖上的 sublinear 類似技術(shù)，這時團隊偶然看到了華盛頓大學(xué)和卡納基梅隆大學(xué)團隊合作的一篇名為DTR論文：《Dynamic Tensor Rematerialization》，目前該論文已被 ICLR 2021接收為 Spotlight。論文鏈接：https://arxiv.org/abs/2006.09616這篇論文主要講了一種動態(tài)圖的顯存優(yōu)化技術(shù)——在前向計算時釋放保存中間結(jié)果的tensor，反向求導(dǎo)時根據(jù)計算歷史恢復(fù)之前釋放的tensor，從而達到訓(xùn)練更大模型的目的，亮點是提出了一種動態(tài)選取釋放的tensor的策略。而MegEngine框架本來就有靜態(tài)圖上的顯存優(yōu)化，功能上與 DTR 類似，經(jīng)過一番理論到工程的實現(xiàn)之后，MegEngine團隊通過DTR成功實現(xiàn)了動態(tài)圖顯存優(yōu)化技術(shù)。至于為什么MegEngine 團隊要針對動態(tài)圖做優(yōu)化，那當(dāng)然是因為動態(tài)圖代碼易寫易調(diào)試、是現(xiàn)在的主流趨勢，且隨著越來越多的深度學(xué)習(xí)框架支持動態(tài)圖模式，能否在動態(tài)圖訓(xùn)練時最大程度地利用有限的顯存資源，成為了評估深度學(xué)習(xí)框架性能的重要指標。然而對于目前業(yè)界的主流框架而言，如果想使用顯存優(yōu)化訓(xùn)練更大的模型必須要先轉(zhuǎn)靜態(tài)圖之后再用Sublinear來優(yōu)化才行。意思就是你動態(tài)圖想直接優(yōu)化顯存不行，必須先走靜態(tài)圖這個橋梁才行，這意味什么呢？舉個例子哈，假如你要去民事局和心愛的對象辦理結(jié)婚，但是民事局卻告訴你們必須先去一趟警察局登記才行，那誰能樂意呢，不知道情況的被誤會了可咋辦，多一趟手續(xù)多少會帶來一些麻煩啊。所以說，直接在簡單易用的動態(tài)圖上做顯存優(yōu)化就成為了一個迫切的需求，可以給到AI研究員和企業(yè)工程師很大的幫助，而未來 MegEngine 也將會全面遷移到動態(tài)圖上，并且會致力于優(yōu)化動態(tài)圖的性能，使得動態(tài)圖和靜態(tài)圖一樣快，甚至更快。那么在代碼里開啟動態(tài)圖優(yōu)化功能會很麻煩嗎？NO NO NO~就像前文已經(jīng)說過的一樣，真的只需添加下面兩行代碼就行：第一行代碼是聲明DTR的參數(shù)，第二行代碼是打開DTR開關(guān)，只要把這兩行放在開頭，后面的模型代碼保證不用再做任何改動。真可謂是一鍵開啟動態(tài)Sublinear顯存優(yōu)化功能，簡直方便到不行！那么優(yōu)化效果如何呢？下面通過一張動態(tài)Sublinear顯存優(yōu)化圖來展示一下：如上圖所示，這是某個AI 模型的訓(xùn)練圖，藍色線是不加任何動態(tài)圖顯存優(yōu)化的訓(xùn)練過程，在峰值條件下需要占用 16GB 的GPU顯存；綠色線是Sublinear優(yōu)化之后的靜態(tài)圖訓(xùn)練過程，在峰值條件下需要占用 8GB的GPU顯存，紅色線是開啟動態(tài)Sublinear顯存優(yōu)化之后的訓(xùn)練過程，在峰值條件下，僅需要占用4GB的GPU顯存，也就是藍線的1/4。對比藍色線和紅色線，很明顯看出原先的16GB顯存占用現(xiàn)在優(yōu)化到來4GB左右！大家都知道2080Ti 的計算能力比 V100 高，但是卻受限于顯存容量比 V100 低，而開啟了 DTR 優(yōu)化之后，相當(dāng)于用多出來的計算能力彌補了顯存容量上的不足，使得原來在 2080Ti 上無法訓(xùn)練的模型可以訓(xùn)練起來。也即等于一張 2080 Ti 可以用出 V100 的感覺！且曠視MegEngine的工程師親自做了測試，發(fā)現(xiàn)在 2080Ti上，ResNet50、ShuffleNet等網(wǎng)絡(luò)模型的最大batchsize可以達到原來的3倍以上。而且這個黑科技魔法是可疊加到多卡上的，所以說十張2080 Ti 同樣可以用出十張V100的效果。好家伙，這簡直就是訓(xùn)練 AI 大模型的神器??！因為并不是每個實驗室都有像谷歌、Fackbook這樣的土豪，而依靠曠視這項 DTR技術(shù)實現(xiàn)，可以大大節(jié)省購買顯卡的資金，能讓更多的大學(xué)、實驗室可以參與到這場不分對錯的“大力出奇跡”的 “軍備競賽”當(dāng)中來。這背后具體是什么原理呢？MegEngine團隊的工程師寫了一篇詳細的文章介紹其DTR的實現(xiàn)與優(yōu)化路徑，感興趣的讀者可以點擊文末的閱讀原文鏈接一鍵直達。實驗數(shù)據(jù)對比
對這個訓(xùn)練 AI 大模型神器的感覺還是不夠直觀？那就不如看一下更直觀的實驗數(shù)據(jù)對比。下圖是MegEngine的DTR實現(xiàn)與原論文在PyTorch中的實現(xiàn)在ResNet1202上的訓(xùn)練情況對比，注意到實驗用的顯卡不同，所以從數(shù)據(jù)上看MegEngine稍快一些。不難看出，在顯存管理上MegEngine要更好一些，因為在11G的顯卡上它仍然能跑batchsize=100 的模型訓(xùn)練。除了論文中嘗試的最大batchsize=140之外，MegEngine團隊還嘗試了更大的batchsize，也都是可以運行的。下面是在MegEngine框架上開啟不同顯存優(yōu)化的訓(xùn)練耗時對比，baseline是在動態(tài)圖模式下不加任何顯存優(yōu)化運行的結(jié)果。首先是兩個常見的模型——ResNet50和ShuffleNet的對比，可以發(fā)現(xiàn)開啟DTR優(yōu)化后極限batchsize超過了靜態(tài)圖Sublinear和baseline，且在batchsize相同時耗時和Sublinear持平。需要說明的是這里拿視覺模型只是為了方便做實驗和舉例，實際上這項優(yōu)化功能適用于CV、NLP等不同領(lǐng)域的任意 AI 模型訓(xùn)練。上面的兩個模型都是比較偏靜態(tài)的，所以可以用靜態(tài)圖的Sublinear顯存優(yōu)化來做對比，而下面這個SPOS網(wǎng)絡(luò)就比較特殊，它是一個從輸入到輸出有多條路徑可以更新的大網(wǎng)絡(luò)。在訓(xùn)練過程中，每一輪會隨機采樣去更新某一條路徑，這就導(dǎo)致每輪執(zhí)行的語句可能不相同。對于這種網(wǎng)絡(luò)，在動態(tài)圖里實現(xiàn)會比較自然。因此，這里只取了動態(tài)圖DTR優(yōu)化的結(jié)果與Baseline比較。不論是單卡還是八卡，動態(tài)圖的極限batchsize都在100，如果打開DTR可以跑到250甚至更大。硬幣都是兩面的，在MegEngine框架上實現(xiàn)的這項DTR 技術(shù)非常之香——2080Ti 可以當(dāng)作V100來用，但固然也有代價——犧牲了計算算力和計算時長：犧牲的計算時長由DTR的參數(shù)決定，最壞情況下所有 Tensor 在不被用到的時候都立即釋放，恢復(fù)每個 Tensor 的時間都是 O(N)的，總的時間就會達到 O(N^2) 級別。但一般情況下只是原來的時間的常數(shù)倍。
例如訓(xùn)練 batchsize=200 的 ResNet50 需要 16G 左右的顯存，每輪的訓(xùn)練耗時是800ms。設(shè)置DTR的閾值為7G時，只需要 11G 的顯存，訓(xùn)練耗時為 898ms；設(shè)置DTR的閾值為3G時，只需要 7.5G 的顯存，訓(xùn)練耗時為 1239ms。
但是要清楚地認識到 DTR 的優(yōu)勢不是體現(xiàn)在耗時上，而是體現(xiàn)在可以訓(xùn)練 batchsize 更大的模型。用了DTR耗時一定是會增加的，但是這個是可以容忍的，因為原本不能訓(xùn)練的大模型現(xiàn)在可以訓(xùn)練了！最后，還要介紹一下這項DTR技術(shù)的最神奇之處——顯存占用越大越吃香！也即是顯卡容量越大，優(yōu)化的效果越好，都能把任意大的內(nèi)存優(yōu)化到原來層數(shù)對應(yīng)的 1/
倍，這里的 O(1/)是對于n層的前饋神經(jīng)網(wǎng)絡(luò)的理論下界。總結(jié)
1、MegEngine是首個實現(xiàn)DTR技術(shù)的開源深度學(xué)習(xí)框架。2、一張 2080 Ti 可以用出 V100 的感覺。3、在2080Ti上，ResNet50、ShuffleNet等網(wǎng)絡(luò)的最大batchsize可以達到原來的3倍以上！4、單卡上的優(yōu)化可以疊加到多卡上一起加倍節(jié)省顯存。6、對于用戶的體驗：運行速度更快、使用體驗更好：API簡潔，用戶只用打開開關(guān)，不用改動訓(xùn)練代碼。7、對于框架開發(fā)者的體驗：得益于MegEngine框架簡潔的底層實現(xiàn)，在 MegEngine 中實現(xiàn)DTR更加簡單自然，并且便于進行擴展。（注：本文開頭二、三兩段為清華大學(xué)劉知遠老師原創(chuàng)，本文已經(jīng)過劉老師本人授權(quán)使用。）雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)
雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
