作者 | 曠視研究院 編輯 | Linda 基于梯度檢查點的亞線性顯存優(yōu)化方法 [1] 由于較高的計算 / 顯存性價比受到關(guān)注。MegEngine 經(jīng)過工程擴(kuò)展和優(yōu)化，發(fā)展出一套行之有效的加強(qiáng)版亞線性顯存優(yōu)化技術(shù)，既可在計算存儲資源受限的條件下，輕松訓(xùn)練更深的模型，又可使用更大 batch size，進(jìn)一步提升模型性能，穩(wěn)定 batchwise 算子。使用 MegEngine 訓(xùn)練 ResNet18/ResNet50，顯存占用分別最高降低 23%/40%；在更大的 Bert 模型上，降幅更是高達(dá) 75%，而額外的計算開銷幾乎不變。該技術(shù)已在 MegEngine 開源，歡迎大家上手使用：https://github.com/MegEngine

深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練是一件復(fù)雜的事情，它體現(xiàn)為模型的時間復(fù)雜度和空間復(fù)雜度，分別對應(yīng)著計算和內(nèi)存；而訓(xùn)練時內(nèi)存占用問題是漂浮在深度學(xué)習(xí)社區(qū)上空的一塊烏云，如何撥云見日，最大降低神經(jīng)網(wǎng)絡(luò)訓(xùn)練的內(nèi)存占用，是一個繞不開的課題。

GPU 顯卡等硬件為深度學(xué)習(xí)提供了必需的算力，但硬件自身有限的存儲，限制了可訓(xùn)練模型的尺寸，尤其是大型深度網(wǎng)絡(luò)，由此誕生出一系列相關(guān)技術(shù)，比如亞線性顯存優(yōu)化、梯度累加、混合精度訓(xùn)練、分布式訓(xùn)練，進(jìn)行 GPU 顯存優(yōu)化。

其中，亞線性顯存優(yōu)化方法 [1] 由于較高的計算 / 顯存性價比備受關(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ù)的工程實現(xiàn)和實驗數(shù)據(jù)，從技術(shù)背景、原理、使用、展望等多個方面進(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)練時有效利用有限的~~ 計算~~ 存儲資源，尤其是減少 GPU 顯存占用，是評估深度學(xué)習(xí)框架性能的重要指標(biāo)。

在計算存儲資源一定的情況下，深度學(xué)習(xí)框架有幾種降低顯存占用的常用方法，其示例如下：

• 通過合適的梯度定義，讓算子的梯度計算不再依賴于前向計算作為輸入，從而 in-place 地完成算子的前向計算，比如 Sigmoid、Relu 等；? ?

• 在生命周期沒有重疊的算子之間共享顯存；? ?

• 通過額外的計算減少顯存占用，比如利用梯度檢查點重新計算中間結(jié)果的亞線性顯存優(yōu)化方法 [1]；? ?

• 通過額外的數(shù)據(jù)傳輸減少顯存占用，比如把暫時不用的數(shù)據(jù)從 GPU 交換到 CPU，需要時再從 CPU 交換回來。

上述顯存優(yōu)化技術(shù)在 MegEngine 中皆有不同程度的實現(xiàn)，這里重點討論基于梯度檢查點的亞線性顯存優(yōu)化技術(shù)。

原 ? 理

一個神經(jīng)網(wǎng)絡(luò)模型所占用的顯存空間大體分為兩個方面：?

1)模型本身的參數(shù)。

2)模型訓(xùn)練臨時占用的空間，包括參數(shù)的梯度、特征圖等。其中最大占比是 2)中以特征圖形式存在的中間結(jié)果，比如，從示例 [1] 可知，根據(jù)實現(xiàn)的不同，從 70% 到 90% 以上的顯存用來存儲特征圖。

這里的訓(xùn)練過程又可分為前向計算，反向計算和優(yōu)化三個方面，其中前向計算的中間結(jié)果最占顯存，還有反向計算的梯度。第 1)方面模型自身的參數(shù)內(nèi)存占用最小。?

MegEngine 加強(qiáng)版亞線性顯存優(yōu)化技術(shù)借鑒了 [1] 的方法，尤其適用于計算存儲資源受限的情況，比如一張英偉達(dá) 2080Ti，只有 11G 的顯存；而更貴的 Tesla V100，最大顯存也只有 32G。

? ?圖 1：亞線性顯存優(yōu)化原理，其中 (b) 保存了 Relu 結(jié)果，實際中 Relu 結(jié)果可用 in-place 計算

圖 1(a) 給出了卷積神經(jīng)網(wǎng)絡(luò)的基本單元，它由 Conv-BN-Relu 組成。可以看到，反向計算梯度的過程依賴于前向計算獲取的中間結(jié)果，一個網(wǎng)絡(luò)需要保存的中間結(jié)果與其大小成正比，即顯存復(fù)雜度為 O(n)。

本質(zhì)上，亞線性顯存優(yōu)化方法是以時間換空間，以計算換顯存，如圖 1(b) 所示，它的算法原理如下：

• 選取神經(jīng)網(wǎng)絡(luò)中 k 個檢查點，從而把網(wǎng)絡(luò)分成 k 個 block，需要注意的是，初始輸入也作為一個檢查點；前向計算過程中只保存檢查點處的中間結(jié)果；? ?

