清華朱軍團隊新作：使用4位整數訓練Transformer，比FP16快2.2倍，加速35.1%，加速AGI到來！

將活化、權重和梯度量化為4位，可望加速神經網路訓練。

然而，現有的4位訓練方法需要自訂數字格式，而現代硬體不支援這種格式。

最近，清華朱軍團隊提出了一種使用INT4演算法實現所有矩陣乘法的Transformer訓練方法。

使用超低INT4精度進行訓練，是非常具有挑戰性的。為了實現這一目標，研究者仔細分析了Transformer中活化和梯度的具體結構，為它們提出專用的量化器。

對於前向傳播，研究者確定了異常值的挑戰，並提出了Hadamard量化器來抑制異常值。

對於後向傳播，他們透過提出位元分割，來利用梯度的結構稀疏性，並利用分數取樣技術來精確量化梯度。

這種新的演算法，在自然語言理解、機器翻譯和圖像分類等廣泛任務上，都實現了具有競爭力的準確性。

原型線性算子運算速度比FP16同類算子快2.2倍，訓練速度提高了35.1%。

圖片

論文網址：https://arxiv.org/abs/2306.11987

#程式碼位址：https://github.com/xijiu9/Train_Transformers_with_INT4

全新的INT 4訓練演算法

訓練神經網路對運算的要求很高。使用低精度算術進行訓練（完全量化訓練/FQT）有望提高計算和記憶體效率。

FQT方法在原來的全精度計算圖中添加了一些量化器和反量化器，並用消耗更小的低精度浮點運算，取代了消耗更高的浮點運算。

FQT的研究旨在降低訓練數值精度，而不犧牲太多的收斂速度或精確度。

所需的數值精確度已從FP16降低到FP8、INT32 INT8和INT8 INT5。

FP8訓練是在Transformer引擎的Nvidia H100 GPU中實現的，加速了大規模Transformer的訓練。最近的訓練數值精度，已經降到4位了。

然而，這些4位訓練方法不能直接用於加速，因為它們需要自訂數字格式，而現代硬體不支援這些格式。

首先，前向傳播中的不可微量化器，會使損失情況變得崎嶇不平，基於梯度的最佳化器很容易陷入局部最優。

其次，梯度僅以低精度近似計算。這種不精確的梯度會減慢訓練過程，甚至導致訓練不穩定或發散。

而在這項工作中，研究者為Transformer提出了一種新穎的INT4訓練演算法。

圖片

訓練Transformer的所有高消耗的線性運算，都可以寫在矩陣乘法（MM）的形式中。

這種MM形式，可以讓我們設計更靈活的量化器，透過利用Transformer中活化、權重和梯度的特定結構，就可以更好地近似於FP32矩陣乘法。

隨機數值線性代數 (RandNLA) 領域的進步，被這種量化器充分利用。

對於前向傳播，研究者發現，活化中的異常值是精確度下降的主要原因。

為了抑制異常值，他們提出了Hadamard量化器，它會對激活矩陣的變換版本進行量化。這種變換是區塊對角Hadamard矩陣，它將離群值中攜帶的資訊傳播到矩陣的鄰近條目，從而縮小了離群值的數值範圍。

對於後向傳播，他們利用了活化梯度的結構稀疏性。研究者發現，有些token的梯度非常大。同時，其餘大多數token的梯度非常均勻，甚至比較大梯度的量化殘差更均勻。

圖片

因此，與其計算所有梯度，不如節省計算較大梯度殘差的運算資源。

為了利用這種稀疏性，研究者提出了位元分割，將每個token的梯度分割為高4位元和低4位元。

然後，透過槓桿分數採樣（leverage score sampling）來選擇資訊最豐富的梯度，這是RandNLA的重要採樣技術。

圖片

結合前向和後向傳播的量化技術，研究者提出了一種使用INT4MM進行Transformer中所有線性運算的演算法， 並且評估了在各種任務上訓練Transformer的演算法，包括自然語言理解、問答、機器翻譯和影像分類。

