多輪對話推理速度提升46%，開源方案打破LLM多輪對話的長度限制

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

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

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

開源地址：https://github.com/hpcaitech/SwiftInfer

StreamingLLM簡介

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

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

常見的KV Cache機制能夠節(jié)約模型計算的時間，但是在多輪對話的情景下，key和value的緩存會消耗大量的內(nèi)存，無法在有限的顯存下無限擴展上下文。同時，訓(xùn)練好的模型在不做二次微調(diào)的前提下也無法很好地泛化到比訓(xùn)練序列長度更長的文本，導(dǎo)致生成效果糟糕。

圖來源：https://arxiv.org/pdf/2309.17453.pdf

StreamingLLM為了解決了這個問題，通過觀察了注意力模塊中Softmax的輸出，發(fā)現(xiàn)了attention sink的現(xiàn)象。我們知道注意力機制會為每一個token分配一個注意力值，而文本最初的幾個token總是會分配到很多無用的注意力。當(dāng)我們使用基于滑動窗口的注意力機制時，一旦這幾個token被踢出了窗口，模型的生成效果就會迅速崩潰。只要一直把這幾個token保留在窗口內(nèi)，模型就能穩(wěn)定地生成出高質(zhì)量的文本。

比起密集注意力（Dense Attention）、窗口注意力（Window Attention）以及帶重計算的滑動窗口注意力(Sliding Window w/ Re-computing)，StreamingLLM基于attention sink的注意力機制無論是在計算復(fù)雜度還是生成效果上都表現(xiàn)優(yōu)異。在不需要重新訓(xùn)練模型的前提下，StreamingLLM能夠直接兼容目前的主流大語言模型并改善推理性能。

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

為了將StreamingLLM這一技術(shù)更好的應(yīng)用到落地場景，Colossal-AI團隊成功地將StreamingLLM方法與TensorRT推理優(yōu)化結(jié)合，不僅繼承了原始StreamingLLM的所有優(yōu)點，而且還具有更高的運行效率。使用TensorRT-LLM的API，我們還能夠獲得接近于PyTorch API的模型編寫體驗。

基于TensorRT-LLM，我們重新實現(xiàn)了KV Cache機制以及帶有位置偏移的注意力模塊。如下圖所示，假設(shè)我們的窗口大小為10個token，隨著生成的token增加（由黃色方塊表示），我們在KV緩存中將中間的token踢出，與此同時，始終保持著文本開始的幾個token（由藍(lán)色方塊表示）。由于黃色方塊的位置會發(fā)生變化，在計算注意力時，我們也需要重新注入位置信息。

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

大模型無限輸入流推理加速46%

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

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

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

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

Colossal-LLaMA-2-13B開源

Colossal-LLaMA-2-13B模型，僅用25B token 數(shù)據(jù)和萬元算力，效果遠(yuǎn)超基于 LLaMA-2 的其他中文漢化模型。即使與其他采用中文語料，可能花費上千萬元成本，從頭預(yù)訓(xùn)練的各大知名模型相比，Colossal-LLaMA-2在同規(guī)模下仍表現(xiàn)搶眼。13B 版本通過構(gòu)建更為完善的數(shù)據(jù)體系，在知識性內(nèi)容掌握程度，自然語言處理任務(wù)理解程度，以及安全性，價值觀等問題上，都有質(zhì)的提升。

Colossal-AI云平臺

Colossal-AI云平臺在整合Colossal-AI系統(tǒng)優(yōu)化和廉價算力的基礎(chǔ)上，近期發(fā)布了AI云主機的功能，方便用戶以近似裸機的方式進(jìn)行AI大模型的開發(fā)和調(diào)試，并提供了多種使用方式，包括：Jupyter Notebook、ssh、服務(wù)本地端口映射和grafana監(jiān)控，全方位的為用戶提供便捷的開發(fā)體驗。同時，還為用戶預(yù)制了含有ColossalAI代碼倉庫和運行環(huán)境的docker鏡像，用戶無需環(huán)境和資源配置，便可一鍵運行ColossalAI代碼倉庫中的代碼樣例。

Colossal-AI開源地址：

https://github.com/hpcaitech/ColossalAI

參考鏈接：

https://hpc-ai.com/blog/Colossal-AI-SwiftInfer

雷峰網(wǎng)(公眾號：雷峰網(wǎng))