mit-han-lab/streaming-llm

这是什么？

核心功能

• 01在 KV 缓存中保留开头的注意力汇聚令牌和最近令牌窗口。
• 02无需微调，即可让模型处理远超训练长度的输入。
• 03附带流式 Llama 聊天机器人示例脚本。
• 04提供困惑度（perplexity）评估代码。
• 05支持 Llama-2、MPT、Falcon 和 Pythia 模型。

适用场景

• 运行需要跨多轮对话持续响应的长时间聊天机器人或日常助手。
• 研究大模型中的注意力汇聚现象和 KV 缓存管理。
• 复现 StreamingLLM 论文中的困惑度实验结果。

快速上手

1. 1

   创建并激活 Python 3.8 的 conda 环境

   conda create -yn streaming python=3.8
2. 2

   激活该环境

   conda activate streaming
3. 3

   安装 PyTorch 及指定版本的 Python 依赖

   pip install torch torchvision torchaudio
4. 4

   在 GPU 0 上运行流式 Llama 聊天机器人演示

   CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py --enable_streaming

使用前须知

• !环境配置需要较旧的 Python 版本（3.8）以及固定的 transformers 版本（4.33.0）。
• !演示需要 GPU，并且需要自行获取 Llama-2 等预训练模型。
• !模型只关注近期 token 和注意力汇聚点，不会扩展上下文，也不提供长期记忆。
• !仓库最近一次推送在 2024 年 7 月，维护活跃度有限。

收录于专题

标签

常见问题

streaming-llm 是什么？

StreamingLLM 研究代码，通过保留注意力汇聚（attention sink）令牌和近期窗口，让预训练大模型无需微调即可处理无限长的文本流。 StreamingLLM 是 MIT Han Lab 发布的研究项目，对应 ICLR 2024 论文《Efficient Streaming Language Models with Attention Sinks》。它针对长时间多轮对话中的两个问题：解码时缓存所有历史 Key 和 Value 状态带来的显存压力，以及输入超过训练长度后模型输出质量下降。

streaming-llm 可以免费使用吗？用的什么许可证？

streaming-llm 以 MIT 许可证开源。用于商业产品前，请先确认许可证条款。

streaming-llm 还在维护吗？

超过半年未提交。streaming-llm 最近一次提交是在 2024年7月11日。

streaming-llm 适合用来做什么？

运行需要跨多轮对话持续响应的长时间聊天机器人或日常助手。 研究大模型中的注意力汇聚现象和 KV 缓存管理。 复现 StreamingLLM 论文中的困惑度实验结果。

streaming-llm 怎么安装和运行？

创建并激活 Python 3.8 的 conda 环境: conda create -yn streaming python=3.8 激活该环境: conda activate streaming 安装 PyTorch 及指定版本的 Python 依赖: pip install torch torchvision torchaudio 在 GPU 0 上运行流式 Llama 聊天机器人演示: CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py --enable_streaming

在哪里获取 streaming-llm？

源代码和安装说明在 GitHub 上：https://github.com/mit-han-lab/streaming-llm