LongLLaMA：面向上下文扩展的聚焦Transformer训练

简述

本仓库包含 LongLLaMA 的研究预览版，这是一款能够处理长达 256k 个标记甚至更多上下文的大型语言模型。

LongLLaMA 基于 OpenLLaMA 构建，并采用 聚焦Transformer (FoT) 方法进行微调。 LongLLaMA Code 则基于 Code Llama 构建。 我们以宽松的许可证（Apache 2.0）发布了一个较小的 3B 基础版本（未进行指令调优）的 LongLLaMA 模型，并在 Hugging Face 上提供了支持更长上下文的推理代码。我们的模型权重可以作为现有实现中 LLaMA 的直接替代品（适用于短上下文，最多 2048 个标记）。此外，我们还提供了评估结果以及与原始 OpenLLaMA 模型的对比。

除此之外，我们还发布了用于 指令调优（PyTorch） 和 FoT 持续预训练（JAX） 的代码。

概述

基础模型

聚焦Transformer：面向上下文扩展的对比学习（FoT）提出了一种简单的方法，使语言模型能够在训练时使用显著较短的输入，却具备处理可能包含数百万个标记的上下文的能力。FoT 允许一部分注意力层访问由 (key, value) 对组成的内存缓存，从而扩展上下文长度。FoT 的独特之处在于其训练过程，借鉴了对比学习的思想。具体而言，我们会有意将记忆注意力层暴露于相关和不相关的 key 中（类似于来自无关文档的负样本）。这一策略促使模型区分与语义多样化的 value 相关的 key，从而增强其结构。这样一来，就能将有效上下文长度外推到远超训练时所见的范围。

LongLLaMA 是一款基于 OpenLLaMA 并采用 FoT 方法微调的模型，其中三层层被用于上下文扩展。关键在于，LongLLaMA 能够将上下文长度外推至远超训练时所见的水平：8k。例如，在密钥检索任务中，它能够处理长度达 256k 的输入。 LongLLaMA Code 是一款基于 Code Llama 并采用 FoT 方法微调的模型。

FoT 持续预训练

在 fot_continued_pretraining 子文件夹中，我们提供了可用于使用 FoT 调整 LLaMA 模型的代码。 该代码使用 JAX 和 Flax 编写，并基于 EasyLM。

指令/聊天调优

在 instruction_fine_tuning 子文件夹中，我们提供了用于创建 LongLLaMA-指令-3Bv1.1 的代码，它是 LongLLaMA-3Bv1.1 的指令调优版本。我们在调优过程中使用了 OpenOrca（指令）和 zetavg/ShareGPT-Processed（聊天）数据集。 该代码利用了 PyTorch 和 Hugging Face 训练器。

推理代码

在 src 子文件夹中，我们提供了 FoT 模型的推理代码。
该代码使用 PyTorch 编写，并基于 Hugging Face 的 LLaMA 实现。
代码应支持标准的 Hugging Face API。更多详情请参阅 使用说明 部分。

使用说明

另请参阅：

• 带有 LongLLaMA-Instruct-3Bv1.1 的 Colab 笔记本。
• 带有基础 LongLLaMA 示例用法的 Colab 笔记本。

需求

pip install --upgrade pip
pip install transformers==4.33.2 sentencepiece accelerate

加载模型

import torch
from transformers import LlamaTokenizer, AutoModelForCausalLM

tokenizer = LlamaTokenizer.from_pretrained("syzymon/long_llama_3b_v1_1")
model = AutoModelForCausalLM.from_pretrained("syzymon/long_llama_3b_v1_1", 
                                            torch_dtype=torch.float32, 
                                            trust_remote_code=True)

输入处理与生成

LongLLaMA 使用 Hugging Face 接口，输入给模型的长文本将被分割成上下文窗口并加载到内存缓存中。

prompt = "My name is Julien and I like to"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
outputs = model(input_ids=input_ids)

在调用模型时，可以提供参数 last_context_length（默认值为 $1024$），用于指定最后一个上下文窗口中保留的标记数。调整此参数可以改善生成效果，因为模型的前几层无法访问内存缓存。详细信息请参阅 LongLLaMA 如何处理长输入。

generation_output = model.generate(
    input_ids=input_ids,
    max_new_tokens=256,
    num_beams=1,
    last_context_length=1792,
    do_sample=True,
    temperature=1.0,
)
print(tokenizer.decode(generation_output[0]))

其他配置

LongLLaMA 还有若干其他参数：

