在大型語(yǔ)言模型（LLM）的世界中，處理多輪對(duì)話一直是一個(gè)挑戰(zhàn)。前不久麻省理工Guangxuan Xiao等人推出的StreamingLLM，能夠在不犧牲推理速度和生成效果的前提下，可實(shí)現(xiàn)多輪對(duì)話總共400萬(wàn)個(gè)token的流式輸入，22.2倍的推理速度提升。

但StreamingLLM使用原生PyTorch實(shí)現(xiàn)，對(duì)于多輪對(duì)話推理場(chǎng)景落地應(yīng)用的低成本、低延遲、高吞吐等需求仍有優(yōu)化空間。

Colossal-AI團(tuán)隊(duì)開(kāi)源了SwiftInfer，基于TensorRT實(shí)現(xiàn)了StreamingLLM，可以進(jìn)一步提升大模型推理性能46%，為多輪對(duì)話推理提供了高效可靠的落地方案。

開(kāi)源地址：https://github.com/hpcaitech/SwiftInfer

StreamingLLM簡(jiǎn)介

大語(yǔ)言模型能夠記住的上下文長(zhǎng)度，直接影響了ChatGPT等大模型應(yīng)用與用戶互動(dòng)的質(zhì)量。

如何讓LLM在多輪對(duì)話場(chǎng)景下保持生成質(zhì)量，對(duì)推理系統(tǒng)提出了更高的要求，因?yàn)長(zhǎng)LM在預(yù)訓(xùn)練期間只能在有限的注意力窗口的限制下進(jìn)行訓(xùn)練。

常見(jiàn)的KV Cache機(jī)制能夠節(jié)約模型計(jì)算的時(shí)間，但是在多輪對(duì)話的情景下，key和value的緩存會(huì)消耗大量的內(nèi)存，無(wú)法在有限的顯存下無(wú)限擴(kuò)展上下文。

同時(shí)，訓(xùn)練好的模型在不做二次微調(diào)的前提下也無(wú)法很好地泛化到比訓(xùn)練序列長(zhǎng)度更長(zhǎng)的文本，導(dǎo)致生成效果糟糕。

來(lái)源：https://arxiv.org/pdf/2309.17453.pdf

StreamingLLM為了解決了這個(gè)問(wèn)題，通過(guò)觀察了注意力模塊中Softmax的輸出，發(fā)現(xiàn)了attention sink的現(xiàn)象。

我們知道注意力機(jī)制會(huì)為每一個(gè)token分配一個(gè)注意力值，而文本最初的幾個(gè)token總是會(huì)分配到很多無(wú)用的注意力。

當(dāng)我們使用基于滑動(dòng)窗口的注意力機(jī)制時(shí)，一旦這幾個(gè)token被踢出了窗口，模型的生成效果就會(huì)迅速崩潰。但只要一直把這幾個(gè)token保留在窗口內(nèi)，模型就能穩(wěn)定地生成出高質(zhì)量的文本。

比起密集注意力（Dense Attention）、窗口注意力（Window Attention）以及帶重計(jì)算的滑動(dòng)窗口注意力(Sliding Window w/ Re-computing)，StreamingLLM基于attention sink的注意力機(jī)制無(wú)論是在計(jì)算復(fù)雜度還是生成效果上都表現(xiàn)優(yōu)異。

在不需要重新訓(xùn)練模型的前提下，StreamingLLM能夠直接兼容目前的主流大語(yǔ)言模型并改善推理性能。

SwiftInfer：基于TensorRT的StreamingLLM實(shí)現(xiàn)

為了將StreamingLLM這一技術(shù)更好的應(yīng)用到落地場(chǎng)景，Colossal-AI團(tuán)隊(duì)成功地將StreamingLLM方法與TensorRT推理優(yōu)化結(jié)合，不僅繼承了原始StreamingLLM的所有優(yōu)點(diǎn)，而且還具有更高的運(yùn)行效率。