與現有的4位訓練演算法相比，他們的演算法實現了有競爭力的或更高的精度。

此外，演算法與GPU等當代硬體相容，因為它不需要FP4或對數格式等自訂的數位格式。

這種原型量化 INT4 MM算子實現，速度比FP16MM基線快2.2倍，並且將訓練速度提高了35.1%。

相關工作

完全量化訓練

完全量化訓練(FQT) 方法透過將激活、權重和梯度量化為低精度來加速訓練，因此訓練期間的線性和非線性算子可以用低精度算術來實現。

FQT的研究設計了新穎的數值格式和量化演算法，可以更好地逼近全精度張量。

目前的研究前沿是4位FQT。由於梯度的數值範圍很大以及從頭開始訓練量化網路的最佳化問題，FQT具有挑戰性。

由於這些挑戰，現有的4位元FQT 演算法在某些任務上的精確度仍然下降了1-2.5%，並且無法支援當代硬體。

圖片

其他有效的訓練方法

混合專家不會增加訓練預算的情況下提高了模型容量。

結構性dropout利用計算有效的方法來正規化模型。高效率的注意力降低了計算注意力的二次時間複雜度。

分散式訓練系統透過利用更多的運算資源，減少了訓練時間。

研究者降低數值精度的工作與這些方向具有正交性。

圖片

前向傳播

##神經網路訓練是一個迭代優化過程，透過前向和後向傳播計算隨機梯度。

研究團隊使用4位元整數(INT4)演算法加速前向和後向傳播。

正向傳播能以線性和非線性(GeLU, normalization, softmax等)算符的組合來實現。

在我們的訓練過程中，我們用INT4算術加速所有線性運算符，並將所有計算量較小的非線性運算符保留在16位元浮點(FP16)格式中。

Transformer中的所有線性運算都可以寫成矩陣乘法(MM)的形式。

為了方便表述，本文考慮以下簡單矩陣乘法的加速：

圖片#

這種MM的最主要用例是全連接層。

考慮一個輸入形狀為(批次大小S，序列長度T，維度D)的Transformer。

全連接層可以表述成上邊的公式，其中X是N = STtoken的激活，W是權重矩陣。

對於注意力層，可能需要批次矩陣乘法(BMMS)。

我們提出的技術可以應用在BMMS上。

學習步長量化（Learned Step Quantization）

為了加速訓練，必須使用整數運算來計算前向傳播。

研究者為此目的，利用學習步長量化器（LSQ）。

LSQ是靜態量化，他的量化尺度不依賴輸入的方法，因此比動態方法消耗更小，量化方法，需要在每次迭代時動態計算量化尺度。

激活異常值

#簡單地將LSQ應用到具有4位元激活/權重的FQT會導致精度下降，因為會啟動異常值。

圖片

如上圖所示，活化有一些離群值條目，它們是其規模比其他條目大得多。

不幸的是，Transformers傾向於將資訊儲存在這些異常值中，而這樣的截斷會嚴重損害準確性。

當訓練任務是在一些新的下游任務上微調預訓練模型時，異常值問題尤其明顯。

因為預訓練模型比隨機初始化包含更多的異常值 。

Hadamard量化

我們提出了Hadamard量化（HQ）來解決例外值問題。

其主要想法是將另一個具有較少異常值的線性空間中的矩陣進行量化。

激活矩陣中的異常值形成了一個特徵結構（feature-wise structure）。

他們通常集中在幾個維度上，也就是說X中只有幾列顯著大於其他列。

哈達瑪轉換（Hardamand transform）是一個線性變換，它可以將異常值分攤到其他條目中。

後向傳播

#現在我們考慮使用INT4運算來加速線性層的後向傳播。

我們將在本節中討論激活梯度/權重梯度的計算。

梯度的結構稀疏性

我們注意到，在訓練過程中梯度矩陣往往非常稀疏。

而且稀疏性具有這樣的結構：

#的幾行（例如tokens）具有較大的條目，而大多數其他行卻接近全零向量。

圖片

這種結構稀疏性源自於現代神經網路的嚴重過度參數化。

