來(lái)源：機(jī)器之心

圖片來(lái)源：由無(wú)界 AI生成

在大型語(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ì)開源了SwiftInfer，基于 TensorRT 實(shí)現(xiàn)了 StreamingLLM，可以進(jìn)一步提升大模型推理性能 46%，為多輪對(duì)話推理提供了高效可靠的落地方案。

開源地址：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)?LLM 在預(yù)訓(xùn)練期間只能在有限的注意力窗口的限制下進(jìn)行訓(xù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 的模型編寫體驗(yàn)。

基于 TensorRT-LLM，我們重新實(shí)現(xiàn)了 KV Cache 機(jī)制以及帶有位置偏移的注意力模塊。如下圖所示，假設(shè)我們的窗口大小為 10 個(gè) token，隨著生成的 token 增加（由黃色方塊表示），我們?cè)?KV 緩存中將中間的 token 踢出，與此同時(shí)，始終保持著文本開始的幾個(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) / 推理的開發(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ù)和萬(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 大模型的開發(fā)和調(diào)試，并提供了多種使用方式，包括：Jupyter Notebook、ssh、服務(wù)本地端口映射和 grafana 監(jiān)控，全方位的為用戶提供便捷的開發(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 開源地址：https://github.com/hpcaitech/ColossalAI

參考鏈接：

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