用于大型Transformer的8-bit矩陣乘法介紹

原文地址：A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using transformers, accelerate and bitsandbytes

相關博客
【深度學習】【分布式訓練】Collective通信操作及Pytorch示例
【自然語言處理】【大模型】大語言模型BLOOM推理工具測試
【自然語言處理】【大模型】GLM-130B：一個開源雙語預訓練語言模型
【自然語言處理】【大模型】用于大型Transformer的8-bit矩陣乘法介紹
【自然語言處理】【大模型】BLOOM：一個176B參數且可開放獲取的多語言模型
【自然語言處理】【大模型】PaLM：基于Pathways的大語言模型
【自然語言處理】【chatGPT系列】大語言模型可以自我改進
【自然語言處理】【ChatGPT系列】FLAN：微調語言模型是Zero-Shot學習器
【自然語言處理】【ChatGPT系列】ChatGPT的智能來自哪里？
【自然語言處理】【ChatGPT系列】Chain of Thought：從大模型中引導出推理能力
【自然語言處理】【ChatGPT系列】InstructGPT：遵循人類反饋指令來訓練語言模型
【自然語言處理】【ChatGPT系列】大模型的涌現能力

一、簡介

? 語言模型正變的越來越大，PaLM已有有540B的參數量，而OPT、GPT-3和BLOOM則大約有176B參數量。下圖是近些年語言模型的尺寸。

? 這些模型很難在常用設備上運行。例如，僅僅推理BLOOM-176B就需要8x 80GB A00 GPUs。而為了微調BLOOM-176B則需要72個GPU。PaLM則需要更多的資源。

? 這些巨型模型需要太多GPUs才能運行，因此需要尋找方法來減少資源需求并保證模型的性能。已經有各種技術用來減小模型尺寸，例如量化、蒸餾等。在完成BLOOM-176B訓練后，HuggingFace和BigScience逐步探索在少量GPU上運行大模型的方法。最終，設計出了Int8量化方法，該方法在不降低大模型性能的情況下，將顯存占用降低了1至2倍。

二、機器學習中常用的數據類型

? 浮點數在機器學習中也被稱為"精度"。模型大小是有參數量及參數精度決定的，通常是float32、float16和bfloat16。

我們開始于不同浮點數的基本理解，在機器學習的背景下也被稱為"精度"。模型的大小由其參數數量和精度決定，通常是float32、float16和bfloat16。下圖：

? Float32(FP32)是標準的IEEE 32-bit浮點數表示，使用這種類型可以表示范圍廣泛的浮點數。在FP32中，8bits被用于"指數"，23bits被用于"尾數", 1 bit則用于符號位。大多數的硬件都支持FP32操作和指令。

? 在Float16(FP16)數據類型中，5 bits被用作"指數"，10 bits用于"尾數"。這使得FP16數的表示范圍明顯小于FP32，導致有上溢和下溢的風險。例如，若你做$10k\times 10k$，最終得到100k。這在FP16中是不可能的，因為其最大表示為64K。因此你會得到NaN的結果，若像神經網絡那樣順序執行，所有先前的工作都會被破壞。通常，loss縮放能夠一定程度上克服這個問題，但并不總是有用。

? 因此，創建了一種新格式bfloat15(BF16)來避免這種問題。在BF16中，8bits被用于表示"指數", 7bits被用于表示"尾數"。這意味著BF16能夠保留和FP32相同的動態范圍，但是損失了3bits的精度。BF16可以表示巨大的數，但是精度上比FP16差。

? 在Ampere架構中，NVIDIA也引入了TensorFloat-32(TF32)精度格式，其僅使用19 bits就合并了BF16的動態范圍和FP16的精度。其目前僅在內部某些操作中使用。

? 在機器學習的術語中FP32被稱為全精度(4 bytes)，BF16和FP16則稱為半精度(2 bytes)。int8(INT8)數據類型則是由8 bits表示的數，其能夠存儲$2^8$個不同的值([0,255]或者[-128, 127])

