在大型语言模型（LLM）的世界中，处理多轮对话一直是一个挑战。前不久麻省理工Guangxuan Xiao等人推出的StreamingLLM，能够在不牺牲推理速度和生成效果的前提下，可实现多轮对话总共400万个token的流式输入，22.2倍的推理速度提升。
但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：基于TensorRT的StreamingLLM实现
为了将StreamingLLM这一技术更好的应用到落地场景，Colossal-AI团队成功地将StreamingLLM方法与TensorRT推理优化结合，不仅继承了原始StreamingLLM的所有优点，而且还具有更高的运行效率。使用TensorRT-LLM的API，我们还能够获得接近于PyTorch API的模型编写体验。
基于TensorRT-LLM，我们重新实现了KV Cache机制以及带有位置偏移的注意力模块。如下图所示，假设我们的窗口大小为10个token，随着生成的token增加（由黄色方块表示），我们在KV缓存中将中间的token踢出，与此同时，始终保持着文本开始的几个token（由蓝色方块表示）。由于黄色方块的位置会发生变化，在计算注意力时，我们也需要重新注入位置信息。
需要注意的是，StreamingLLM不会直接提高模型能访问的上下文窗口，而是能够在支持流式超多轮对话的同时保证模型的生成效果。
大模型无限输入流推理加速46%
原版本的StreamingLLM可以可靠地实现超过400万个token的流式输入，实现了比带重计算的滑动窗口注意力机制高出22.2倍的速度提升。
Colossal-AI团队发布的SwiftInfer可以进一步提升推理性能，最多带来额外的最多46%的推理吞吐速度提升，为大模型多轮对话推理提供低成本、低延迟、高吞吐的最佳实践。TensorRT-LLM团队也在同期对StreamingLLM进行了类似支持。
Colossal-AI社区动态
Colossal-AI目前已获得GitHub星数三万五千多颗，位列全球TOP400，细分赛道排名世界第一，可通过高效多维并行、异构内存等，降低AI大模型训练/微调/推理的开发与应用成本，提升模型任务表现，降低GPU需求。作为主流开源AI大模型系统社区，Colossal-AI生态在多方面保持活跃更新。
Colossal-LLaMA-2-13B开源
Colossal-LLaMA-2-13B模型，仅用25B token 数据和万元算力，效果远超基于 LLaMA-2 的其他中文汉化模型。即使与其他采用中文语料，可能花费上千万元成本，从头预训练的各大知名模型相比，Colossal-LLaMA-2在同规模下仍表现抢眼。13B 版本通过构建更为完善的数据体系，在知识性内容掌握程度，自然语言处理任务理解程度，以及安全性，价值观等问题上，都有质的提升。
Colossal-AI云平台
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
雷峰网(公众号：雷峰网)
雷峰网版权文章，未经授权禁止转载。详情见转载须知。