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

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

本篇文章轉(zhuǎn)自于“機器之心”

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

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

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

AI 訓(xùn)練撞上「內(nèi)存墻」

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

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

節(jié)省訓(xùn)練內(nèi)存的方法

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

ActNN：2 bit 激活壓縮訓(xùn)練

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

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

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

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

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

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

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

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

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

實現(xiàn)

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

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

實驗結(jié)果

因為 ActNN 進行的是有損壓縮，所以最重要的一點是先驗證 ActNN 是否會影響模型的精度。下圖是使用 ActNN 在 ImageNet 上訓(xùn)練 ResNet-50 的結(jié)果。FP 代表普通的 fp32 訓(xùn)練， BLPA 是來自 NeurIPS 2019 的一個相關(guān)工作。可以看到，在 ActNN 的 2-bit 壓縮模式下，模型幾乎沒有損失精度。在更極限的 1.25 bit 的情況下，ActNN 也能收斂，只不過會損失一些精度。而之前的工作 BLPA 在小于 4 bit 的情況就下無法收斂。

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

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

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

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

兩行代碼即可在 PyTorch 中使用

import actnn
model = actnn.QModel(model)

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

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