InfLLM

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.1
transformers >= 4.37.2
fschat >= 0.2.35
datasets >= 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 项目通过创新的记忆机制，展现了大语言模型在处理超长序列时的强大潜力，为相关领域的研究和应用提供了新的思路和方法。