微軟6頁論文爆火：三進(jìn)制LLM，真香！

十三 2024-02-29 13:47:33 來源：量子位

金磊 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

現(xiàn)在，大語言模型（LLM）迎來了“1-bit時(shí)代”。

這就是由微軟和中國科學(xué)院大學(xué)在最新一項(xiàng)研究中所提出的結(jié)論——

所有的LLM，都將是1.58 bit的。

具體而言，這項(xiàng)研究提出的方法叫做BitNet b1.58，可以說是從大語言模型“根兒”上的參數(shù)下手。

將傳統(tǒng)以16位浮點(diǎn)數(shù)（如FP16或BF16）形式的存儲，統(tǒng)統(tǒng)變成了三進(jìn)制，也就是 {-1, 0, 1}。

值得注意的是，這里的“1.58 bit”并不是指每個參數(shù)占用1.58字節(jié)的存儲空間，而是指每個參數(shù)可以用1.58位的信息來表示。

在如此轉(zhuǎn)換之后，矩陣中的計(jì)算就只會涉及到整數(shù)的加法，因此會讓大模型在保持一定精度的同時(shí)，顯著減少所需的存儲空間和計(jì)算資源。

例如BitNet b1.58在3B模型大小時(shí)與Llama做比較，速度提高了2.71倍的同時(shí)，GPU內(nèi)存使用幾乎僅是原先的四分之一。

而且當(dāng)模型的規(guī)模越大時(shí)（例如70B），速度上的提升和內(nèi)存上的節(jié)省就會更加顯著！

這種顛覆傳統(tǒng)的思路著實(shí)是讓網(wǎng)友們眼前一亮，論文在X上也是受到了高度的關(guān)注：

網(wǎng)友們驚嘆“改變游戲規(guī)則”的同時(shí)，還玩起了谷歌attention論文的老梗：

1 bit is all YOU need.

那么BitNet b1.58具體又是如何實(shí)現(xiàn)的？我們繼續(xù)往下看。

把參數(shù)變成三進(jìn)制

這項(xiàng)研究實(shí)則是原班人馬在此前發(fā)表的一篇論文基礎(chǔ)之上做的優(yōu)化，即在原始BitNet的基礎(chǔ)上增加了一個額外的0值。

整體來看，BitNet b1.58依舊是基于BitNet架構(gòu)（一種Transformer），用BitLinear替換了nn.Linear。

至于細(xì)節(jié)上的優(yōu)化，首先就是我們剛才提到的“加個0”，即權(quán)重量化（weight quantization）。

BitNet b1.58模型的權(quán)重被量化為三元值{-1, 0, 1}，這相當(dāng)于在二進(jìn)制系統(tǒng)中使用了1.58 bit來表示每個權(quán)重。這種量化方法減少了模型的內(nèi)存占用，并簡化了計(jì)算過程。

其次，在量化函數(shù)設(shè)計(jì)方面，為了將權(quán)重限制在-1、0或+1之間，研究者們采用了一種稱為absmean的量化函數(shù)。

這個函數(shù)先會根據(jù)權(quán)重矩陣的平均絕對值進(jìn)行縮放，然后將每個值四舍五入到最接近的整數(shù)（-1, 0, +1）。

接下來就到了激活量化（activation quantization）這一步。

激活值的量化與BitNet中的實(shí)現(xiàn)相同，但在非線性函數(shù)之前不將激活值縮放到[0, Qb]的范圍內(nèi)。相反，激活值被縮放到[?Qb, Qb]的范圍，以此來消除零點(diǎn)量化。

值得一提的是，研究團(tuán)隊(duì)為了BitNet b1.58與開源社區(qū)兼容，采用了LLaMA模型的組件，如RMSNorm、SwiGLU等，使得它可以輕松集成到主流開源軟件中。

最后，在實(shí)驗(yàn)的性能比較上，團(tuán)隊(duì)將BitNet b1.58與FP16 LLaMA LLM在不同大小的模型上進(jìn)行了比較。

結(jié)果顯示，BitNet b1.58在3B模型大小時(shí)開始與全精度LLaMA LLM在困惑度上匹配，同時(shí)在延遲、內(nèi)存使用和吞吐量方面有顯著提升。

而且當(dāng)模型規(guī)模越大時(shí)，這種性能上提升就會越發(fā)顯著。

網(wǎng)友：能在消費(fèi)級GPU跑120B大模型了

正如上文所言，這篇研究獨(dú)特的方法在網(wǎng)上引發(fā)了不小的熱議。

DeepLearning.scala作者楊博表示：

BitNet b1.58相比原版BitNet，最大的特點(diǎn)就是允許0參數(shù)。我覺得稍微修改一下量化函數(shù)，也許可以控制0參數(shù)的比例。當(dāng)0參數(shù)的比例很大時(shí)，可以用稀疏格式存儲權(quán)重，使得平均每個參數(shù)的顯存占用甚至低于1比特。這就相當(dāng)于權(quán)重級別的MoE了。我覺得比一般的MoE更優(yōu)雅。

與此同時(shí)，他也提出了關(guān)于BitNet的缺點(diǎn)：

BitNet最大的缺點(diǎn)在于雖然能減少推理時(shí)的顯存開銷，但優(yōu)化器狀態(tài)和梯度仍然要用浮點(diǎn)數(shù)，訓(xùn)練仍然很費(fèi)顯存。我覺得如果能把BitNet和訓(xùn)練時(shí)節(jié)省顯存的技術(shù)結(jié)合起來，那么相比傳統(tǒng)半精度網(wǎng)絡(luò)，同等算力和顯存下支持更多參數(shù)，優(yōu)勢就很大了。

目前能節(jié)省優(yōu)化器狀態(tài)的顯存開銷的辦法是offloading。能節(jié)省梯度的顯存占用的辦法可能是ReLoRA。但是ReLoRA的論文實(shí)驗(yàn)只用了十億參數(shù)的模型，并沒有證據(jù)表明能不能推廣到百億、千億參數(shù)的模型。

△圖源：知乎，經(jīng)授權(quán)引用

不過也有網(wǎng)友分析認(rèn)為：

若論文成立，那么我們就能在24GB消費(fèi)級GPU上跑120B的大模型了。

那么你覺得這種新方法如何呢？

參考鏈接：
[1]https://arxiv.org/abs/2402.17764
[2]https://twitter.com/_akhaliq/status/1762729757454618720
[3]https://www.zhihu.com/question/646359036/answer/3413044355