文 | jxyxiangyu
編 | 小軼

“小夕，小夕！又出來了個 SOTA 模型！趕緊 follow ！”

小夕看了看新模型的參數量， 然后看了看實驗室服務器的幾張小破卡。

小夕，陷入了沉默。

自從人們發現越大的模型性能越好后，神經網絡模型的參數量就在越來越大的道路上一去不復返了。從XX-large到GPT3，再到5300億參數的Megatron Turing-NLG，深度學習越來越像是只有財大氣粗的大公司才能玩得起的玩具。如果，我們想要在實驗室“簡陋”的環境下，嘗試更大的模型，有什么行之有效的方法呢？

最近，Facebook 推出了支持 pytorch 的 8 位優化器，在減小內存占用的同時，竟然還能保持和32位優化器相當的準確性。不得不說 facebook yyds。那么，下面就讓我們一起來看看具體是怎么做的吧。

論文題目:
8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION

論文鏈接:
https://arxiv-download.xixiaoyao.cn/pdf/2110.02861.pdf

開源鏈接:
https://github.com/facebookresearch/bitsandbytes

量化

在介紹論文作者的解決方法之前，先補充一點關于量化的基本概念。通常意義上來說，量化是指將信號的連續取值近似為有限多個離散值的過程。具體到計算機系統，指的是將浮點數值映射到低bit數值的操作^[1]。

一般來說，我們可以通過以下手段應用量化

量化模型參數來壓縮模型；

量化模型某些層的激活值來減少內存占用*；

(注：參數和梯度也會占用一定的內存空間，但相對于激活值而言，占用比例不大，一般來說，量化參數和梯度帶來的內存收益沒有量化激活值的大)

▲量化示意圖

上圖是 Song Han 在 ICLR'2016 上提出的量化方法。將模型參數分別聚類到幾個質心，并將參數量化到對應的質心，在更新參數時，是將同一質心對應的梯度累加用于更新該質心對應的參數

可以看到，量化通過將參數（浮點值）映射到二值、三值或線性量化到一個區間（一般是低比特數值）的方式，減小了模型大小，在某些硬件上面，低比特數值運算速度高于浮點數值，一定程度上可以加速模型的訓練和預測；除此之外，模型在訓練和預測的時候，模型參數本身只占用了內存的一小部分，大部分存儲了模型的激活值，如果將量化應用到激活值上，一定程度也減小了內存占用，這樣，我們就可以嘗試更大的模型和設置更大的mini-batch了。

當然，量化這么好，也不是沒有缺點的，量化后的模型或多或少會引入精度損失；并且目前學術界多采用 pytorch 框架，好巧不巧的是 pytorch 框架對量化的支持沒有 tensorflow 好，這總不能為了體驗大模型的快感再轉到 tensorflow 上面去吧，想想 tensorflow 混亂的 api 就頭疼(╯°Д°)╯ ┻━┻

最近，Facebook 推出了支持 pytorch 的 8 位優化器，在減小內存占用的同時，竟然還能保持和32位優化器相當的準確性。

狀態優化器

再簡單介紹下帶有狀態的優化器（stateful optimizer）。和普通的隨機梯度下降（SGD）相比，為了加速優化而提出的帶有梯度統計信息的優化器，就是狀態優化器。常見的例如帶動量的 SGD 和 Adam 。計算公式如下：其中， 和 是平滑因子， 是非常小的常量， 是學習率。

作者認為，狀態優化器會維護歷史梯度數據，一定程度上占用了內存。通過量化這些梯度，可以有效地降低內存占用。

非線性量化

前述已經介紹了量化就是將信號的連續取值近似為有限多個離散值的過程，在降低模型參數量的同時，也會帶來一定的精度損失，為減小精度損失，多采用非線性的量化方式，大致可以歸納為三個步驟：

對于輸入張量，計算歸一化因子；

將張量通過歸一化后，找到在量化空間中距離最近的值；

將量化后的張量的每個元素的索引存儲下來

那么，我們就可以遍歷存儲的索引并通過下式得到反量化張量：，其中，是反量化映射

為了使不同元素值量級一致，一般會將張量歸一化到的區間范圍，這時，取的是輸入張量中絕對值的最大值，即，然后通過二分查找的方式找到量化空間中距離該值最近的量化值

動態樹量化

上一節看到，非線性量化在歸一化時會嚴重依賴輸入張量中的最值，像某些特別大或特別小的異常值，對量化會產生較大的精度影響。動態樹量化（dynamic tree quantization）就是一種以較低的量化精度損失處理這種情況的方法。

與浮點數的存儲方式類似，動態樹量化以這類方式解釋存儲在內存中的數值，以此實現量化，具體由四部分組成：

首位是符號位

符號位后連續的0的數量表示指數大小

再之后的第一位是指示位，如果指示位為1表示后續剩余的位為線性量化區域

線性量化區域

其中，指示位是可以動態移動的。通過移動指示位，可以表示指數為或者精度為的數值，表示范圍為

8位優化器

有了前面的知識鋪墊，下面就可以詳細地說明作者提出的8位優化器了。該8位優化器由三部分構成：

逐塊量化(block-wise quantization)

動態量化(dynamic quantization)

穩定的詞嵌入層(stable embedding layer)

應用上述組件，將8位優化器的狀態反量化為32位并更新狀態和參數，然后將這些狀態量化回8位進行存儲。由于是在寄存器中進行8位和32位的轉換，一定程度上可以減小內存占用并加速訓練。

逐塊量化

