節省顯存新思路，在PyTorch里使用2 bit激活壓縮訓練神經(jīng)網(wǎng)絡(luò )

隨著(zhù)超大規模深度學(xué)習模型逐漸成為 AI 的趨勢，如何在有限的 GPU 內存下訓練這些模型成為了一個(gè)難題。

本篇文章轉自于“機器之心”

本文將介紹來(lái)自加州伯克利大學(xué)的 ActNN，一個(gè)基于 PyTorch 的激活壓縮訓練框架。在同樣的內存限制下，ActNN 通過(guò)使用 2 bit 激活壓縮，可以將 batch size 擴大 6-14 倍，將模型尺寸或者輸入圖片擴大 6-10 倍。ActNN 相關(guān)論文已被 ICML 2021 接收為 Long Talk，代碼開(kāi)源于 github。

論文 https://arxiv.org/abs/2104.14129

代碼 https://github.com/ucbrise/actnn

AI 訓練撞上「內存墻」

從 AlexNet，ResNet 到 GPT-3，深度學(xué)習性能的突破都離不開(kāi)模型規模的瘋狂增長(cháng)。大模型有更好的性能已經(jīng)成為業(yè)界的共識。過(guò)去幾年，不僅訓練一個(gè)最先進(jìn)模型需要的算力在指數增長(cháng)，訓練一個(gè)最先進(jìn)模型需要的內存也在指數增長(cháng)。如下圖所示，大型 Transformer 模型的參數量以每?jì)赡攴?240 倍的速度指數增長(cháng)。但是，單個(gè) GPU 的內存卻只以每?jì)赡攴?2 倍的速度在緩慢增長(cháng)。另外，在訓練模型時(shí)，不光要存儲模型參數，還要存儲中間結果激活值和優(yōu)化器狀態(tài)，所需要的內存更多。如何在有限的 GPU 內存下訓練這些大規模模型成為了挑戰。 

source：Gholami A, Yao Z, Kim S, Mahoney MW, Keutzer K. AI and Memory Wall. RiseLab Medium Blog Post, University of California Berkeley

節省訓練內存的方法

目前，節省訓練內存的方法主要有三類(lèi)：1. 重計算（Gradient checkpointing/Rematerialization)  2. 使用 CPU 內存進(jìn)行交換 (swapping)  和 3. 使用分布式訓練將 Tensor 分散存儲在多個(gè) GPU 上。這三類(lèi)方法互相不沖突，可以結合使用。大部分機器學(xué)習框架對這些方法都提供了一些支持，也有不少相關(guān)的論文。但是，想要高效、自動(dòng)化地實(shí)現這些策略并不容易。與已有方法不同，我們提出了 ActNN，一個(gè)新的基于壓縮的內存節省框架。在提供理論證明的同時(shí)，我們基于 PyTorch 提供了一個(gè)高效易用的實(shí)現。Table.1 比較了 ActNN 和已有的一些內存節省系統。ActNN 支持 PyTorch 的動(dòng)態(tài)圖執行模式，并且不需要預先進(jìn)行復雜的策略搜索。ActNN 作為一個(gè)獨立的 Python 庫，使用時(shí) import 即可，不需要修改或重新編譯 PyTorch。與已有的工作相比，ActNN 靈活且易于使用。同時(shí)，ActNN 在理論上也可以和已有的技術(shù)相互疊加。

ActNN：2 bit 激活壓縮訓練

在訓練一個(gè)多層神經(jīng)網(wǎng)絡(luò )時(shí)，在前向傳播中，每一層的中間結果都要被存下來(lái)用于計算反向傳播的梯度。這些中間結果，又被叫做「激活值」（activation)，實(shí)際上占據了大部分的內存消耗，尤其是在 batch size 較大或者輸入圖片較大的時(shí)候。ActNN 的原理是就是壓縮這些激活值來(lái)節省內存。如下圖所示，左圖表示的是普通的前向傳播和反向傳播，前向傳播時(shí)會(huì )存下所有層的 fp32 激活值用于反向傳播，內存使用在計算 loss 的時(shí)候達到峰值。右圖表示的是 ActNN 的訓練方法：在前向傳播時(shí)，通過(guò)一個(gè)壓縮操作 Q 將激活值壓縮后再存儲；反向傳播時(shí)，通過(guò)解壓縮操作 Q^-1 將激活值解壓再計算梯度。

