我們知道，將激活、權重和梯度量化為 4-bit 對于加速神經網絡訓練非常有價值。但現有的 4-bit 訓練方法需要自定義數字格式，而當代硬件不支持這些格式。在本文中，清華朱軍等人提出了一種使用 INT4 算法實現所有矩陣乘法的 Transformer 訓練方法。

模型訓練得快不快，這與激活值、權重、梯度等因素的要求緊密相關。

神經網絡訓練需要一定計算量，使用低精度算法（全量化訓練或 FQT 訓練）有望提升計算和內存的效率。FQT 在原始的全精度計算圖中增加了量化器和去量化器，并將昂貴的浮點運算替換為廉價的低精度浮點運算。

對 FQT 的研究旨在降低訓練數值精度，同時降低收斂速度和精度的犧牲。所需數值精度從 FP16 降到 FP8、INT32+INT8 和 INT8+INT5。FP8 訓練通過有 Transformer 引擎的 Nvidia H100 GPU 完成，這使大規模 Transformer 訓練實現了驚人的加速。

最近訓練數值精度已被壓低到 4 位（ 4 bits）。Sun 等人成功訓練了幾個具有 INT4 激活 / 權重和 FP4 梯度的當代網絡；Chmiel 等人提出自定義的 4 位對數數字格式，進一步提高了精度。然而，這些 4 位訓練方法不能直接用于加速，因為它們需要自定義數字格式，這在當代硬件上是不支持的。

在 4 位這樣極低的水平上訓練存在著巨大的優化挑戰，首先前向傳播的不可微分量化器會使損失函數圖不平整，其中基于梯度的優化器很容易卡在局部最優。其次梯度在低精度下只能近似計算，這種不精確的梯度會減慢訓練過程，甚至導致訓練不穩定或發散的情況出現。

本文為流行的神經網絡 Transformer 提出了新的 INT4 訓練算法。訓練 Transformer 所用的成本巨大的線性運算都可以寫成矩陣乘法（MM）的形式。MM 形式使研究人員能夠設計更加靈活的量化器。這種量化器通過 Transformer 中的特定的激活、權重和梯度結構，更好地近似了 FP32 矩陣乘法。本文中的量化器還利用了隨機數值線性代數領域的新進展。

圖片

論文地址：https://arxiv.org/pdf/2306.11987.pdf

研究表明，對前向傳播而言，精度下降的主要原因是激活中的異常值。為了抑制該異常值，研究提出了 Hadamard 量化器，用它對變換后的激活矩陣進行量化。該變換是一個分塊對角的 Hadamard 矩陣，它將異常值所攜帶的信息擴散到異常值附近的矩陣項上，從而縮小了異常值的數值范圍。

對反向傳播而言，研究利用了激活梯度的結構稀疏性。研究表明，一些 token 的梯度非常大，但同時，其余大多數的 token 梯度又非常小，甚至比較大梯度的量化殘差更小。因此，與其計算這些小梯度，不如將計算資源用于計算較大梯度的殘差。

結合前向和反向傳播的量化技術，本文提出一種算法，即對 Transformer 中的所有線性運算使用 INT4 MMs。研究評估了在各種任務上訓練 Transformer 的算法，包括自然語言理解、問答、機器翻譯和圖像分類。與現有的 4 位訓練工作相比，研究所提出的算法實現了相媲美或更高的精度。此外，該算法與當代硬件 (如 GPU) 是兼容的，因為它不需要自定義數字格式 (如 FP4 或對數格式)。并且研究提出的原型量化 + INT4 MM 算子比 FP16 MM 基線快了 2.2 倍，將訓練速度提高了 35.1%。

前向傳播

在訓練過程中，研究者利用 INT4 算法加速所有的線性算子，并將所有計算強度較低的非線性算子設置為 FP16 格式。Transformer 中的所有線性算子都可以寫成矩陣乘法形式。為了便于演示，他們考慮了如下簡單的矩陣乘法加速。

圖片

這種矩陣乘法的最主要用例是全連接層。

學得的步長量化

加速訓練必須使用整數運算來計算前向傳播。因此，研究者利用了學得的步長量化器（LSQ）。作為一種靜態量化方法，LSQ 的量化規模不依賴于輸入，因此比動態量化方法成本更低。相較之下，動態量化方法需要在每次迭代時動態地計算量化規模。

給定一個 FP 矩陣 X，LSQ 通過如下公式 (2) 將 X 量化為整數。

圖片

激活異常值

簡單地將 LSQ 應用到具有 4-bit 激活 / 權重的 FQT（fully quantized training，全量化訓練）中，會由于激活異常值而導致準確度下降。如下圖 1 (a) 所示，激活的有一些異常值項，其數量級比其他項大得多。

在這種情況下，步長 s_X 在量化粒度和可表示數值范圍之間進行權衡。如果 s_X 很大，則可以很好地表示異常值，同時代價是以粗略的方式表示其他大多數項。如果 s_X 很小，則必須截斷 [?Q_Ns_X, Q_Ps_X] 范圍之外的項。

Hadamard 量化

研究者提出使用 Hadamard 量化器（HQ）來解決異常值問題，它的主要思路是在另一個異常值較少的線性空間中量化矩陣。

激活矩陣中的異常值可以形成特征級結構。這些異常值通常集中在幾個維度上，也就是 X 中只有幾列顯著大于其他列。作為一種線性變換，Hadamard 變換可以將異常值分攤到其他項中。具體地，Hadamard 變換 H_k 是一個 2^k × 2^k 矩陣。

為了抑制異常值，研究者對 X 和 W 的變換版本進行量化。

