LLM 也许 LongLM：无需微调即可自扩展 LLM 上下文窗口

在 LLM 也许 LongLM：无需微调即可自扩展 LLM 上下文窗口 中提出的 Self-Extend 的实现。

更新

• [2024年5月31日]:🎉 SelfExtend 在 Google I/O 演示中被重点介绍，地点为 YouTube，展示了 Gemma 的长上下文能力！！！

• [2024年5月1日]:🎉 SelfExtend 已被 ICML 2024 接受！维也纳见！

• [2024年4月19日]:💡 我们增加了对 LLama-3 和 transformers==4.40 的支持。若要使用 transformers==4.40，可将 Llama_4_40.py 文件名改为 Llama.py，以替换 现有补丁文件。

  • （尽管 Llama-3 官方模型库推荐使用 transformers==4.40，但我们发现使用 transformers==4.38.2 时，Llama-3 同样可以良好运行）

• [2024年4月6日]: 我们添加了使用 SelfExtend 进行超参数搜索的结果，您可在此查看 这里

• [2024年3月24日]: 我们添加了 Triton 实现的闪速自扩展功能。现在，您可以使用我们 Triton 实现的 FlashSelfExtend 来体验自扩展功能！

• • 我们添加了 SelfExtend 的 FlashAttention 实现，特别感谢 Qingquan Song！该实现使用了 Tri Dao 提供的原生 flash_attn。
  • 我们还尝试使用 Triton 为 SelfExtend 实现 FlashAttention。但目前，它仅适用于预填充阶段，无法在解码阶段良好工作。 我们正在积极调试中。欢迎提出任何建议！
  • 我们增加了对 Qwen1.5 的支持。更多详情请参阅 代码。
  • 我们重新整理了此仓库，并重构了几处文件。现在，所有代码均在 transformers==4.38.2 和 flash_attn==2.5.6 环境下运行。旧版代码/示例/README 已打包至 legacy_patch_before_4_38。我们建议使用我们的 Docker 镜像：hoytjin/selfextend_docker:v0.1，以避免任何环境问题。
  • 我们更新了 API，现在只需简单调用 SelfExtend.apply(加载的模型, 组大小, 窗口大小) 即可享受我们的自扩展功能！请查看并运行提供的 示例 以获取更多细节！
  • 我们新增了一个包含 32k 上下文长度的密码示例，以及一个更具挑战性的 10 位密码示例。
  • 欢迎加入我们的 Discord 讨论群！🔥🔥

• [2024年2月22日]: 我们添加了 Google 新型 LLM Gemma 的实现！！！欢迎大家试用和测试！！

• [2024年1月19日]: 我们增加了 基于 transformers 4.36.2 的 Llama 实现以及 基于 transformers 4.37 的 Microsoft 官方 phi-2 实现。还有一个好消息：闪速注意力版本将在几天内推出！💥

• [2024年1月11日]: 我们已测试了 phi-2 的实现。效果很好。您可在 Reddit 帖子 中找到相关结果，详细信息也可见 X 平台帖子。

• [2024年1月7日]: 添加 Mistral 的实现

• [2024年1月5日]: 我们的提案方法已在 Reddit 帖子 中讨论。

可能与 Self-Extend 无关的问题：

• Gemma-7b 必须以 bfloat16 格式加载。而 Gemma-2b 使用 float16 格式仍可正常运行。
• 如果使用 transformers 4.36，Llama 默认使用的注意力机制是 LlamaSpdaAttention，而非 LlamaSpdaAttention。请注意这一点，并确保将前向传播方法替换为正确的类。
• 使用 Self-Extend 时，应禁用 Mistral 的滑动窗口功能。为何不应使用 SWA 的原因可在我们的论文中找到。

第三方实现

Llama.cpp https://github.com/ggerganov/llama.cpp

Llama.cpp 对 Self-Extend 的实现和集成非常出色！不妨一试！😄

1. 概述

本研究旨在激发 LLM 在无需微调的情况下处理长上下文的能力。由于训练过程中序列长度受限，大型语言模型 (LLMs) 在推理阶段往往难以应对长输入序列。在本工作中，我们认为现有的 LLM 本身具备处理长上下文的内在能力。基于这一观点，我们提出通过自我扩展 LLM 的上下文窗口来充分发挥其内在潜力。我们提出了 Self-Extend 方法，以激发 LLM 处理长上下文的潜能。其基本思想是构建双层注意力信息：组级和邻域级。这两层均由模型原有的自注意力机制计算得出，因此该方法无需任何训练。