如果只是為了節省內存，這里可以使用各種壓縮算法，但是大部分現有的壓縮算法并不能高效地運行在 GPU 上，會(huì )引入較大的開(kāi)銷(xiāo)。ActNN 選擇了使用 2-bit 量化作為這里的壓縮算法。量化操作的代價(jià)較小，而且有一些好的數學(xué)性質(zhì)允許我們使用有損壓縮達到較大的壓縮比。

把 fp32 浮點(diǎn)數量化為 2-bit 整數是一個(gè)有損壓縮，會(huì )引入一些誤差。論文從理論上分析了量化引入的誤差是如何影響訓練的收斂性的。

第一，存在一個(gè)隨機化的量化策略，使得使用有損量化壓縮后，估計出的有損梯度是原梯度的一個(gè)無(wú)偏估計。

在這一條件下，我們套用已有的隨機梯度下降收斂性定理，得出最后收斂時(shí)的誤差會(huì )被梯度的方差所限制。

第二，我們推導出了使用量化壓縮之后，隨機梯度下降計算出的梯度的方差。

等號右邊的第一項是隨機梯度下降在 minibatch 采樣時(shí)產(chǎn)生的方差，等號右邊的第二項是有損壓縮額外引入的方差。這條公式顯示地刻畫(huà)了有損壓縮帶來(lái)的影響。注意到，當有損量化壓縮帶來(lái)的方差遠小于原來(lái)隨機梯度下降自帶的方差時(shí)，ActNN 引入的有損壓縮就不會(huì )影響訓練的收斂性。更多關(guān)于公式的推導和可視化參見(jiàn)文末的論文鏈接。論文對不同的算子（conv2d，batch norm，linear等）都提供了詳細的分析。

由上述公式啟發(fā)，我們提出了一些新的量化技巧用于降低有損壓縮引入的額外方差。我們引入了新的量化技巧 （ Per-group Quantization，Fine-Grained Mixed-Precision，Runtime Adaptation） 來(lái)利用梯度在不同樣本，不同緯度，不同層之間的異構特性。最后的壓縮算法會(huì )分配更多的 bit 給更重要的激活值。平均每個(gè)浮點(diǎn)數分配到 2 bit。

在具體實(shí)現壓縮算法時(shí)，還有很多可以調節的參數。這里產(chǎn)生了一個(gè)內存節省和訓練速度的取舍。一般來(lái)說(shuō)，使用更復雜的壓縮算法可以節省更多的內存，但是也會(huì )引入更多額外的開(kāi)銷(xiāo)，使訓練速度變慢。為了給用戶(hù)較大的靈活性，ActNN 提供了 5 個(gè)優(yōu)化等級 L1-L5 供用戶(hù)選擇。低的優(yōu)化等級節省的內存較少，但是運行速度快。高的優(yōu)化等級節省的內存多，但是運行也更慢。在最高優(yōu)化等級 L5 下，ActNN 會(huì )結合一個(gè)簡(jiǎn)單的內存交換策略，將壓縮后的激活值移到 CPU 內存上，進(jìn)一步節省內存。

實(shí)現

要在 PyTorch 實(shí)現 ActNN 算法非常簡(jiǎn)單。對于一個(gè) PyTorch nn Module，我們只需要在其 forward 函數里加入量化壓縮，在其 backward 函數里加入解壓縮操作。所有的計算還是在 fp32 下進(jìn)行，與原來(lái)一樣，偽代碼如下圖所示。

