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