2. 如何使用 SelfExtend

2.1 设置

对于当前的 Llama 实现，所使用的 Python 包如下：

transformers==4.38.2
flash_attn==2.5.6 

我们建议使用此 Docker 镜像：hoytjin/selfextend_docker:v0.1

我们此前已为多个模型提供了补丁。您可查看 legacy_patch_before_4_38。其中包含旧版补丁（Llama、Mistral、Phi 等）及 README。

安装

将仓库克隆到您的机器上，并将您的模型文件复制到克隆后的目录中。

2.2 运行

import SelfExtend

# 加载您的模型，例如：loaded_model = AutoModelForCausalLM.from_pretrained(model_path) 

# 组大小、邻域窗口。 

SelfExtend.apply(loaded_model, 组大小, 窗口大小, enable_flash_attention=False)

# 推理，例如：loaded_model.generate(...)

默认情况下，enable_flash_attention=False。如果模型已启用 FlashAttention，则可将其设置为 True。

我们以密码检索为例说明如何使用 SelfExtend。您可查看 example.py：

python example.py

3. 如何选择 group_size 和 neighbor_window

以下建议基于我们的经验：

• 以 Llama-2 为基模型时，group_size 的合理范围是 2~64；neighbor_window 的可行范围是 512~1536。不过，在许多情况下，较大的 group_size 配合较小的 neighbor_window 也能取得不错的效果。

• 选择 group_size 和 neighbor_window 的一般原则是：确保输入序列长度不超过最大扩展窗口大小（对于 Llama-2 而言，即 (4096 - neighbor_window) * group_size + neighbor_window）。

• 我们在实验中并未对 group_size 进行细致调优。通常认为，对于同一序列，较小的 group 应该表现更好。然而，在某些实验中我们发现这一规律并不完全成立：

  有时，较大的 group_size 反而更有利。这可能是因为相对位置较大的部分在预训练阶段未得到充分训练，而较大的 group_size 能够利用那些训练更充分的较小位置来辅助上下文窗口的扩展。不过，较小的 group_size 往往具有更高的精度。因此，这里存在一个权衡。更多细节请参阅消融实验部分。
  
  例如：
  若 QA 任务的输入长度为 15,800，且 neighbor window 设置为 1,024，则 group_size 可设为 5，因为 5 * (4,096 - 1,024) + 1,024 = 16,384，大于 15,800。然而，将 group_size 设为 6，甚至更大如 8 或 16，可能会进一步提升模型性能。当 group_size 为 5 时，Self-Extend 会使用位置 1,025 至 3,979 来扩展上下文窗口；而若 group_size 设为 8，则 Self-Extend 仅使用位置 1,025 至 2,871 进行扩展。尽管 group_size 为 8 的精度不如 5，但 group_size 为 5 所使用的 2,872 至 3,979 位置在预训练过程中训练程度较低，这可能会影响扩展效果。

• 对于长度为 L 的序列，可以先尝试最小的 group_size [计算公式为：G * (L - w_n) + w_n]，然后再测试更大的 group 是否能带来更好的效果。

SelfExtend 在“大海捞针”任务上的表现

经验法则：

设预训练上下文窗口为 $L$，目标扩展长度为 $N$，邻域窗口为 $W$，分组大小为 $G$，则超参数选择的经验法则应满足以下不等式：$(\frac{1}{2} \sim \frac{2}{3}) \times L > W + \frac{N-W}{G}$。这一规则是基于经验得出的，我们认为其背后的原因在于：相对位置较大的部分在预训练中并未得到充分训练。因此，只有约一半到三分之二的位置训练得较好，Self-Extend 应当仅利用这些训练较好的相对位置来进行扩展。这一发现也解释了： 过小的 group_size 会导致性能下降，因为虽然它提供了精确的位置信息，却要求 Self-Extend 使用训练不足的相对位置进行扩展。 过大的 neighbor window 也会降低性能，因为它虽然提供了更多的邻近信息，但也需要依赖训练较差的相对位置来完成扩展。 实验结果表明，Self-Extend 对超参数的选择并不敏感。采用预先设定的启发式值作为 group_size 和 neighbor window 的取值，通常即可获得令人满意的效果。

