MegEngine亞線性顯存優(yōu)化
MegEngine經(jīng)過工程擴(kuò)展和優(yōu)化，發(fā)展出一套行之有效的加強(qiáng)版亞線性顯存優(yōu)化技術(shù)，既可在計(jì)算存儲(chǔ)資源受限的條件下，輕松訓(xùn)練更深的模型，又可使用更大batch size，進(jìn)一步提升模型性能，穩(wěn)定batchwise算子。使用MegEngine訓(xùn)練ResNet18/ResNet50，顯存占用分別最高降低23%/40%；在更大的Bert模型上，降幅更是高達(dá)75%，而額外的計(jì)算開銷幾乎不變。
基于梯度檢查點(diǎn)的亞線性顯存優(yōu)化方法[1]由于較高的計(jì)算/顯存性價(jià)比受到關(guān)注。MegEngine經(jīng)過工程擴(kuò)展和優(yōu)化，發(fā)展出一套行之有效的加強(qiáng)版亞線性顯存優(yōu)化技術(shù)，既可在計(jì)算存儲(chǔ)資源受限的條件下，輕松訓(xùn)練更深的模型，又可使用更大batch size，進(jìn)一步提升模型性能，穩(wěn)定batchwise算子。使用MegEngine訓(xùn)練ResNet18/ResNet50，顯存占用分別最高降低23%/40%；在更大的Bert模型上，降幅更是高達(dá)75%，而額外的計(jì)算開銷幾乎不變。該技術(shù)已在MegEngine開源，歡迎大家上手使用：https://github.com/MegEngine。
深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練是一件復(fù)雜的事情，它體現(xiàn)為模型的時(shí)間復(fù)雜度和空間復(fù)雜度，分別對(duì)應(yīng)著計(jì)算和內(nèi)存；而訓(xùn)練時(shí)內(nèi)存占用問題是漂浮在深度學(xué)習(xí)社區(qū)上空的一塊烏云，如何撥云見日，最大降低神經(jīng)網(wǎng)絡(luò)訓(xùn)練的內(nèi)存占用，是一個(gè)繞不開的課題。
GPU顯卡等硬件為深度學(xué)習(xí)提供了必需的算力，但硬件自身有限的存儲(chǔ)，限制了可訓(xùn)練模型的尺寸，尤其是大型深度網(wǎng)絡(luò)，由此誕生出一系列相關(guān)技術(shù)，比如亞線性顯存優(yōu)化、梯度累加、混合精度訓(xùn)練、分布式訓(xùn)練，進(jìn)行GPU顯存優(yōu)化。
其中，亞線性顯存優(yōu)化方法[1]由于較高的計(jì)算/顯存性價(jià)比備受關(guān)注；曠視基于此，經(jīng)過工程擴(kuò)展和優(yōu)化，發(fā)展出加強(qiáng)版的MegEngine亞線性顯存優(yōu)化技術(shù)，輕松把大模型甚至超大模型裝進(jìn)顯存，也可以毫無壓力使用大batch訓(xùn)練模型。
這里將圍繞著深度學(xué)習(xí)框架MegEngine亞線性顯存優(yōu)化技術(shù)的工程實(shí)現(xiàn)和實(shí)驗(yàn)數(shù)據(jù)，從技術(shù)背景、原理、使用、展望等多個(gè)方面進(jìn)行首次深入解讀。
背景
在深度學(xué)習(xí)領(lǐng)域中，隨著訓(xùn)練數(shù)據(jù)的增加，需要相應(yīng)增加模型的尺寸和復(fù)雜度，進(jìn)行模型「擴(kuò)容」；而ResNet [2] 等技術(shù)的出現(xiàn)在算法層面掃清了訓(xùn)練深度模型的障礙。不斷增加的數(shù)據(jù)和持續(xù)創(chuàng)新的算法給深度學(xué)習(xí)框架帶來了新挑戰(zhàn)，能否在模型訓(xùn)練時(shí)有效利用有限的計(jì)算存儲(chǔ)資源，尤其是減少GPU顯存占用，是評(píng)估深度學(xué)習(xí)框架性能的重要指標(biāo)。
在計(jì)算存儲(chǔ)資源一定的情況下，深度學(xué)習(xí)框架有幾種降低顯存占用的常用方法，其示例如下：
? 通過合適的梯度定義，讓算子的梯度計(jì)算不再依賴于前向計(jì)算作為輸入，從而in-place地完成算子的前向計(jì)算，比如Sigmoid、Relu等；
? 在生命周期沒有重疊的算子之間共享顯存；
? 通過額外的計(jì)算減少顯存占用，比如利用梯度檢查點(diǎn)重新計(jì)算中間結(jié)果的亞線性顯存優(yōu)化方法[1]；
? 通過額外的數(shù)據(jù)傳輸減少顯存占用，比如把暫時(shí)不用的數(shù)據(jù)從GPU交換到CPU，需要時(shí)再從CPU交換回來。
上述顯存優(yōu)化技術(shù)在MegEngine中皆有不同程度的實(shí)現(xiàn)，這里重點(diǎn)討論基于梯度檢查點(diǎn)的亞線性顯存優(yōu)化技術(shù)。
原理
一個(gè)神經(jīng)網(wǎng)絡(luò)模型所占用的顯存空間大體分為兩個(gè)方面：1）模型本身的參數(shù)，2）模型訓(xùn)練臨時(shí)占用的空間，包括參數(shù)的梯度、特征圖等。其中最大占比是 2）中以特征圖形式存在的中間結(jié)果，比如，從示例[1]可知，根據(jù)實(shí)現(xiàn)的不同，從70%到90%以上的顯存用來存儲(chǔ)特征圖。這里的訓(xùn)練過程又可分為前向計(jì)算，反向計(jì)算和優(yōu)化三個(gè)方面，其中前向計(jì)算的中間結(jié)果最占顯存，還有反向計(jì)算的梯度。第 1）方面模型自身的參數(shù)內(nèi)存占用最小。 MegEngine加強(qiáng)版亞線性顯存優(yōu)化技術(shù)借鑒了[1]的方法，尤其適用于計(jì)算存儲(chǔ)資源受限的情況，比如一張英偉達(dá)2080Ti，只有11G的顯存；而更貴的Tesla V100，最大顯存也只有32G。