• mem_layers 用于指定具有记忆功能的层（应为空列表或检查点描述中列出的所有记忆层的列表）。
• mem_dtype 允许更改内存缓存的数据类型。
• mem_attention_grouping 可以在速度和减少内存使用之间进行权衡。
  当设置为 (4, 2048) 时，记忆层每次最多可同时处理 $4 \times 2048$ 个查询（每组 4 个注意力头，每个头处理 2048 个查询）。

import torch
from transformers import LlamaTokenizer, AutoModelForCausalLM

tokenizer = LlamaTokenizer.from_pretrained("syzymon/long_llama_3b_v1_1")
model = AutoModelForCausalLM.from_pretrained(
    "syzymon/long_llama_3b_v1_1", torch_dtype=torch.float32, 
    mem_layers=[], 
    mem_dtype='bfloat16',
    trust_remote_code=True,
    mem_attention_grouping=(4, 2048),
)

作为 LLaMA 代码的直接替换

LongLLaMA 检查点也可以作为 Hugging Face 的 LLaMA 实现 中 LLaMA 检查点的直接替换使用，但在这种情况下，它们将被限制为原始的 $2048$ 标记上下文长度。

from transformers import LlamaTokenizer, LlamaForCausalLM
import torch

tokenizer = LlamaTokenizer.from_pretrained("syzymon/long_llama_3b_v1_1")
model = LlamaForCausalLM.from_pretrained("syzymon/long_llama_3b_v1_1", torch_dtype=torch.float32)

LongLLaMA 如何处理长输入

超过 $lctx=2048$（对于 LongLLaMA Code 为 $lctx=4096$）标记的输入会自动被分割成窗口 $w_1, \ldots, w_m$。前 $m-2$ 个窗口每个包含 $lctx$ 个标记，$w_{m-1}$ 包含不超过 $lctx$ 个标记，而 $w_m$ 则包含由 last_context_length 指定的标记数。模型会逐个处理这些窗口，并在每次处理后扩展内存缓存。如果 use_cache 设置为 True，则最后一个窗口不会被加载到内存缓存中，而是会被加载到本地（生成）缓存中。

内存缓存为指定的记忆层中的每个注意力头存储 $(key, value)$ 对。此外，它还存储注意力掩码。

如果 use_cache=True（即在生成过程中），LongLLaMA 将使用两种缓存：用于指定层的记忆缓存，以及用于所有层的本地（生成）缓存。当本地缓存中的元素数量超过 $lctx$ 时，其内容会被移动到记忆层的内存缓存中。

为简化起见，本仓库中通过内存缓存和全注意力机制实现了上下文扩展。未来可以通过像 Faiss 这样的系统，用外部数据库中的 KNN 搜索来替代这一简单机制，从而进一步扩展上下文长度。我们将此作为未来的工作方向。

LongLLaMA 性能

我们展示了一些 LongLLaMA 结果的示例。更多细节请参阅我们的论文 Focused Transformer: Contrastive Training for Context Scaling。

我们在来自 Landmark Attention: Random-Access Infinite Context Length for Transformers 的密码检索任务上取得了良好的效果。生成提示和运行模型的代码位于 examples/passkey.py 中。

我们的 LongLLaMA 3B 模型在两个下游任务——TREC 问题分类和 WebQS 问答——上也表现出使用长上下文时的改进效果。

LongLLaMA 在不需要长上下文的任务上也能保持性能。

特别是，经过代码微调后的 LongLLaMA-Code 7B 在推理能力（GSM8K）和知识储备（MMLU）方面都有所提升：

我们在零样本设置下，使用 lm-evaluation-harness 对比了 OpenLLaMA 和 LongLLaMA 的表现。

从v1.1版本的模型开始，我们决定使用EleutherAI对lm-evaluation-harness的实现，并做了一点小修改：在输入序列的开头添加了<bos>标记。结果如下表所示。

我们还提供了人工评估的结果。我们在以下任一标记后截断生成的文本：

• "\ndef "
• "\nclass "
• "\nif __name__"

作者

• Szymon Tworkowski
• Konrad Staniszewski
• Mikołaj Pacek
• Henryk Michalewski
• Yuhuai Wu
• Piotr Miłoś

引用

若要引用本工作，请使用以下 BibTeX 格式：

