在大型语言模型（LLM）的世界中，处理多轮对话一直是一个挑战。前不久麻省理工Guangxuan Xiao等人推出的StreamingLLM，能够在不牺牲推理速度和生成效果的前提下，可实现多轮对话总共400万个token的流式输入，22.2倍的推理速度提升。
但StreamingLLM使用原生PyTorch实现，对于多轮对话推理场景落地应用的低成本、低延迟、高吞吐等需求仍有优化空间。
Colossal-AI团队开源了SwiftInfer，基于TensorRT实现了StreamingLLM，可以进一步提升大模型推理性能46%，为多轮对话推理提供了高效可靠的落地方案。
开源地址：
https://github.com/hpcaitech/SwiftInfer
https://huggingface.co/hpcai-tech
StreamingLLM简介
大语言模型能够记住的上下文长度，直接影响了ChatGPT等大模型应用与用户互动的质量。
如何让LLM在多轮对话场景下保持生成质量，对推理系统提出了更高的要求，因为LLM在预训练期间只能在有限的注意力窗口的限制下进行训练。
常见的KV Cache机制能够节约模型计算的时间，但是在多轮对话的情景下，key和value的缓存会消耗大量的内存，无法在有限的显存下无限扩展上下文。同时，训练好的模型在不做二次微调的前提下也无法很好地泛化到比训练序列长度更长的文本，导致生成效果糟糕。
StreamingLLM为了解决了这个问题，通过观察了注意力模块中Softmax的输出，发现了attention sink的现象。我们知道注意力机制会为每一个token分配一个注意力值，而文本最初的几个token总是会分配到很多无用的注意力。当我们使用基于滑动窗口的注意力机制时，一旦这几个token被踢出了窗口，模型的生成效果就会迅速崩溃。只要一直把这几个token保留在窗口内，模型就能稳定地生成出高质量的文本。
比起密集注意力（Dense Attention）、窗口注意力（Window Attention）以及带重计算的滑动窗口注意力(Sliding Window w/ Re-computing)，StreamingLLM基于attention sink的注意力机制无论是在计算复杂度还是生成效果上都表现优异。在不需要重新训练模型的前提下，StreamingLLM能够直接兼容目前的主流大语言模型并改善推理性能。
SwiftInfer：基于TensorRT的StearmingLLM实现
为了将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 版本通过构建更为完善的数据体系，在知识性内容掌握程度，自然语言处理任务理解程度，以及安全性，价值观等问题上，都有质的提升。
MoE
MoE模型在处理输入数据时只有少数专家模型被激活或者使用，因此可在训练和推理上表现出显著优势，达到节省计算资源的效果。GPT-4据传也采用MoE架构，Mistral AI、MiniMax等知名大模型创企最近都透露了MoE进展。Colossal-AI通过EZ-MoE等优化，可提升MoE模型训练效率9倍，并开源相关代码与模型。
本文由 Hugging Face 中文社区内容共建项目提供，稿件由社区成员投稿，经授权发布于 Hugging Face 公众号。文章内容不代表官方立场，文中介绍的产品和服务等均不构成投资建议。了解更多请关注公众号:
如果你有与开源 AI、Hugging Face 相关的技术和实践分享内容，以及最新的开源 AI 项目发布，希望通过我们分享给更多 AI 从业者和开发者们，请通过下面的链接投稿与我们取得联系:
相关文章