? 理想情況下，訓練和推理應該在FP32上進行，但是其比FP16/BF16慢兩倍。因此，采用一種混合精度的方法，模型權重仍然是FP32，前向和后向傳播則使用FP16/BF16，從而加快訓練速度。P16/BF16被用來更新FP32權重。

? 可以通過參數量乘以浮點數精度的大小來計算模型所占用的bytes量。例如，若模型使用bfloat16版本的BLOOM-176B模型，那么模型大小為176*10**9 x 2 bytes = 352GB！這個量級對于適配少量GPU來說相當有挑戰。

? 但是我們是否可以使用不同的數據類型以更少的存儲空間來保存這些權重？一種稱為量化的方法被廣泛的應用于Deep Learning。

三、模型量化介紹

? 通過實驗發現，在推理中使用2 bytes的BF16/FP16精度能夠幾乎達到4 bytes的FP32精度相同的效果，而且模型尺寸可以減少一半。若能夠進一步削減那就太棒了，但是在更低的精度上推理質量開始急劇下降。為了解決這個問題，我們引入了8 bits量化。該方法使用四分之一的精度，這樣僅需要1/4的模型尺寸！但是，其不是通過丟掉另外一半bits來實現的。

? 量化基本上是從一種數據類型"舍入"為另一種數據類型來完成的。例如，若一個數量類型范圍0…9，另一個范圍則是0…4。那么第一個數據類型中的"4"將會被舍入為第二種數據類型中的"2"。然而，若第一種數據類型中的"3"，其會位于第二種數據類型的1和2之間，然后通常會被舍入為"2"。也就是說第一種數據類型中的"4"和"3"都會對應第二種數據類型中的"2"。這表明量化是可能帶來信息丟失的噪音過程，一種有損壓縮。

? 有兩種常見的8-bit量化技術：zero-point量化和absolute maximum(absmax)量化。zero-point量化和absmax量化會將浮點數值映射至更加緊湊的int8(1 byte)值。這些方法首先會將輸入按照量化常數進行縮放，從而實現規范化。

? 舉例來說，在zero-point量化中，若范圍是$[?1.0\dots 1.0]$并希望量化至范圍$[?127\dots 127]$。那么應該按照因子127進行縮放，然后四舍五入至8-bit精度。為了還原原始值，需要將int8的值除以量化因子127。例如，0.3被縮放為$0.3\times 127=38.1$，然后四舍五入為38。若要恢復，則$38/127=0.2992$。在這個例子中量化誤差為0.008。隨著這些微小的誤差在模型各個層中傳播，會逐步積累和增長并導致性能下降。

? 再來看看absmax量化的細節。為了在absmax量化中完成fp16和int8的映射，需要先除以張量中的絕對最大值(令整個張量介于-1至1之間)，然后在乘以目標數據類型的總范圍。例如，在一個向量上應用absmax量化，該向量為
$$v=\left[\begin{array}{cccccccc}1.2& ?0.5& ?4.3& 1.2& ?3.1& 0.8& 2.4& 5.4\end{array}\right]$$
從向量中選擇最大值，即5.4。而int8的范圍為[-127,127]，所以量化過程為$v/5.4\times 127=v\times \frac{127}{5.4}\approx v\times 23.5$，即整個向量乘以縮放因子23.5。最終得到的量化后向量為
$$\left[\begin{array}{cccccccc}28& ?12& ?101& 28& ?73& 19& 56& 127\end{array}\right]$$

? 為了還原原始值，可以使用全精度的int8數除以量化因子23.5。但是由于四舍五入的原因，會丟失一些精度。

? 這些技巧能夠以多種方式組合。例如，當涉及矩陣乘法時，ow-wise或者vector-wise量化可以使得結果更加準確。以矩陣乘法$A\times B=C$為例，相對于使用每個張量的絕對最大值來規范張量，vector-vise量化則會尋找矩陣A每行的絕對最大值和矩陣B每列的絕對最大值。然后通過除以這些絕對最大值向量來規范化矩陣A和B。然后執行$A\times B$來得到C。為了最終返回FP16精度的值，通過計算A和B絕對最大值向量的外積來反規范化。