@misc{tworkowski2023focused,
      title={Focused Transformer: Contrastive Training for Context Scaling}, 
      author={Szymon Tworkowski and Konrad Staniszewski and Mikołaj Pacek and Yuhuai Wu and Henryk Michalewski and Piotr Miłoś},
      year={2023},
      eprint={2307.03170},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

许可证

源代码和基础 LongLLaMA 3B 模型检查点采用 Apache License, Version 2.0 许可证授权。
经过指令/聊天微调的模型仅用于研究目的。
关于 LongLLaMA-Code 7B，请参阅 codellama/CodeLlama-7b-hf 的许可证。
LongLLaMA-Code 7B Instruct 是在 TIGER-Lab/MathInstruct、OpenOrca 和 ShareGPT-Processed 数据集上对 LongLLaMA-Code 7B 进行微调得到的。
部分示例使用了外部代码（请参阅文件头部的版权声明和许可证信息）。

致谢

我们衷心感谢 TPU Research Cloud 计划，该计划通过提供重要的计算资源，对我们的研究起到了关键作用。同时，我们也感谢 Xinyang Geng 和 Hao Liu 发布了 OpenLLaMA 检查点以及 EasyLM 库。

特别感谢 Keiran Paster，他为 LongLLaMA-Code 的预训练数据提供了极其宝贵的建议。

此外，我们还要感谢 Xiaosong,He，他为我们改进跨批次代码的解释提出了许多有益的建议。

LongLLaMA 快速上手指南

LongLLaMA 是一个基于 OpenLLaMA 和 Code Llama 构建的大语言模型，通过 Focused Transformer (FoT) 技术进行微调，能够处理长达 256k tokens 甚至更长的上下文。它支持作为现有 LLaMA 实现的直接替代品（在短上下文下），并提供了专门的推理代码以支持长文本输入。

环境准备

系统要求

• Python: 建议 Python 3.8+
• GPU: 推荐使用支持 CUDA 的 NVIDIA GPU（显存需求视模型大小而定，3B 版本约需 6-8GB，7B 版本需更高显存或使用量化/Offload 技术）
• 操作系统: Linux / macOS / Windows (WSL)

前置依赖

请确保已安装 pip 并更新至最新版本。本项目主要依赖 transformers, sentencepiece 和 accelerate。

国内加速提示：如果下载 Hugging Face 模型或依赖包速度较慢，建议配置国内镜像源。

• 设置 pip 镜像：pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
• 设置 HF 镜像：在代码中设置环境变量 export HF_ENDPOINT=https://hf-mirror.com

安装步骤

使用 pip 安装必要的 Python 库。注意 transformers 版本需为 4.33.2 以确保兼容性。

pip install --upgrade pip
pip install transformers==4.33.2 sentencepiece accelerate

基本使用

以下示例展示如何加载 LongLLaMA-3Bv1.1 模型并进行文本生成。该模型会自动将长输入分割为上下文窗口并加载到内存缓存中以支持长文本处理。

1. 加载模型与分词器

import torch
from transformers import LlamaTokenizer, AutoModelForCausalLM

# 加载分词器和模型
# trust_remote_code=True 是必须的，因为使用了自定义的 FoT 实现
tokenizer = LlamaTokenizer.from_pretrained("syzymon/long_llama_3b_v1_1")
model = AutoModelForCausalLM.from_pretrained(
    "syzymon/long_llama_3b_v1_1", 
    torch_dtype=torch.float32, 
    trust_remote_code=True
)

# 如果有 GPU，将模型移至 GPU
if torch.cuda.is_available():
    model = model.to("cuda")

2. 文本生成示例

LongLLaMA 支持通过 last_context_length 参数优化生成长度。该参数指定最后一个上下文窗口保留的 token 数量，适当调整可提升生成质量。

prompt = "My name is Julien and I like to"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

if torch.cuda.is_available():
    input_ids = input_ids.to("cuda")

# 执行生成
generation_output = model.generate(
    input_ids=input_ids,
    max_new_tokens=256,
    num_beams=1,
    last_context_length=1792,  # 调整此参数以优化长上下文表现
    do_sample=True,
    temperature=1.0,
)

# 解码并输出结果
print(tokenizer.decode(generation_output[0]))

3. 高级配置（可选）

如需优化显存占用或指定特定的记忆层，可在加载模型时传入额外参数：

model = AutoModelForCausalLM.from_pretrained(
    "syzymon/long_llama_3b_v1_1", 
    torch_dtype=torch.float32, 
    trust_remote_code=True,
    mem_layers=[],             # 指定带有记忆功能的层，空列表表示使用默认配置
    mem_dtype='bfloat16',      # 更改记忆缓存的数据类型以节省显存
    mem_attention_grouping=(4, 2048), # 权衡速度与显存：(头数, 每头查询数)
)

注意：如果仅用于短上下文（< 2048 tokens），可以直接使用标准的 LlamaForCausalLM 加载，此时无需 trust_remote_code=True，但将无法利用长上下文扩展能力。