節省視訊記憶體新思路，在 PyTorch 裡使用 2 bit 啟用壓縮訓練神經網路

機器之心專欄

作者：鄭憐憫、陳鍵飛

隨著超大規模深度學習模型逐漸成為 AI 的趨勢，如何在有限的 GPU 記憶體下訓練這些模型成為了一個難題。

本文將介紹來自加州伯克利大學的 ActNN，一個基於 PyTorch 的啟用壓縮訓練框架。在同樣的記憶體限制下，ActNN 通過使用 2 bit 啟用壓縮，可以將 batch size 擴大 6-14 倍，將模型尺寸或者輸入圖片擴大 6-10 倍。ActNN 相關論文已被 ICML 2021 接收為 Long Talk，程式碼開源於 github。

AI 訓練撞上「記憶體牆」

從 AlexNet，ResNet 到 GPT-3，深度學習效能的突破都離不開模型規模的瘋狂增長。大模型有更好的效能已經成為業界的共識。過去幾年，不僅訓練一個最先進模型需要的算力在指數增長，訓練一個最先進模型需要的記憶體也在指數增長。如下圖所示，大型 Transformer 模型的參數量以每兩年翻 240 倍的速度指數增長。但是，單個 GPU 的記憶體卻只以每兩年翻 2 倍的速度在緩慢增長。另外，在訓練模型時，不光要儲存模型參數，還要儲存中間結果啟用值和優化器狀態，所需要的記憶體更多。如何在有限的 GPU 記憶體下訓練這些大規模模型成為了挑戰。

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

節省訓練記憶體的方法

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

ActNN：2 bit 啟用壓縮訓練

在訓練一個多層神經網路時，在前向傳播中，每一層的中間結果都要被存下來用於計算反向傳播的梯度。這些中間結果，又被叫做「啟用值」（activation)，實際上佔據了大部分的記憶體消耗，尤其是在 batch size 較大或者輸入圖片較大的時候。ActNN 的原理是就是壓縮這些啟用值來節省記憶體。如下圖所示，左圖表示的是普通的前向傳播和反向傳播，前向傳播時會存下所有層的 fp32 啟用值用於反向傳播，記憶體使用在計算 loss 的時候達到峰值。右圖表示的是 ActNN 的訓練方法：在前向傳播時，通過一個壓縮操作 Q 將啟用值壓縮後再儲存；反向傳播時，通過解壓縮操作 Q^-1 將啟用值解壓再計算梯度。

如果只是為了節省記憶體，這裡可以使用各種壓縮演算法，但是大部分現有的壓縮演算法並不能高效地運行在 GPU 上，會引入較大的開銷。ActNN 選擇了使用 2-bit 量化作為這裡的壓縮演算法。量化操作的代價較小，而且有一些好的數學性質允許我們使用有失真壓縮達到較大的壓縮比。

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

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

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

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

等號右邊的第一項是隨機梯度下降在 minibatch 取樣時產生的方差，等號右邊的第二項是有失真壓縮額外引入的方差。這條公式顯示地刻畫了有失真壓縮帶來的影響。注意到，當有損量化壓縮帶來的方差遠小於原來隨機梯度下降自帶的方差時，ActNN 引入的有失真壓縮就不會影響訓練的收斂性。更多關於公式的推導和視覺化參見文末的論文連結。論文對不同的運算元（conv2d，batch norm，linear等）都提供了詳細的分析。

由上述公式啟發，我們提出了一些新的量化技巧用於降低有失真壓縮引入的額外方差。我們引入了新的量化技巧 （ Per-group Quantization，Fine-Grained Mixed-Precision，Runtime Adaptation） 來利用梯度在不同樣本，不同緯度，不同層之間的異構特性。最後的壓縮演算法會分配更多的 bit 給更重要的啟用值。平均每個浮點數分配到 2 bit。

在具體實現壓縮演算法時，還有很多可以調節的參數。這裡產生了一個記憶體節省和訓練速度的取捨。一般來說，使用更復雜的壓縮演算法可以節省更多的記憶體，但是也會引入更多額外的開銷，使訓練速度變慢。為了給使用者較大的靈活性，ActNN 提供了 5 個優化等級 L1-L5 供使用者選擇。低的優化等級節省的記憶體較少，但是運行速度快。高的優化等級節省的記憶體多，但是運行也更慢。在最高優化等級 L5 下，ActNN 會結合一個簡單的記憶體交換策略，將壓縮後的啟用值移到 CPU 記憶體上，進一步節省記憶體。

實現

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

ActNN 為大部分常用的 PyTorch nn.Module 實現了使用量化壓縮的版本。使用者只需將模型裡的所有 PyTorch nn.Module 替換成 ActNN 對應的 Module （如把 nn.Conv2d 替換成 actnn.Conv2d)，即可節省記憶體，不需要更改其他程式碼。ActNN 同時也提供了一個 wrapper 實現一行程式碼自動替換。

實驗結果

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

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

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

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

我們還對更多的網路進行了測試。在同樣的記憶體限制下，ActNN 可以將 batch size 擴大 6-14 倍，將模型尺寸或者輸入圖片擴大 6-10 倍。詳細的實驗設定和結果參見文末的論文連結。

兩行程式碼即可在 PyTorch 中使用

import actnnmodel = actnn.QModule(model)

ActNN 提供了一個自動模型轉換封裝。只需在訓練腳本里插入兩行程式碼，即可將普通的 PyTorch 模型轉換為使用 ActNN 的模型。同時，ActNN 也提供了更高階的 API 支援定製化的使用場景。

更多的例子參見 github 連結。我們提供了在影象識別、影象分割、物體檢測，以及自監督學習等多個任務上使用 actnn 的完整例子和訓練記錄，歡迎試用！