什么 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ù)討論。”
沙哥其人
最后再來(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
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。