InfLLM 项目介绍
项目背景
现如今,面对长篇文本输入的应用程序,特别是由大模型(LLMs)驱动的智能代理,已经在现实世界中扮演起重要角色。然而,现有的大语言模型由于在训练时仅对特定长度的序列进行学习,因此在应对长序列时常出现无法泛化和注意力分散的问题。为了减轻这些问题,现有的方法通常采用滑动注意窗口,甚至丢弃远端的tokens来处理超长序列。然而,这样的方法往往不能有效捕捉序列中长距离的语义关系。
InfLLM 项目旨在通过一种基于记忆而不需要训练的新方法,揭示大语言模型在处理流动长序列时的内在能力。具体来说,InfLLM 将远端上下文存储在额外的记忆单元中,并采用一种高效的机制查找与当前token相关的记忆单元进行注意力计算。这样,InfLLM 使得大语言模型在保持长距离依赖捕捉能力的同时,能够高效地处理长序列。
项目亮点
-
训练无关的记忆能力:不需要额外训练,即可让大语言模型在处理长序列时表现优异,甚至当序列长度扩展到1,024K时仍能有效捕捉长距离依赖。
-
低硬件要求:通过改进代码,提升了推理速度并减少了GPU内存的使用,使得在标准硬件配置下就能运行。
-
支持最新的模型和技术:包括对 LLaMA 3 的支持,并利用了 faiss 的 topk 检索功能,进一步强化了性能。
使用方法
配置
InfLLM 使用 YAML 文件进行配置,用户可以在 config/
目录下找到示例配置文件。主要配置项包括模型路径、使用的注意力模式、记忆单元配置等。所有的配置选项都可以调整,以适应具体的应用场景。
运行评估
评估过程中,项目提供了 InfiniteBench 和 LongBench 两个评估基准。用户可以运行项目提供的脚本下载所需的数据集,并通过简单的命令在单个或多个 GPU 上进行评估。
bash scripts/download.sh
bash scripts/[infinitebench,longbench].sh
聊天机器人集成
InfLLM 集成了 fastchat 的 CLI 接口,用户可以通过命令行启动聊天机器人,与模型进行交互。
python -m inf_llm.chat \
--model-path mistralai/Mistral-7B-Instruct-v0.2 \
--inf-llm-config-path config/mistral-inf-llm.yaml
依赖环境
项目对底层环境的要求如下:
torch
>= 1.13.1transformers
>= 4.37.2fschat
>= 0.2.35datasets
>= 2.17.0- 其他依赖详见项目代码
项目引用
如果您觉得 InfLLM 项目有用,欢迎在相关研究中引用我们:
@article{xiao2024infllm,
author = {Chaojun Xiao and Pengle Zhang and Xu Han and Guangxuan Xiao and Yankai Lin and Zhengyan Zhang and Zhiyuan Liu and Song Han and Maosong Sun},
title = {InfLLM: Unveiling the Intrinsic Capacity of LLMs for Understanding
Extremely Long Sequences with Training-Free Memory},
journal = {arXiv},
year = {2024}
}
InfLLM 项目通过创新的记忆机制,展现了大语言模型在处理超长序列时的强大潜力,为相关领域的研究和应用提供了新的思路和方法。