? 這些技術雖然能夠量化模型，但是在較大模型上會帶來性能下降。Hugging Face Transformers和Accelerate庫集成了一種稱為LLM.int8()的8-bit量化算法，能夠在176B參數量模型上使用且不降低模型效果。

四、LLM.int8()簡介

? 理解Transformer中與規模相關的涌現特性對于理解為什么傳統量化方式在大模型中失敗至關重要。性能的下降是由異常特征值導致的，會在后面解釋這一情況。LLM.int8()算法本質上可以由三個步驟來完成矩陣乘法：

對輸入的hidden states逐列的提取異常值(即大于某個閾值的值)；

分別對FP16中的異常值和INT8中的非異常值執行矩陣乘法；

對非異常的結果進行反量化，并將兩者結果合并來獲得最終的FP16結果；

三個步驟如下圖所示：

1. 異常值特征

? 在整個分布之外的值，稱為異常值。異常值檢測被廣泛使用，而擁有特征分布的先驗知識有助于異常值檢測任務。

? 具體來說，我們觀察到經典的量化算法在超過6B參數量的transformer模型上失效了。雖然在較小的模型上也能觀測到較大的異常值特征。但是，我們觀察到一個參數量的閾值，transformer中的異常值會系統性的出現在每個層中。

? 由于8-bit精度的局限性，因此僅使用幾個特別大的值來量化向量將導致非常差的結果。此外，transformer架構的內在特征就是將所有的元素連接在一起，這將導致錯誤跨越多層傳播并被加劇。因此，開發出了混合精度分解來實現這種極端異常值的量化。

2. MatMul內部

? 一旦得到hidden state，使用自定義閾值來抽取異常值并分解矩陣為上述兩部分。我們發現使用6作為閾值進行抽取可以完整的恢復推理性能。異常值部分以fp16實現，所以是經典的矩陣乘法；而8-bit則是通過vector-wise量化將模型權重和hidden state量化至8-bit的精度。即hidden-state使用row-wise量化，模型權重使用column-wise量化。經過這個步驟后，再將結果反量化并以半精度返回。

3. 零退化意味著什么

? 如何評估性能下降？8-bit模型到底損失了多少性能？這里在8-bit模型和native模型上運行了常見的基準，分別針對OPT-175B和BLOOM-176B。

• 對于OPT-175B

• 對于BLOOM-176B

? 可以看到這些模型的性能下降為0，因為這些指標的絕對差值小于標準誤差。

4. 比native模型更快？

? LLM.int8()方法的主要目標在不降低性能的情況下，使得大模型更容易被使用。但是，如果該方法非常的慢則就不實用了。所以，我們對多個模型的生成速度進行了基準測試。實驗發現使用LLM.int8()的BLOOM-176B要比fp16版本慢15%至23%，這是一個可以接受的范圍。但是較小的模型下降會更多。開發人員正在逐步優化這個問題。

五、Transformers集成

1. 使用

? 本文重點描述的模塊是Linear8bitLt，你可以直接從bitsandbytes庫中引入。其來自于經典的torch.nn模塊，并使用下面的代碼來輕易的使用和部署。

? 下面是一個使用bitsandbytes將一個小模型轉換為int8類型。

• 正確的引入

import torch import torch.nn as nnfrom bitsandbytes.nn import Linear8bitLt

• 先定義一個fp16的模型

fp16_model = nn.Sequential(nn.Linear(64, 64),nn.Linear(64, 64))

• 假設該模型已經完成訓練，保存模型

torch.save(fp16_model.state_dict(), "model.pt")

• 現在再定義一個int8模型