• 反向計算梯度的過程中，首先從相應(yīng)檢查點出發(fā)，重新計算單個 block 需要的中間結(jié)果，然后計算 block 內(nèi)部各個 block 的梯度；不同 block 的中間結(jié)果計算共享顯存。這種方法有著明顯的優(yōu)點，即大幅降低了模型的空間復(fù)雜度，同時缺點是增加了額外的計算：? ?

• 顯存占用從 O(n) 變成 O(n/k)+ O(k)，O(n/k) 代表計算單個節(jié)點需要的顯存，O(k) 代表 k 個檢查點需要的顯存， 取 k=sqrt(n)，O(n/k)+ O(k)~O(sqrt(n))，可以看到顯存占用從線性變成了亞線性；? ?

• 因為在反向梯度的計算過程中需要從檢查點恢復(fù)中間結(jié)果，整體需要額外執(zhí)行一次前向計算。

工 ? 程

在 [1] 的基礎(chǔ)上，MegEngine 結(jié)合自身實踐，做了工程擴(kuò)展和優(yōu)化，把亞線性顯存優(yōu)化方法擴(kuò)展至任意的計算圖，并結(jié)合其它常見的顯存優(yōu)化方法，發(fā)展出一套行之有效的加強(qiáng)版亞線性顯存優(yōu)化技術(shù)。

亞線性優(yōu)化方法采用簡單的網(wǎng)格搜索(grid search)選擇檢查點，MegEngine 在此基礎(chǔ)上增加遺傳算法，采用邊界移動、塊合并、塊分裂等策略，實現(xiàn)更細(xì)粒度的優(yōu)化，進(jìn)一步降低了顯存占用。

如圖 2 所示，采用型號為 2080Ti 的 GPU 訓(xùn)練 ResNet50，分別借助基準(zhǔn)、亞線性、亞線性 + 遺傳算法三種顯存優(yōu)化策略，對比了可使用的最大 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 的對比：ResNet50

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

圖 3：遺傳算法收斂示意圖

此外，MegEngine 亞線性優(yōu)化技術(shù)通過工程改良，不再局限于簡單的鏈狀結(jié)構(gòu)和同質(zhì)計算節(jié)點, 可用于任意的計算圖，計算節(jié)點也可異質(zhì)，從而拓展了技術(shù)的適用場景；并可配合上述顯存優(yōu)化方法，進(jìn)一步降低模型的顯存占用。

實 ?驗

MegEngine 基于亞線性顯存技術(shù)開展了相關(guān)實驗，這里固定 batch size=64，在 ResNet18 和 ResNet50 兩個模型上，考察模型訓(xùn)練時的顯存占用和計算時間。? ?

如圖 4 所示，相較于基準(zhǔn)實現(xiàn)，使用 MegEngine 亞線性顯存技術(shù)訓(xùn)練 ResNet18 時，顯存占用降低 32%， 計算時間增加 24%；在較大的 ReNet50 上，顯存占用降低 40%，計算時間增加 25%。同時經(jīng)過理論分析可知，模型越大，亞線性顯存優(yōu)化的效果越明顯，額外的計算時間則幾乎不變。

圖 4：MegEngine 亞線性優(yōu)化技術(shù)實驗顯存 / 時間對比：ReNet18/ReNet50

在更大模型 Bert 上實驗數(shù)據(jù)表明，借助 MegEngine 亞線性顯存技術(shù)，顯存占用最高降低 75%，而計算時間僅增加 23%，這與理論分析相一致。有興趣的同學(xué)可前往 MegEngine ModeHub 試手更多模型實驗：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ù)非常便捷，無需手動設(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 在編譯計算圖和訓(xùn)練模型時，雖有少量的額外時間開銷，但會顯著緩解顯存不足問題。下面以 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 mge
import megengine.functional as F
import megengine.hub as hub
import megengine.optimizer as optim
from megengine.jit import trace, SublinearMemoryConfig
import numpy as np

print(
"Run with batch_size={}, enable_sublinear={}, genetic_nr_iter={}".format(
batch_size, enable_sublinear, genetic_nr_iter
)
)
# 使用 GPU 運(yùn)行這個例子
assert mge.is_cuda_available(), "Please run with GPU"
try:
# 我們從 megengine hub 中加載一個 resnet50 模型。
resnet = hub.load("megengine/models", "resnet50")

optimizer = optim.SGD(resnet.parameters(), lr=0.1,)

config = None
if 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")
return

print("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()
# 報錯顯存不足
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ù)通過額外做一次前向計算，即可達(dá)到 O(sqrt(n)) 的空間復(fù)雜度。如果允許做更多次的前向計算，對整個網(wǎng)絡(luò)遞歸地調(diào)用亞線性顯存算法，有望在時間復(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 未來的一個探索方向。

了解更多信息請查詢：

• MegEngine GitHub：https://github.com/MegEngine

• MegEngine 官網(wǎng)：https://megengine.org.cn

• MegEngine ModelHub：https://megengine.org.cn/model-hub/

參考文獻(xiàn)

Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174. ?

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é)

以上是生活随笔為你收集整理的深度学习 占用gpu内存 使用率为0_深度解析MegEngine亚线性显存优化技术的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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