此外，使用TensorRT-LLM的API，還能夠獲得接近于PyTorch API的模型編寫(xiě)體驗(yàn)?；赥ensorRT-LLM，團(tuán)隊(duì)重新實(shí)現(xiàn)了KV Cache機(jī)制以及帶有位置偏移的注意力模塊。

如下圖所示，假設(shè)窗口大小為10個(gè)token，隨著生成的token增加（由黃色方塊表示），我們?cè)贙V緩存中將中間的token踢出，與此同時(shí)，始終保持著文本開(kāi)始的幾個(gè)token（由藍(lán)色方塊表示）。由于黃色方塊的位置會(huì)發(fā)生變化，在計(jì)算注意力時(shí)，我們也需要重新注入位置信息。

需要注意的是，StreamingLLM不會(huì)直接提高模型能訪問(wèn)的上下文窗口，而是能夠在支持流式超多輪對(duì)話的同時(shí)保證模型的生成效果。

大模型無(wú)限輸入流推理加速46%

原版本的StreamingLLM可以可靠地實(shí)現(xiàn)超過(guò)400萬(wàn)個(gè)token的流式輸入，實(shí)現(xiàn)了比帶重計(jì)算的滑動(dòng)窗口注意力機(jī)制高出22.2倍的速度提升。

Colossal-AI團(tuán)隊(duì)發(fā)布的SwiftInfer可以進(jìn)一步提升推理性能，最多帶來(lái)額外的最多46%的推理吞吐速度提升，為大模型多輪對(duì)話推理提供低成本、低延遲、高吞吐的最佳實(shí)踐。TensorRT-LLM團(tuán)隊(duì)也在同期對(duì)StreamingLLM進(jìn)行了類似支持。

Colossal-AI社區(qū)動(dòng)態(tài)

Colossal-AI目前已獲得GitHub星數(shù)三萬(wàn)五千多顆，位列全球TOP400，細(xì)分賽道排名世界第一，可通過(guò)高效多維并行、異構(gòu)內(nèi)存等，降低AI大模型訓(xùn)練/微調(diào)/推理的開(kāi)發(fā)與應(yīng)用成本，提升模型任務(wù)表現(xiàn)，降低GPU需求。作為主流開(kāi)源AI大模型系統(tǒng)社區(qū)，Colossal-AI生態(tài)在多方面保持活躍更新。

Colossal-LLaMA-2-13B開(kāi)源

Colossal-LLaMA-2-13B模型，僅用25B token數(shù)據(jù)和萬(wàn)元算力，效果遠(yuǎn)超基于 LLaMA-2 的其他中文漢化模型。

即使與其他采用中文語(yǔ)料，可能花費(fèi)上千萬(wàn)元成本，從頭預(yù)訓(xùn)練的各大知名模型相比，Colossal-LLaMA-2在同規(guī)模下仍表現(xiàn)搶眼。

13B 版本通過(guò)構(gòu)建更為完善的數(shù)據(jù)體系，在知識(shí)性內(nèi)容掌握程度，自然語(yǔ)言處理任務(wù)理解程度，以及安全性，價(jià)值觀等問(wèn)題上，都有質(zhì)的提升。

Colossal-AI云平臺(tái)

Colossal-AI云平臺(tái)在整合Colossal-AI系統(tǒng)優(yōu)化和廉價(jià)算力的基礎(chǔ)上，近期發(fā)布了AI云主機(jī)的功能，方便用戶以近似裸機(jī)的方式進(jìn)行AI大模型的開(kāi)發(fā)和調(diào)試，并提供了多種使用方式，包括：Jupyter Notebook、ssh、服務(wù)本地端口映射和grafana監(jiān)控，全方位的為用戶提供便捷的開(kāi)發(fā)體驗(yàn)。

同時(shí)，還為用戶預(yù)制了含有ColossalAI代碼倉(cāng)庫(kù)和運(yùn)行環(huán)境的docker鏡像，用戶無(wú)需環(huán)境和資源配置，便可一鍵運(yùn)行ColossalAI代碼倉(cāng)庫(kù)中的代碼樣例。

Colossal-AI開(kāi)源地址：https://github.com/hpcaitech/ColossalAI