400万token上下文、推理再加速46%！最新开源方案升级MIT成果，推理成本再降低

明敏 发自 凹非寺
量子位 | 公众号 QbitAI

22倍加速还不够，再来提升46%，而且方法直接开源！

这就是开源社区改进MIT爆火项目StreamingLLM的最新成果。

StreamingLLM可以在不牺牲生成效果、推理速度的前提下，实现多轮对话共400万个token，22.2倍推理速度提升。

该项目在上线不到3个月时间内，GitHub项目标星达到5.7k star。

不过，StreamingLLM使用原生PyTorch实现，对于多轮对话推理场景落地应用的低成本、低延迟、高吞吐等需求仍有优化空间。

Colossal-AI团队开源了SwiftInfer，基于TensorRT的StreamingLLM，可以进一步提升大模型推理性能46%，有效解决如上问题。

具体如何实现？一起来看。

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

StreamingLLM如何实现超长多轮对话？

大语言模型能够记住的上下文长度，直接影响了ChatGPT等大模型应用与用户互动的质量。

如何让LLM在多轮对话场景下保持生成质量，对推理系统提出了更高的要求，因为LLM在预训练期间只能在有限的注意力窗口的限制下进行训练。

常见的KV Cache机制能够节约模型计算的时间，但是在多轮对话的情景下，key和value的缓存会消耗大量的内存，无法在有限的显存下无限扩展上下文。

同时，训练好的模型在不做二次微调的前提下也无法很好地泛化到比训练序列长度更长的文本，导致生成效果糟糕。

△图来源：https://arxiv.org/pdf/2309.17453.pdf

StreamingLLM为了解决了这个问题，通过观察了注意力模块中Softmax的输出，发现了attention sink的现象。

要知道，注意力机制会为每一个token分配一个注意力值，而文本最初的几个token总是会分配到很多无用的注意力。

当使用基于滑动窗口的注意力机制时，一旦这几个token被踢出了窗口，模型的生成效果就会迅速崩溃。

只要一直把这几个token保留在窗口内，模型就能稳定地生成出高质量的文本。

比起密集注意力（Dense Attention）、窗口注意力（Window Attention）以及带重计算的滑动窗口注意力(Sliding Window w/ Re-computing)，StreamingLLM基于attention sink的注意力机制无论是在计算复杂度还是生成效果上都表现优异。

在不需要重新训练模型的前提下，StreamingLLM能够直接兼容目前的主流大语言模型并改善推理性能。

SwiftInfer如何升级StreamingLLM？

Colossal-AI团队的方法是，将StreamingLLM方法与TensorRT推理优化结合。

这样的好处是，不仅能继承原始StreamingLLM的所有优点，还具有更高的运行效率。

使用TensorRT-LLM的API，还能够获得接近于PyTorch API的模型编写体验。

具体来看，基于TensorRT-LLM可以重新实现KV Cache机制以及带有位置偏移的注意力模块。

如下图所示，假设窗口大小为10个token，随着生成的token增加（由黄色方块表示），在KV缓存中将中间的token踢出，与此同时，始终保持着文本开始的几个token（由蓝色方块表示）。

由于黄色方块的位置会发生变化，在计算注意力时，也需要重新注入位置信息。

需要注意的是，StreamingLLM不会直接提高模型能访问的上下文窗口，而是能够在支持流式超多轮对话的同时保证模型的生成效果。

经过升级后，在原版本StreamingLLM基础上，Colossal-AI团队发布的SwiftInfer可带来额外的最多46%的推理吞吐速度提升。

能为大模型多轮对话推理提供低成本、低延迟、高吞吐的最佳实践。TensorRT-LLM团队也在同期对StreamingLLM进行了类似支持。

团队刚开源13B大模型

Colossal-AI目前已获得GitHub星数三万五千多颗，位列全球TOP400，细分赛道排名世界第一，可通过高效多维并行、异构内存等，降低AI大模型训练/微调/推理的开发与应用成本，提升模型任务表现，降低GPU需求。作为主流开源AI大模型系统社区，Colossal-AI生态在多方面保持活跃更新。

值得一提的是，Colossal-AI最近还开源了另一个新模型Colossal-LLaMA-2-13B。

它仅用25B token 数据和万元算力，效果远超基 LLaMA-2 的其他中文汉化模型。

即使与其他采用中文语料，可能花费上千万元成本，从头预训练的各大知名模型相比，Colossal-LLaMA-2在同规模下表现仍旧不错。

13B 版本通过构建更为完善的数据体系，在知识性内容掌握程度，自然语言处理任务理解程度，以及安全性，价值观等问题上，都有质的提升。

另外在云平台方面，Colossal-AI云平台在整合Colossal-AI系统优化和廉价算力的基础上，近期发布了AI云主机的功能。

它能方便用户以近似裸机的方式进行AI大模型的开发和调试，并提供了多种使用方式，包括：Jupyter Notebook、ssh、服务本地端口映射和grafana监控，全方位的为用户提供便捷的开发体验。

同时，还为用户预制了含有ColossalAI代码仓库和运行环境的docker镜像，用户无需环境和资源配置，便可一键运行ColossalAI代码仓库中的代码样例。

Colossal-AI开源地址：https://github.com/hpcaitech/ColossalAI

参考链接：https://hpc-ai.com/blog/Colossal-AI-SwiftInfer

— 完 —

点这里👇关注我，记得标星哦～

一键三连「分享」、「点赞」和「在看」

科技前沿进展日日相见 ~