LLM 可能是 LongLM:无需调优的自我扩展 LLM 上下文窗口
实现了 LLM 可能是 LongLM:无需调优的自我扩展 LLM 上下文窗口 中提出的自我扩展方法。
更新
-
[2024/05/31]:🎉 SelfExtend 在 YouTube 的 Google I/O 会议中被提到,用于展示 Gemma 的长上下文能力!
-
[2024/05/01]:🎉 SelfExtend 已被 ICML 2024 接收!维也纳见!
-
[2024/04/19]:💡 我们增加了对 transformers==4.40 的 LLama-3 的支持。要在 transformers==4.40 上使用,您可以将 Llama_4_40.py 文件重命名为
Llama.py
来替换 现有的补丁文件。- (尽管 Llama-3 的官方模型库建议使用 transformers==4.40,但我们发现使用 transformers==4.38.2,Llama-3 也可以正常工作)
-
[2024/04/06]:我们添加了一些使用 SelfExtend 的超参数搜索结果,您可以在这里查看。
-
[2024/03/24]:我们增加了 Triton 实现的闪存自我扩展。现在,您可以使用我们的 Triton 实现的 FlashSelfExtend 来享受自我扩展!
-
[2024/03/20]:我们做了许多更新:
- 我们增加了自我扩展的 FlashAttention 实现,感谢 Qingquan Song!这个实现使用了 Tri Dao 的原始 flash_attn。
- 我们也尝试用 Triton 实现了 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(loaded model, group size, window size)
来享受我们的自我扩展!查看并运行提供的 示例 了解更多详情! - 我们添加了一个 32k 上下文长度的新密钥示例,以及一个更具挑战性的 10 位密钥。
- 请加入我们的 Discord 进行讨论!🔥🔥
-
[2024/02/22]:我们添加了 Google 新 LLM Gemma 的实现!欢迎尝试和测试!
-
[2024/01/19]:我们添加了 Llama with transformers 4.36.2 的实现 和 Microsoft 官方 phi-2 with transformers 4.37 的实现。另一个好消息是:闪存注意力版本将在几天内发布!💥
-
[2024/01/11]:我们测试了 phi-2 的实现。它能正常工作。您可以在这个 Reddit 帖子 上找到一些结果,并在这个 X 帖子 上找到详细信息。
-
[2024/01/08]:添加了第三方实现部分。
-
[2024/01/07]:添加了 Mistral 的实现。
-
[2024/01/05]:我们提出的方法在这个 Reddit 帖子 中进行了讨论。
可能与 Self-Extend 无关的问题:
- Gemma-7b 必须以 bfloat16 加载。但 Gemma-2b 仍然可以在 float16 下正常工作。
- 如果使用 transformers 4.36,Llama 默认使用的注意力是
LlamaSpdaAttention
而不是LlamaSpdaAttention
。请注意这一点,并确保用正确的类替换 forward 方法。 - 要使用 Self-Extend,必须禁用 Mistral 的滑动窗口。我们在论文中解释了为什么不应该使用 SWA。
第三方实现
Llama.cpp https://github.com/ggerganov/llama.cpp
Llama.cpp 对自我扩展有很好的实现和集成!值得一试!😄
1. 概述
这项工作引出了 LLMs 处理长上下文的内在能力,而无需进行微调。训练期间有限的训练序列长度可能限制大型语言模型(LLMs)在推理时对长输入序列的应用。在这项工作中,我们认为现有的 LLMs 本身具有处理长上下文的内在能力。基于这一论点,我们建议通过 LLMs 自身来扩展其上下文窗口,以充分利用其内在能力。我们提出 Self-Extend 来激发 LLMs 处理长上下文的潜力。基本思想是构建双层注意力信息:组级和邻居级。这两个层次由原始模型的自注意力计算,这意味着该方法不需要任何训练。
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, group_size, window_size, enable_flash_attention=False)
# 推理,例如:loaded_model.generate(...)
enable_flash_attention 默认为 False,如果模型加载时启用了 FlashAttention,您可以将 enable_flash_attention 设置为 True。
我们使用 passkeyretrieval 作为示例来展示如何使用 self-extend。您可以查看 example.py:
python example.py
3. 如何选择 group_size 和 neighbor_window
以下想法基于我们的经验:
-
以 Llama-2 为基础模型,2~64 是合理的 group_size;512~1536 是可行的 neighbor_window。但在许多情况下,更大的 group_size 和更小的 neighbor_window 也很好。
-
选择 group_size 和 neighbor_window 的一般规则是:确保输入序列长度在最大扩展窗口