夢(mèng)晨 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

什么AI應(yīng)用每秒處理20000個(gè)AI推理請(qǐng)求，達(dá)到2024年谷歌搜索流量的1/5？

答案是獨(dú)角獸Character.ai，由Transformer作者Noam Shazeer（后面簡(jiǎn)稱(chēng)沙哥）創(chuàng)辦。

剛剛，沙哥公布了推理優(yōu)化獨(dú)門(mén)秘訣，迅速引起業(yè)界熱議。

具體來(lái)說(shuō)Character.ai在整個(gè)服務(wù)堆棧中實(shí)現(xiàn)了如下成績(jī)：

• 內(nèi)存高效架構(gòu)設(shè)計(jì)：將KV緩存大小減少20倍以上，而不會(huì)降低質(zhì)量
• Attention狀態(tài)緩存：95%請(qǐng)求無(wú)需重算
• 直接用in8精度量化訓(xùn)練：推理零損失還省顯存

Character.AI通過(guò)以上種種優(yōu)化，已經(jīng)把推理成本降低到最初的1/33，如果用市場(chǎng)上最好的商業(yè)API來(lái)支撐這種級(jí)別的流量，成本會(huì)比現(xiàn)在高出13.5倍!

眾多公布的方法中，原生int8訓(xùn)練是最受關(guān)注的。

雖然大多數(shù)技巧都來(lái)自公開(kāi)研究，但是正如網(wǎng)友所說(shuō)，知道如何把它們高效整合在一起實(shí)現(xiàn)的團(tuán)隊(duì)才是真正的護(hù)城河。

秘訣1：高效利用顯存，attention 參數(shù)量降低20倍

大模型的一大痛點(diǎn)是顯存占用高，導(dǎo)致無(wú)法支持大批量推理。Attention 層中的 Key-Value(KV)緩存便是罪魁禍?zhǔn)字弧?/p>

為了降低顯存占用，Character.AI在Attention層大動(dòng)手術(shù):

• 全面采用MQA（Multi-Query Attention）

與大多數(shù)開(kāi)源模型中采用的GQA（Grouped-Query Attention）相比，將KV緩存大小減少了 8 倍。

而MQA正是沙哥本人2019年在谷歌期間提出的，有網(wǎng)友評(píng)價(jià)“當(dāng)一個(gè)人能在生產(chǎn)環(huán)境中引用自己的論文，就達(dá)到了一個(gè)新的高度”。

• 混合注意力視野

將局部注意力與全局注意力層交織在一起，使用滑動(dòng)窗口訓(xùn)練局部注意力，將復(fù)雜度從 O(length^2 ) 降低到 O(length)。

團(tuán)隊(duì)發(fā)現(xiàn)，將大多數(shù)注意力層的注意力范圍減少到1024不會(huì)對(duì)評(píng)估指標(biāo)產(chǎn)生重大影響，包括長(zhǎng)上下文大海撈針基準(zhǔn)。在Character.ai生產(chǎn)模型中，每6層中只有1層使用全局注意力。

• 跨層KV共享

團(tuán)隊(duì)將KV緩存綁定在相鄰的注意力層上，這進(jìn)一步將 KV緩存大小減少了 2-3 倍。

對(duì)于全局注意力，跨塊綁定多個(gè)全局層的KV緩存，因?yàn)槿肿⒁饬釉陂L(zhǎng)上下文用例中主導(dǎo)KV緩存大小，團(tuán)隊(duì)發(fā)現(xiàn)跨層共享KV不會(huì)降低質(zhì)量。

下圖中左半部分是標(biāo)準(zhǔn)Transformer設(shè)計(jì)，每個(gè)注意力都是全局注意力。右半部分為Character.ai的設(shè)計(jì)，藍(lán)色框表示全局注意力，綠色框表示局部注意力，連線(xiàn)表示KV共享。

這一套組合拳下來(lái)，KV緩存大小減少20倍以上，顯存再也不是瓶頸了。

秘訣2：巧用狀態(tài)緩存，95%請(qǐng)求無(wú)需重算

Character.AI還有一招神來(lái)之筆，就是在不同對(duì)話(huà)之間緩存Attention狀態(tài)。

作為聊天機(jī)器人角色扮演服務(wù)，Character.AI上大部分對(duì)話(huà)都是連續(xù)多輪的，平均每個(gè)對(duì)話(huà)包含180條消息。如果每次都要重新計(jì)算前面的狀態(tài)，成本可想而知。