int8_model = nn.Sequential(Linear8bitLt(64, 64, has_fp16_weights=False),Linear8bitLt(64, 64, has_fp16_weights=False) )

添加參數has_fp16_weights很重要。默認值為True，其被用于Int8/FP16混合精度訓練。然而，這里關注的是推理，所以將其設置為False。

• 現在將fp16的模型加載至int8模型中

int8_model.load_state_dict(torch.load("model.pt")) # print(int8_model[0].weight) int8_model = int8_model.to("cuda:0") # 執行該代碼時會進行量化 # print(int8_model[0].weight)

通過輸出print(int8_model[0].weight)可以看到模型被量化為Int8類型，那么怎么還原為FP16權重呢？

(int8_model[0].weight.CB * int8_model[0].weight.SCB) / 127

• 使用int8模型進行推理

input_ = torch.randn((1,64), dtype=torch.float16) hidden_states = int8_model(input_.to(torch.device('cuda:0')))

2. 你只需要accelerate

? 當使用大模型時，acceleate庫包含了有用的程序。init_empty_weights方法特別有用，因為任何模型(無論大小)都可以作為上下文管理器使用此方法進行初始化，而無需為模型權重分配任何內存。

import torch.nn as nn from accelerate import init_empty_weightswith init_empty_weights():model = nn.Sequential(*[nn.Linear(100000, 100000) for _ in range(1000)])

這個初始化的模型會被放置至Pytorch的元設備上，其是一種不用分配存儲空間來表示shape和dtype的潛在機制。

起初，該函數在.from_pretrained函數中被調用，并將所有參數重寫為torch.nn.Parameter。但是，這不符合我們的需求，因為希望在Linear8bitLt模塊中保留Int8Params類。因此我們將

module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))

修改為

param_cls = type(module._parameters[name]) kwargs = module._parameters[name].__dict__ module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)

通過這個修改，我們可以通過自定義函數在沒有任何內存消耗的情況下，利用這個上下文管理器將所有的nn.Linear替換為bnb.nn.Linear8bitLt。

def replace_8bit_linear(model, threshold=6.0, module_to_not_convert="lm_head"):for name, module in model.named_children():if len(list(module.children())) > 0:# 遞歸replace_8bit_linear(module, threshold, module_to_not_convert)if isinstance(module, nn.Linear) and name != module_to_not_convert:with init_empty_weights():model._modules[name] = bnb.nn.Linear8bitLt(module.in_features,module.out_features,module.bias is not None,has_fp16_weights=False,threshold=threshold,)return model

該函數會遞歸的將元設備上的所有nn.Linear替換為Linear8bitLt模塊。屬性has_fp16_weights必須被設置為False，以便加載int8權重和量化信息。

3. 如何在transformers中使用

from transformers import AutoTokenizer, AutoModelForCausalLMdef inference(payload, model, tokenizer):input_ids = tokenizer(payload, return_tensors="pt").input_ids.to(model.device)print(f"輸入:\n {payload}")logits = model.generate(input_ids, num_beams=1, max_new_tokens=128)print(f"生成:\n {tokenizer.decode(logits[0].tolist()[len(input_ids[0]):])}")model_name = "bigscience/bloomz-7b1-mt" payload = "一個傳奇的開端，一個不滅的神話，這不僅僅是一部電影，而是作為一個走進新時代的標簽，永遠彪炳史冊。你認為這句話的立場是贊揚、中立還是批評？" tokenizer = AutoTokenizer.from_pretrained(model_name) model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True) model_native = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") # 比較推理結果 inference(payload, model_8bit, tokenizer) inference(payload, model_native, tokenizer) # 計算顯存節約程度 mem_fp16 = model_native.get_memory_footprint() mem_int8 = model_8bit.get_memory_footprint() print(mem_fp16/mem_int8)

總結

以上是生活随笔為你收集整理的【自然语言处理】【大模型】用于大型Transformer的8-bit矩阵乘法介绍的全部內容，希望文章能夠幫你解決所遇到的問題。

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