作者 | 陳大鑫、青暮
話說人工智能江湖門派眾多且繁雜，好似那大腦中的神經網絡一般，但繁雜中卻仍然有著一統的跡象......許久之前，ML派坐落美利堅合眾山中，百年來武學奇才輩出，隱然成江湖第一大名門正派，門內有三套入門武功，曰：圖模型加圈，神經網加層，優化目標加正則。有童謠為證：熟練ML入門功，不會作文也會謅。不期數年，北方DL神教異軍突起，內修表示學習，外練神經網絡，心法眾多，曰門，曰注意，曰記憶，曰對抗，曰增強。經ImageNet一役威震武林，豢Alpha犬一匹無人可近。一時家家筑丹爐，人人煉丹忙，門徒云集，依附者眾，有一統江湖之勢。有童謠為證：左手大數據，右手英偉達，每逢頂會煉丹忙。轉至今時，竟知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顯存以上的大模型了。“魔法”從何而來
曠視MegEngine 團隊在2020年下半年的時候希望由靜態圖徹底向動態圖遷移，所以開始探索動態圖上的 sublinear 類似技術，這時團隊偶然看到了華盛頓大學和卡納基梅隆大學團隊合作的一篇名為DTR論文：《Dynamic Tensor Rematerialization》，目前該論文已被 ICLR 2021接收為 Spotlight。論文鏈接：https://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的效果。因為并不是每個實驗室都有像谷歌、Fackbook這樣的土豪，而依靠曠視這項 DTR技術實現，可以大大節省購買顯卡的資金，能讓更多的大學、實驗室可以參與到這場不分對錯的“大力出奇跡”的 “軍備競賽”當中來。這背后具體是什么原理呢？MegEngine團隊的工程師寫了一篇詳細的文章介紹其DTR的實現與優化路徑，感興趣的讀者可以點擊文末的閱讀原文鏈接一鍵直達。實驗數據對比
下圖是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層的前饋神經網絡的理論下界。總結
1、MegEngine是首個實現DTR技術的開源深度學習框架。2、一張 2080 Ti 可以用出 V100 的感覺。3、在2080Ti上，ResNet50、ShuffleNet等網絡的最大batchsize可以達到原來的3倍以上！4、單卡上的優化可以疊加到多卡上一起加倍節省顯存。6、對于用戶的體驗：運行速度更快、使用體驗更好：API簡潔，用戶只用打開開關，不用改動訓練代碼。7、對于框架開發者的體驗：得益于MegEngine框架簡潔的底層實現，在 MegEngine 中實現DTR更加簡單自然，并且便于進行擴展。（注：本文開頭二、三兩段為清華大學劉知遠老師原創，本文已經過劉老師本人授權使用。）
雷峰網原創文章，未經授權禁止轉載。詳情見轉載須知。