圖1：亞線性顯存優(yōu)化原理，其中 (b) 保存了Relu結(jié)果，實(shí)際中Relu結(jié)果可用in-place計(jì)算
圖 1(a) 給出了卷積神經(jīng)網(wǎng)絡(luò)的基本單元，它由Conv-BN-Relu組成。可以看到，反向計(jì)算梯度的過程依賴于前向計(jì)算獲取的中間結(jié)果，一個(gè)網(wǎng)絡(luò)需要保存的中間結(jié)果與其大小成正比，即顯存復(fù)雜度為O(n)。本質(zhì)上，亞線性顯存優(yōu)化方法是以時(shí)間換空間，以計(jì)算換顯存，如圖 1(b) 所示，它的算法原理如下：
? 選取神經(jīng)網(wǎng)絡(luò)中k個(gè)檢查點(diǎn)，從而把網(wǎng)絡(luò)分成k個(gè)block，需要注意的是，初始輸入也作為一個(gè)檢查點(diǎn)；前向計(jì)算過程中只保存檢查點(diǎn)處的中間結(jié)果；
? 反向計(jì)算梯度的過程中，首先從相應(yīng)檢查點(diǎn)出發(fā)，重新計(jì)算單個(gè)block需要的中間結(jié)果，然后計(jì)算block內(nèi)部各個(gè)block的梯度；不同block的中間結(jié)果計(jì)算共享顯存。
這種方法有著明顯的優(yōu)點(diǎn)，即大幅降低了模型的空間復(fù)雜度，同時(shí)缺點(diǎn)是增加了額外的計(jì)算：
? 顯存占用從O(n)變成O(n/k)+ O(k)，O(n/k)代表計(jì)算單個(gè)節(jié)點(diǎn)需要的顯存，O(k)代表k個(gè)檢查點(diǎn)需要的顯存， 取k=sqrt(n)，O(n/k)+ O(k)~O(sqrt(n))，可以看到顯存占用從線性變成了亞線性；
? 因?yàn)樵诜聪蛱荻鹊挠?jì)算過程中需要從檢查點(diǎn)恢復(fù)中間結(jié)果，整體需要額外執(zhí)行一次前向計(jì)算。
工程
在[1]的基礎(chǔ)上，MegEngine結(jié)合自身實(shí)踐，做了工程擴(kuò)展和優(yōu)化，把亞線性顯存優(yōu)化方法擴(kuò)展至任意的計(jì)算圖，并結(jié)合其它常見的顯存優(yōu)化方法，發(fā)展出一套行之有效的加強(qiáng)版亞線性顯存優(yōu)化技術(shù)。亞線性優(yōu)化方法采用簡單的網(wǎng)格搜索（grid search）選擇檢查點(diǎn)，MegEngine在此基礎(chǔ)上增加遺傳算法，采用邊界移動(dòng)、塊合并、塊分裂等策略，實(shí)現(xiàn)更細(xì)粒度的優(yōu)化，進(jìn)一步降低了顯存占用。
如圖2所示，采用型號(hào)為2080Ti的GPU訓(xùn)練ResNet50，分別借助基準(zhǔn)、亞線性、亞線性+遺傳算法三種顯存優(yōu)化策略，對(duì)比了可使用的最大batch size。僅使用亞線性優(yōu)化，batch size從133增至211，是基準(zhǔn)的1.6x；而使用亞線性+遺傳算法聯(lián)合優(yōu)化，batch size進(jìn)一步增至262，較基準(zhǔn)提升2x。