通過結合量化后的矩陣，研究者得到如下。

其中逆變換彼此之間相互抵消，并且 MM 可以實現如下。

圖片

反向傳播

研究者使用 INT4 運算來加速線性層的反向傳播。公式 (3) 中定義的線性算子 HQ-MM 具有四個輸入，分別是激活 X、權重 W 以及步長 s_X 和 s_W。給定關于損失函數 L 的輸出梯度?_YL，他們需要計算這四個輸入的梯度。

梯度的結構稀疏性

研究者注意到，訓練過程中梯度矩陣?_Y 往往非常稀疏。稀疏性結構是這樣的：?_Y 的少數行（即 tokens）具有較大的項，而大多數其他行接近全零向量。他們在下圖 2 中繪制了所有行的 per-row 范數∥(?_Y)_i:∥的直方圖。

圖片

Bit 拆分和平均分數采樣

研究者討論了如何設計梯度量化器，從而利用結構稀疏性在反向傳播期間準確計算 MM。高級的思路是，很多行的梯度非常的小，因而對參數梯度的影響也很小，但卻浪費了大量計算。此外，大行無法用 INT4 準確地表示。

為利用這種稀疏性，研究者提出 bit 拆分，將每個 token 的梯度拆分為更高的 4bits 和更低的 4bits。然后再通過平均分數采樣選擇信息量最大的梯度，這是 RandNLA 的一種重要性采樣技術。

實驗結果

研究在各種任務中評估了 INT4 訓練算法，包括語言模型微調、機器翻譯和圖像分類。研究使用了 CUDA 和 cutlass2 實現了所提出的 HQ-MM 和 LSS-MM 算法。除了簡單地使用 LSQ 作為嵌入層外，研究用 INT4 替換了所有浮點線性運算符，并保持最后一層分類器的全精度。并且，在此過程中，研究人員對所有評估模型采用默認架構、優化器、調度器和超參數。

收斂模型精度。下表 1 展示了收斂模型在各任務上的精度。

圖片

語言模型微調。與 LSQ+LUQ 相比，研究提出的算法在 bert-base 模型上提升了 5.5% 的平均精度、，在 bert-large 模型上提升了 25% 的平均精度。

研究團隊還展示了算法在 SQUAD、SQUAD 2.0、Adversarial QA、CoNLL-2003 和 SWAG 數據集上進一步展示了結果。在所有任務上，與 LSQ+LUQ 相比，該方法取得了更好的性能。與 LSQ+LUQ 相比，該方法在 SQUAD 和 SQUAD 2.0 上分別提高了 1.8% 和 3.6%。在更困難的對抗性 QA 中，該方法的 F1 分數提高了 6.8%。在 SWAG 和 CoNLL-2003 上，該方法分別提高了 6.7%、4.2% 的精度。

機器翻譯。研究還將所提出的方法用于預訓練。該方法在 WMT 14 En-De 數據集上訓練了一個基于 Transformer 的 [51] 模型用于機器翻譯。

HQ+LSS 的 BLEU 降解率約為 1.0%，小于 Ultra-low 的 2.1%，高于 LUQ 論文中報道的 0.3%。盡管如此，HQ+LSS 在這項預訓練任務上的表現仍然與現有方法相當，并且它支持當代硬件。

圖像分類。研究在 ImageNet21k 上加載預訓練的 ViT 檢查點，并在 CIFAR-10、CIFAR-100 和 ImageNet1k 上對其進行微調。

與 LSQ+LUQ 相比，研究方法將 ViT-B/32 和 ViT-L/32 的準確率分別提高了 1.1% 和 0.2%。在 ImageNet1k 上，該方法與 LSQ+LUQ 相比，ViT-B/32 的精度提高了 2%，ViT-L/32 的精度提高了 2.6%，ViT-L/32 的精度提高了 0.2%。

研究團隊進一步測試了算法在 ImageNet1K 上預訓練 DeiT-Small 模型的有效性，其中 HQ+LSS 與 LSQ+LUQ 相比仍然可以收斂到相似的精度水平，同時對硬件更加友好。

消融研究

研究者進行消融研究，以獨立地在挑戰性 CoLA 數據集上展示前向和反向方法的有效性。為了研究不同量化器對前向傳播的有效性，他們將反向傳播設置為 FP16。結果如下圖 3 (a) 所示。

對于反向傳播，研究者比較了簡單的極小極大量化器、LUQ 和他們自己的 LSS，并將前向傳播設置為 FP16。結果如下圖 3 (b) 所示，雖然位寬高于 2，但 LSS 取得的結果與 LUQ 相當，甚至略高于后者。

圖片

計算和內存效率

研究者比較自己提出的 HQ-MM (HQ)、計算權重梯度的 LSS（LSSWeight）、計算激活梯度的 LSS（LSSAct）的吞吐量、它們的平均吞吐量（INT4）及下圖 4 中英偉達 RTX 3090 GPU 上 cutlass 提供的基線張量核心 FP16 GEMM 實現（FP16），它的峰值吞吐量為 142 FP16 TFLOPs 和 568 INT4 TFLOPs。

圖片

研究者還比較 FP16 PyTorch AMP 以及自己 INT4 訓練算法在 8 個英偉達 A100 GPU 上訓練類 BERT 和類 GPT 語言模型的訓練吞吐量。他們改變了隱藏層大小、中間全連接層大小和批大小，并在下圖 5 中繪制了 INT4 訓練的加速比。

結果顯示，INT4 訓練算法對于類 BERT 模型實現了最高 35.1% 的加速，對于類 GPT 模型實現了最高 26.5% 的加速。

圖片