幾乎在整個訓練過程中，網路都以超參數化方案運行，除了一些困難的例子之外，它可以很好地適應大多數訓練資料。

因此，對於擬合良好的數據點，（激活）梯度將接近零。

研究人員發現對於預訓練任務，例如，經過幾個訓練週期後，結構稀疏性很快就會出現。

對於微調任務，梯度整個訓練過程中總是稀疏的。

位元分割（Bit Splitting）和槓桿分數取樣(Leverage Score Sampling)

如何設計梯度量化器，以利用結構稀疏性在反向傳播期間準確計算MM呢？

進階的想法是：梯度的許多行都是如此小，對參數梯度影響很小，但浪費了大量的計算量。

另一方面，大行無法用INT4精確表示。

我們放棄掉一些小行並使用節省下來的運算能力來更準確地表示大行。

實驗

研究人員在包含語言模型在內的各種任務上評估我們的INT4訓練演算法微調、機器翻譯和影像分類。

研究人員用CUDA和cutlass執行了他們提出的HQ-MM和LSS-MM演算法。

研究人員用INT4實作取代所有浮點線性運算符，但沒有簡單地使用LSQ來嵌入層，並保持最後一個分類器層的精確度。

最後研究人員對所有評估的模型採用了預設架構、最佳化器、調度器和超參數。

收斂模型精確度

研究人員在下表中比較了收斂模型在各種任務上的準確度。

圖片

作為對照的方法包括全精度訓練（FP）、INT8訓練（INT8）、FP4訓練（“超低」），使用LSQ進行活化和權重(LSQ LUQ)的4 位對數量化，以及我們這種利用HQ進行前向傳播，利用LSS進行反向傳播（HQ LSS）的演算法。

「超低」沒有公開的實現，因此我們僅列出了它在機器上的原始論文中的效能翻譯任務。

除了大型機器翻譯任務和大型視覺Transformer任務之外，我們將每次運行重複三次，並將標準差報告為表中的下標。

研究人員沒有進行任何類型的知識蒸餾或資料增強。

消融實驗

研究者進行的消融實驗目的是展示前向和後向方法的有效性。

研究不同量化器的前向傳播的有效性，我們將後向傳播留在FP16。

結果如下圖所示。

圖片

計算與記憶體效率

最後，研究人員透過評估他們的原型實現，展示了他們的方法加速神經網路訓練的潛力。

而且他們的實作還沒有完全最佳化。

研究人員也沒有將線性算子與非線性和歸一化進行融合。

因此，結果無法完全反映INT4訓練演算法的潛力。

完全最佳化的實作需要大量工程，超出了我們論文的討論範圍。

結論

研究人員提出了一種對硬體很友善的Transformer INT4的訓練方法。

透過分析Transformer中MM的屬性，研究人員提出了HQ和LSS方法來量化活化和梯度，同時保持準確性。

在幾個重要任務上，我們的方法與現有的INT4方法表現相當，甚至更好。

研究人員的這些工作可能會擴展到除了Transformers之外的其他MM架構中，例如 MLP-Mixer、圖神經網路和循環神經網路網路。

這是他們未來的研究方向。

更廣泛的影響：研究人員的演算法可以提高效率並減少訓練神經網路的能源消耗，這有助於減少深度學習造成的碳排放。

但是，高效的訓練演算法也可能促進那些，對於人來安全存在隱患的大語言模型和惡意人工智慧應用程式的開發。

例如，會被用於虛假內容產生的相關模型和應用。

限制：這項工作的主要限制是它只能加速具有較大規模的矩陣乘法（線性層）的大模型，但不能加速卷積層。

而且，所提出的方法還不能很好地適用於OPT-175B等超大模型。

據我們所知，即使是INT8訓練對於這些超大型模型來說仍然是尚待解決的問題。

以上是清華朱軍團隊新作：使用4位整數訓練Transformer，比FP16快2.2倍，加速35.1%，加速AGI到來！的詳細內容。更多資訊請關注PHP中文網其他相關文章！