圖2：三種顯存優(yōu)化方法優(yōu)化batch size的對(duì)比：ResNet50
通過選定同一模型、給定batch size，可以更好地觀察遺傳算法優(yōu)化顯存占用的情況。如圖3所示，隨著迭代次數(shù)的增加，遺傳算法逐漸收斂顯存占用，并在第5次迭代之后達(dá)到一個(gè)較穩(wěn)定的狀態(tài)。

圖3：遺傳算法收斂示意圖
此外，MegEngine亞線性優(yōu)化技術(shù)通過工程改良，不再局限于簡單的鏈狀結(jié)構(gòu)和同質(zhì)計(jì)算節(jié)點(diǎn), 可用于任意的計(jì)算圖，計(jì)算節(jié)點(diǎn)也可異質(zhì)，從而拓展了技術(shù)的適用場景；并可配合上述顯存優(yōu)化方法，進(jìn)一步降低模型的顯存占用。
實(shí)驗(yàn)
MegEngine基于亞線性顯存技術(shù)開展了相關(guān)實(shí)驗(yàn)，這里固定batch size=64，在ResNet18和ResNet50兩個(gè)模型上，考察模型訓(xùn)練時(shí)的顯存占用和計(jì)算時(shí)間。如圖4所示，相較于基準(zhǔn)實(shí)現(xiàn)，使用MegEngine亞線性顯存技術(shù)訓(xùn)練ResNet18時(shí)，顯存占用降低32%， 計(jì)算時(shí)間增加24%；在較大的ReNet50上，顯存占用降低40%，計(jì)算時(shí)間增加25%。同時(shí)經(jīng)過理論分析可知，模型越大，亞線性顯存優(yōu)化的效果越明顯，額外的計(jì)算時(shí)間則幾乎不變。

圖4：MegEngine亞線性優(yōu)化技術(shù)實(shí)驗(yàn)顯存/時(shí)間對(duì)比：ReNet18/ReNet50
在更大模型Bert上實(shí)驗(yàn)數(shù)據(jù)表明，借助MegEngine亞線性顯存技術(shù)，顯存占用最高降低75%，而計(jì)算時(shí)間僅增加23%，這與理論分析相一致。有興趣的同學(xué)可前往MegEngine ModeHub試手更多模型實(shí)驗(yàn)：https://megengine.org.cn/model-hub/。
使用
MegEngine官網(wǎng)提供了亞線性顯存優(yōu)化技術(shù)的使用文檔。當(dāng)你的GPU顯存有限，苦于無法訓(xùn)練較深、較大的神經(jīng)網(wǎng)絡(luò)模型，或者無法使用大batch進(jìn)一步提升深度神經(jīng)網(wǎng)絡(luò)的性能，抑或想要使batchwise算子更加穩(wěn)定，那么，MegEngine亞線性顯存優(yōu)化技術(shù)正是你需要的解決方案。
上手MegEngine亞線性優(yōu)化技術(shù)非常便捷，無需手動(dòng)設(shè)定梯度檢查點(diǎn)，通過幾個(gè)簡單的參數(shù)，輕松控制遺傳算法的搜索策略。具體使用時(shí)，在MegEngine靜態(tài)圖接口中調(diào)用SublinearMemoryConfig設(shè)置trace的參數(shù)sublinear_memory_config，即可打開亞線性顯存優(yōu)化：