常見的量化需要將原始的張量在張量級別歸一化，這樣可能會引入核之間的多次信息通信和同步，造成額外的時間開銷，而逐塊量化則是將張量分成多個小塊并在塊級別歸一化，減小了核之間的通信開銷，除此之外，還可以將張量元素中的異常值的影響限制在單個塊中。假設為有個元素的張量，分成每個大小為的塊，那么，可以分成個塊，對每個塊分別做歸一化，歸一化因子為，每個塊分別通過下式進行量化操作：其中，為塊索引，為塊中元素的索引

動態量化

8位優化器的動態量化部分是對前面提到的動態樹量化的擴展，對于像的第二個狀態這種嚴格為正的數值，符號位就顯得有些多余，而在語言模型的訓練過程中，作者發現的變化范圍在3~5個數量級，小于動態樹量化的7個數量級，因此，可以用固定的位將只會用到的位劃分開，進一步減小內存占用。對于其他帶符號的狀態張量，則繼續使用動態樹量化。

穩定的詞嵌入層

為了確保nlp任務中模型的穩定訓練，作者還添加了穩定的詞嵌入層。作者使用Xavier uniform對詞嵌入層進行初始化，并且在與位置向量合并前進行層歸一化操作，這樣可以使參數在初始化和訓練期間保持1左右的方差。詞嵌入層的優化器狀態用32位存儲，權重和梯度用16位存儲。

8位優化器 vs 32位優化器

作者在多個任務（包括機器翻譯、大規模語言模型的預訓練以及微調、圖像分類和圖像預訓練以及微調）上比較了8位優化器和32位優化器的性能，比較的優化器包括、或，實驗中，除了將32位優化器替換為8位優化器外，沒有改動超參和權重、梯度以及激活值的精度。

除了GLUE任務之外，其余的NLP任務均使用了作者提出的穩定詞嵌入層。為確保實驗結果的可信度，還在不同隨機數種子下多次實驗，選取了實驗結果的中位數作為最終性能。實驗結果如下所示：

可以看到，8位優化器在多個任務上均達到甚至是超過了32位優化器的性能，與此同時，還能大幅減小內存開銷并加速訓練。此外，作者還列出了在同等顯存大小的條件下，使用8位優化器和32位優化器可以支持訓練的模型。可以說，非常貼心了ヾ(●゜ⅴ゜)ﾉ

消融研究

作者基于語料庫訓練了多個模型，用于研究8位優化器中各個組件的影響。實驗結果如下：其中，32位優化器(baseline)采用的是線性量化。

為測試優化器的穩定性，對于小規模的模型，作者分別訓練了不同的超參數下的模型，超參數為 {1e-8, 1e-7, 1e-6}, {0.90, 0.87, 0.93}, {0.999, 0.99, 0.98}以及學習率方面的改動，而對于超過1B的大規模模型，則是在相同超參下采用不同的隨機數種子多次運行。所有的結果均是選擇可以成功訓練完（沒有因梯度爆炸或彌散而無法訓練）的模型性能的中值。

可以看出，逐塊量化、動態量化和穩定的詞嵌入層對結果都有正向影響。

此外，作者還對比了32位優化器和8位優化器對超參的敏感程度，比較了32位和8位優化器在、、和的變化下的走勢

可以看到，8位優化器和32位相比，困惑度走勢基本一致，表明對超參不敏感，在將32位優化器替換為8位優化器后，超參不需要進一步的調整

局限性

從實驗結果可以看出，8位優化器完全可以作為32位優化器的替代品。當然，8位優化器也存在一些局限性：

8位優化器需要穩定的詞嵌入層來達到32位優化器的性能；

8位優化器減小內存的大小與模型參數量成正比，對于像cnn這種激活值比參數占內存多得多的模型，8位優化器并沒有特別明顯的內存減小，反而更適合transformer這種架構的大規模模型

總結

不得不說，Facebook的8位優化器簡直是我等“窮困”煉丹黨的福音。現在，8位優化器已經開源，開源地址已經在文章開頭提到。目前，8位優化器已經支持Adam, AdamW, RMSProp, LARS, LAMB優化器。使用時，需要安裝并導入包bitsandbytes-cudaXXX，其中，XXX是本地環境的cuda工具包版本號，注釋掉原有的優化器，調用8位優化器就可以了。

import?bitsandbytes?as?bnb#?adam?=?torch.optim.Adam(model.parameters(),?lr=0.001,?betas=(0.9,?0.995))?#?comment?out?old?optimizer adam?=?bnb.optim.Adam8bit(model.parameters(),?lr=0.001,?betas=(0.9,?0.995))?#?add?bnb?optimizer adam?=?bnb.optim.Adam(model.parameters(),?lr=0.001,?betas=(0.9,?0.995),?optim_bits=8)?#?equivalenttorch.nn.Embedding(...)?->??bnb.nn.StableEmbedding(...)?#?recommended?for?NLP?models

據官網描述，僅僅需要改動兩行代碼，就可以節省75%的內存！小伙伴們，還不想抓緊時間上車體驗一下嘛？

▲沒時間解釋了，快上車

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

?

[1] 商湯科技SenseTime, 模型量化了解一下？(https://zhuanlan.zhihu.com/p/132561405)

[2] Song, H. , ?H. Mao , and ?W. J. Dally . "Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding." ICLR 2016. (https://arxiv-download.xixiaoyao.cn/pdf/1510.00149.pdf)

總結

以上是生活随笔為你收集整理的一训练就显存爆炸？Facebook 推出 8 比特优化器，两行代码拯救你的显存！的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。