ActNN 為大部分常用的 PyTorch nn.Module 實(shí)現了使用量化壓縮的版本。用戶(hù)只需將模型里的所有 PyTorch nn.Module 替換成 ActNN 對應的 Module （如把 nn.Conv2d 替換成 actnn.Conv2d)，即可節省內存，不需要更改其他代碼。ActNN 同時(shí)也提供了一個(gè) wrapper 實(shí)現一行代碼自動(dòng)替換。

實(shí)驗結果

因為 ActNN 進(jìn)行的是有損壓縮，所以最重要的一點(diǎn)是先驗證 ActNN 是否會(huì )影響模型的精度。下圖是使用 ActNN 在 ImageNet 上訓練 ResNet-50 的結果。FP 代表普通的 fp32 訓練， BLPA 是來(lái)自 NeurIPS 2019 的一個(gè)相關(guān)工作?？梢钥吹?，在 ActNN 的 2-bit 壓縮模式下，模型幾乎沒(méi)有損失精度。在更極限的 1.25 bit 的情況下，ActNN 也能收斂，只不過(guò)會(huì )損失一些精度。而之前的工作 BLPA 在小于 4 bit 的情況就下無(wú)法收斂。

我們還在圖像分割，物體檢測，以及自監督學(xué)習等多個(gè)任務(wù)上進(jìn)行了實(shí)驗。ActNN 都能在 2-bit 壓縮模式下達到和普通 fp32 幾乎一樣的結果。在部分任務(wù)上，因為 ActNN 可以使用更大的 batch size，甚至可以取得更好的測試結果。詳細的實(shí)驗結果和訓練記錄參見(jiàn)文末的論文與 github 鏈接。

之后，我們對比了 ActNN 與普通 fp32 訓練的實(shí)際內存使用情況。如下表所示，ActNN 可以將激活值占用的內存壓縮 12 倍，將訓練使用的總內存壓縮 4 - 7 倍。這一實(shí)際內存壓縮效果符合理論推導。為什么激活值壓縮倍率是 12 而不是 32 bit / 2 bit = 16？主要是因為 ActNN 不能使用 inplace 的 ReLU，以及需要存儲少量額外的 min 和 scale 用于解壓縮。

最后，我們測試了 ActNN 的訓練速度。因為 ActNN 在訓練過(guò)程中進(jìn)行了壓縮，這些壓縮在節省內存的同時(shí)也會(huì )引入額外的計算開(kāi)銷(xiāo)。一般來(lái)說(shuō)，省得內存越多，進(jìn)入的額外開(kāi)銷(xiāo)就越多，訓練也就越慢。我們在 NVIDIA T4 (16 GB 內存) 上對比了 ActNN 和已有內存節省系統的訓練速度。如下圖所示，DTR （ICLR 2020），BLPA （NeurIPS 2019）和 swap 分別是基于重計算，壓縮和內存交換的三種方法，紅叉代表 Out-of-memory。y 軸是訓練吞吐量 (images per second)，越高越好。綠色的曲線(xiàn)是綜合 ActNN 在不同優(yōu)化等級下的最優(yōu)結果?？梢钥吹?，ActNN 不僅能開(kāi)到最大的 batch size（即最省內存），同時(shí)在所有 batch size 下都比 baseline 的訓練速度更快。

我們還對更多的網(wǎng)絡(luò )進(jìn)行了測試。在同樣的內存限制下，ActNN 可以將 batch size 擴大 6-14 倍，將模型尺寸或者輸入圖片擴大 6-10 倍。詳細的實(shí)驗設置和結果參見(jiàn)文末的論文鏈接。

兩行代碼即可在 PyTorch 中使用

import actnn
model = actnn.QModel(model)

ActNN 提供了一個(gè)自動(dòng)模型轉換封裝。只需在訓練腳本里插入兩行代碼，即可將普通的 PyTorch  模型轉換為使用 ActNN 的模型。同時(shí)，ActNN 也提供了更高級的 API 支持定制化的使用場(chǎng)景。

更多的例子參見(jiàn) github 鏈接。我們提供了在圖像識別、圖像分割、物體檢測，以及自監督學(xué)習等多個(gè)任務(wù)上使用 actnn 的完整例子和訓練記錄，歡迎試用！