from megengine.jit import trace, SublinearMemoryConfig

config = SublinearMemoryConfig()

@trace(symbolic=True, sublinear_memory_config=config)

def train_func(data, label, *, net, optimizer):

...

MegEngine在編譯計(jì)算圖和訓(xùn)練模型時(shí)，雖有少量的額外時(shí)間開銷，但會(huì)顯著緩解顯存不足問題。下面以ResNet50為例，說明MegEngine可有效突破顯存瓶頸，訓(xùn)練batch size從100最高增至200：

import os

from multiprocessing import Process

def train_resnet_demo(batch_size, enable_sublinear, genetic_nr_iter=0):

import megengine as mgeimport megengine.functional as Fimport megengine.hub as hubimport megengine.optimizer as optimfrom megengine.jit import trace, SublinearMemoryConfigimport numpy as npprint("Run with batch_size={}, enable_sublinear={}, genetic_nr_iter={}".format(batch_size, enable_sublinear, genetic_nr_iter))# 使用GPU運(yùn)行這個(gè)例子assert mge.is_cuda_available(), "Please run with GPU"try:# 從 megengine hub 中加載一個(gè) resnet50 模型。resnet = hub.load("megengine/models", "resnet50")optimizer = optim.SGD(resnet.parameters(), lr=0.1,)config = Noneif enable_sublinear:config = SublinearMemoryConfig(genetic_nr_iter=genetic_nr_iter)@trace(symbolic=True, sublinear_memory_config=config)def train_func(data, label, *, net, optimizer):pred = net(data)loss = F.cross_entropy_with_softmax(pred, label)optimizer.backward(loss)resnet.train()for i in range(10):batch_data = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)batch_label = np.random.randint(1000, size=(batch_size,)).astype(np.int32)optimizer.zero_grad()train_func(batch_data, batch_label, net=resnet, optimizer=optimizer)optimizer.step()except:print("Failed")returnprint("Sucess")

以下示例結(jié)果在2080Ti GPU運(yùn)行得到，顯存容量為 11 GB

不使用亞線性內(nèi)存優(yōu)化，允許的batch_size最大為 100 左右

p = Process(target=train_resnet_demo, args=(100, False))

p.start()

p.join()

報(bào)錯(cuò)顯存不足

p = Process(target=train_resnet_demo, args=(200, False))

p.start()

p.join()

使用亞線性內(nèi)存優(yōu)化，允許的batch_size最大為 200 左右

p = Process(target=train_resnet_demo, args=(200, True, 20))

p.start()

p.join()

展望
如上所述，MegEngine的亞線性顯存優(yōu)化技術(shù)通過額外做一次前向計(jì)算，即可達(dá)到O(sqrt(n))的空間復(fù)雜度。如果允許做更多次的前向計(jì)算，對(duì)整個(gè)網(wǎng)絡(luò)遞歸地調(diào)用亞線性顯存算法，有望在時(shí)間復(fù)雜度為O(n log n)的情況下，達(dá)到 O(log n)的空間復(fù)雜度。
更進(jìn)一步，MegEngine還將探索亞線性顯存優(yōu)化技術(shù)與數(shù)據(jù)并行/模型并行、混合精度訓(xùn)練的組合使用問題，以期獲得更佳的集成效果。最后，在RNN以及GNN、Transformer等其他類型網(wǎng)絡(luò)上的使用問題，也是MegEngine未來的一個(gè)探索方向。
歡迎訪問
? MegEngine GitHub（歡迎star）：https://github.com/MegEngine
? MegEngine官網(wǎng)：https://megengine.org.cn
? MegEngine ModelHub：https://megengine.org.cn/model-hub/
參考文獻(xiàn)

1. Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174.
2. He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).

總結(jié)

以上是生活随笔為你收集整理的MegEngine亚线性显存优化的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。