類GPT模型訓(xùn)練提速26.5%，清華朱軍等人用INT4算法加速神經(jīng)網(wǎng)絡(luò)訓(xùn)練

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

模型訓(xùn)練得快不快，這與激活值、權(quán)重、梯度等因素的要求緊密相關(guān)。

神經(jīng)網(wǎng)絡(luò)訓(xùn)練需要一定計(jì)算量，使用低精度算法（全量化訓(xùn)練或 FQT 訓(xùn)練）有望提升計(jì)算和內(nèi)存的效率。FQT 在原始的全精度計(jì)算圖中增加了量化器和去量化器，并將昂貴的浮點(diǎn)運(yùn)算替換為廉價(jià)的低精度浮點(diǎn)運(yùn)算。

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

最近訓(xùn)練數(shù)值精度已被壓低到 4 位（ 4 bits）。Sun 等人成功訓(xùn)練了幾個(gè)具有 INT4 激活 / 權(quán)重和 FP4 梯度的當(dāng)代網(wǎng)絡(luò)；Chmiel 等人提出自定義的 4 位對數(shù)數(shù)字格式，進(jìn)一步提高了精度。然而，這些 4 位訓(xùn)練方法不能直接用于加速，因?yàn)樗鼈冃枰远x數(shù)字格式，這在當(dāng)代硬件上是不支持的。

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

本文為流行的神經(jīng)網(wǎng)絡(luò) Transformer 提出了新的 INT4 訓(xùn)練算法。訓(xùn)練 Transformer 所用的成本巨大的線性運(yùn)算都可以寫成矩陣乘法（MM）的形式。MM 形式使研究人員能夠設(shè)計(jì)更加靈活的量化器。這種量化器通過 Transformer 中的特定的激活、權(quán)重和梯度結(jié)構(gòu)，更好地近似了 FP32 矩陣乘法。本文中的量化器還利用了隨機(jī)數(shù)值線性代數(shù)領(lǐng)域的新進(jìn)展。

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

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

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

結(jié)合前向和反向傳播的量化技術(shù)，本文提出一種算法，即對 Transformer 中的所有線性運(yùn)算使用 INT4 MMs。研究評估了在各種任務(wù)上訓(xùn)練 Transformer 的算法，包括自然語言理解、問答、機(jī)器翻譯和圖像分類。與現(xiàn)有的 4 位訓(xùn)練工作相比，研究所提出的算法實(shí)現(xiàn)了相媲美或更高的精度。此外，該算法與當(dāng)代硬件 (如 GPU) 是兼容的，因?yàn)樗恍枰远x數(shù)字格式 (如 FP4 或?qū)?shù)格式)。并且研究提出的原型量化 + INT4 MM 算子比 FP16 MM 基線快了 2.2 倍，將訓(xùn)練速度提高了 35.1%。

前向傳播

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

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

學(xué)得的步長量化

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

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

激活異常值

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

在這種情況下，步長 s_X 在量化粒度和可表示數(shù)值范圍之間進(jìn)行權(quán)衡。如果 s_X 很大，則可以很好地表示異常值，同時(shí)代價(jià)是以粗略的方式表示其他大多數(shù)項(xiàng)。如果 s_X 很小，則必須截?cái)?[?Q_Ns_X, Q_Ps_X] 范圍之外的項(xiàng)。

Hadamard 量化

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

激活矩陣中的異常值可以形成特征級結(jié)構(gòu)。這些異常值通常集中在幾個(gè)維度上，也就是 X 中只有幾列顯著大于其他列。作為一種線性變換，Hadamard 變換可以將異常值分?jǐn)偟狡渌?xiàng)中。具體地，Hadamard 變換 H_k 是一個(gè) 2^k × 2^k 矩陣。

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

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

其中逆變換彼此之間相互抵消，并且 MM 可以實(shí)現(xiàn)如下。

反向傳播

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

梯度的結(jié)構(gòu)稀疏性

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

Bit 拆分和平均分?jǐn)?shù)采樣

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

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

機(jī)器翻譯。研究還將所提出的方法用于預(yù)訓(xùn)練。該方法在 WMT 14 En-De 數(shù)據(jù)集上訓(xùn)練了一個(gè)基于 Transformer 的 [51] 模型用于機(jī)器翻譯。

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

圖像分類。研究在 ImageNet21k 上加載預(yù)訓(xùn)練的 ViT 檢查點(diǎn)，并在 CIFAR-10、CIFAR-100 和 ImageNet1k 上對其進(jìn)行微調(diào)。

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

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

消融研究

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

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

計(jì)算和內(nèi)存效率

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

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

結(jié)果顯示，INT4 訓(xùn)練算法對于類 BERT 模型實(shí)現(xiàn)了最高 35.1% 的加速，對于類 GPT 模型實(shí)現(xiàn)了最高 26.5% 的加速。

更多技術(shù)和實(shí)驗(yàn)細(xì)節(jié)請參閱原論文。

THE END