于是團(tuán)隊(duì)設(shè)計(jì)了一個(gè)緩存機(jī)制，把每個(gè)對(duì)話(huà)的Prefix和生成的消息都緩存在內(nèi)存中，供后續(xù)調(diào)用。

借鑒RadixAttention的思路，樹(shù)狀結(jié)構(gòu)的LRU緩存組織緩存的KV張量。緩存的KV值由前綴token的Rolling Hash速檢索最長(zhǎng)匹配的緩存，即使前綴只有部分匹配也能命中。

更妙的是，他們還用會(huì)話(huà)保持(Sticky Session)把同一對(duì)話(huà)路由到同一個(gè)服務(wù)器，進(jìn)一步提高緩存命中率。最終做到95%的請(qǐng)求都能復(fù)用已有緩存，大幅降低了計(jì)算成本。

下圖中，藍(lán)色框表示主機(jī)內(nèi)存上的緩存張量。綠色和黃色框表示CUDA內(nèi)存上的KV緩存。當(dāng)新查詢(xún)到達(dá)時(shí)，它檢索最長(zhǎng)匹配前綴的KV緩存，Rolling Hash系統(tǒng)允許檢索部分匹配消息的緩存。

秘訣3：直接量化訓(xùn)練，推理零損失還省顯存

最后一招，Character.AI沒(méi)有采用常見(jiàn)的“訓(xùn)練后量化”，而是直接用Int8精度訓(xùn)練模型。

這種格式雖然表達(dá)精度降低，但通過(guò)精心設(shè)計(jì)定制的矩陣乘和 Attention 內(nèi)核，不僅把訓(xùn)練效率提高了好幾倍，而且還能無(wú)損用于推理。

不過(guò)沙哥在這里暫時(shí)留了一手，表示“量化訓(xùn)練本身就是一個(gè)復(fù)雜的話(huà)題，將在以后的文章中繼續(xù)討論?！?/p>

沙哥其人

最后再來(lái)介紹一下傳奇人物Noam Shazeer本人。

他1994年拿了信息學(xué)奧賽IOI金牌，后來(lái)畢業(yè)于杜克大學(xué)。

2000年加入谷歌，當(dāng)時(shí)全公司只有200人左右，他參與了谷歌搜索的拼寫(xiě)糾正功能，后來(lái)也負(fù)責(zé)過(guò)早期廣告系統(tǒng)。

據(jù)知情人透露，在當(dāng)初面試谷歌時(shí)，沙哥就被問(wèn)到如何實(shí)現(xiàn)拼寫(xiě)糾正。他描述了一種根據(jù)其他用戶(hù)的輸入輸入記錄，進(jìn)行統(tǒng)計(jì)驗(yàn)證的方法。

面試官Gmail之父Paul Buchheit意識(shí)到，沙哥的方案比谷歌當(dāng)時(shí)使用的要好。沙哥成功入職之后就把他的面試方案寫(xiě)出來(lái)了。

在Transformer開(kāi)山之作《Attention is All You Need》研究中，沙哥最后一個(gè)加入團(tuán)隊(duì)，一來(lái)就負(fù)責(zé)重新編寫(xiě)了整個(gè)代碼。

在沙哥出手之前，Transformer早期原型性能并沒(méi)有超越當(dāng)時(shí)流行的LSTM方案，是他把早期設(shè)計(jì)中的卷積等模塊都拿掉，給出了一個(gè)極簡(jiǎn)主義方案。最終破了BLEU測(cè)試的記錄，同時(shí)計(jì)算效率也更高。

隊(duì)友用“他是一個(gè)巫師”來(lái)評(píng)價(jià)他的工程和代碼能力。

除此之外，沙哥還有驚人的遠(yuǎn)見(jiàn)。在Transformer架構(gòu)問(wèn)世不久，他就給谷歌高層寫(xiě)信，提議公司放棄整個(gè)搜索索引，并用Transformer架構(gòu)訓(xùn)練一個(gè)巨大的神經(jīng)網(wǎng)絡(luò)替代。

2021年，沙哥離開(kāi)谷歌后創(chuàng)辦了Character.AI，讓玩家簡(jiǎn)單自創(chuàng)個(gè)性化AI陪聊，目前估值約50億美元。

最近有消息稱(chēng)，Meta與馬斯克的都在爭(zhēng)取與他們合作，把聊天機(jī)器人引入社交平臺(tái)。

參考鏈接：
[1]https://research.character.ai/optimizing-inference/
[2]https://x.com/NoamShazeer/status/1803790708358410380