[简要总结] Self-Extend 对超参数的选择并不敏感。可以通过一个代表性任务来确定合适的超参数，或者直接遵循我们的经验不等式：$(\frac{1}{2} \sim \frac{2}{3}) \times L > W + \frac{N-W}{G}$。

---

如果您认为我们的方法有用，请引用我们的论文：

@misc{jin2024llm,
      title={LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning}, 
      author={Hongye Jin and Xiaotian Han and Jingfeng Yang and Zhimeng Jiang and Zirui Liu and Chia-Yuan Chang and Huiyuan Chen and Xia Hu},
      year={2024},
      eprint={2401.01325},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

4. 贡献

我们欢迎研究社区为提升 Self-Extend 的效率做出贡献。如果您有任何想法或想报告 bug，请提交 issue 或 pull request。

5. 许可证

代码采用 MIT 许可证发布。

LongLM (Self-Extend) 快速上手指南

LongLM (Self-Extend) 是一种无需微调即可扩展大语言模型（LLM）上下文窗口的技术。它通过构建“组级别”和“邻居级别”的双层注意力机制，激发模型固有的长文本处理能力。

1. 环境准备

系统要求

• Python: 建议 Python 3.8+
• GPU: 支持 CUDA 的 NVIDIA GPU（推荐使用 FlashAttention 加速）
• Docker (推荐): 为避免环境依赖冲突，强烈建议使用官方提供的 Docker 镜像。
  • 镜像地址：hoytjin/selfextend_docker:v0.1

前置依赖

当前主要实现基于 Llama 系列模型，核心依赖版本如下：

• transformers==4.38.2
• flash_attn==2.5.6

注意：若使用其他版本（如 4.40 支持 Llama-3），需替换对应的补丁文件，详见项目更新日志。

2. 安装步骤

方案 A：使用 Docker（推荐）

直接拉取预配置好的环境镜像，避免手动安装复杂的 CUDA 和 FlashAttention 依赖。

docker pull hoytjin/selfextend_docker:v0.1
docker run -it --gpus all -v $(pwd):/workspace hoytjin/selfextend_docker:v0.1

方案 B：手动安装

克隆仓库并安装依赖：

# 1. 克隆仓库
git clone https://github.com/datamllab/LongLM.git
cd LongLM

# 2. 安装核心依赖 (确保 torch 已正确安装且匹配 CUDA 版本)
pip install transformers==4.38.2
pip install flash-attn==2.5.6 --no-build-isolation

# 3. (可选) 国内加速源
# pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers==4.38.2
# pip install -i https://pypi.tuna.tsinghua.edu.cn/simple flash-attn==2.5.6 --no-build-isolation

3. 基本使用

核心代码示例

只需导入 SelfExtend 模块，加载模型后调用 apply 方法即可启用长上下文扩展。

import SelfExtend
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 加载模型 (以 Llama-2 为例)
model_path = "meta-llama/Llama-2-7b-hf" # 替换为你的本地模型路径
loaded_model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 2. 应用 Self-Extend
# 参数说明:
# group_size: 组大小 (经验值 2~64)
# window_size: 邻居窗口大小 (经验值 512~1536)
# enable_flash_attention: 若模型加载时启用了 FlashAttention 则设为 True
group_size = 8
window_size = 1024
SelfExtend.apply(loaded_model, group_size, window_size, enable_flash_attention=False)

# 3. 进行长文本推理
input_text = "你的超长输入文本..." 
inputs = tokenizer(input_text, return_tensors="pt").to(loaded_model.device)

# 生成响应
outputs = loaded_model.generate(**inputs, max_new_tokens=512)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

运行官方示例

项目提供了基于 Passkey Retrieval 任务的完整示例脚本，可直接运行测试：

python example.py

参数选择建议

• group_size (G): 建议范围 2 ~ 64。较小的值精度更高，但可能利用到训练不足的位置；较大的值容错性更好。
• window_size (W): 建议范围 512 ~ 1536。
• 经验法则: 确保目标长度 $N$ 满足不等式 $(\frac{1}{2} \sim \frac{2}{3}) \times L > W + \frac{N-W}{G}$，其中 $L$ 为模型原始上下文长度（如 4096）。通常无需精细搜索，使用启发式默认值即可获得良好效果。