YaRN

此仓库包含 YaRN 上下文窗口扩展方法的代码和数据。

论文

论文（ICLR 2024）：YaRN：大型语言模型的高效上下文窗口扩展
旧预印本 (arXiv)

模型

LLaMA

我们发布了使用 YaRN 微调的 Llama 2 变体，其上下文窗口长度分别为 32K、64K 和 128K。这些模型在 🤗 Hugging Face 上以 Llama 2 许可证提供。

规模	上下文长度	链接
7B	64K	NousResearch/Yarn-Llama-2-7b-64k
7B	128K	NousResearch/Yarn-Llama-2-7b-128k
13B	64K	NousResearch/Yarn-Llama-2-13b-64k
13B	128K	NousResearch/Yarn-Llama-2-13b-128k
70B	32K	NousResearch/Yarn-Llama-2-70b-32k

此外，我们还发布了使用 NTK-aware 和 YaRN 微调的 Llama 2 7B 的 8K 上下文窗口版本（见会议论文中的表 1）。

Mistral

随着我们论文 v2 版本的发布，我们也发布了 Mistral 7B v0.1 的 64K 和 128K 变体。

规模	上下文长度	链接
7B	64K	NousResearch/Yarn-Mistral-7b-64k
7B	128K	NousResearch/Yarn-Mistral-7b-128k

SOLAR

SOLAR 10.7B v1.0 模型利用 深度增强缩放 技术，在 Mistral 7B v0.1 的基础上增加了层数，这可能在每参数的基础上提升长上下文性能。我们发布了 32K 和 64K 的变体。

规模	上下文长度	链接
10.7B	32K	NousResearch/Yarn-Solar-10b-32k
10.7B	64K	NousResearch/Yarn-Solar-10b-64k

复现

我们坚信开放科学，因此公开了所有用于复现论文结果的代码和数据。要复现，请克隆仓库并进行本地安装。

git clone https://github.com/jquesnelle/yarn
cd yarn
pip install -e .

训练

要训练模型，请运行 accelerate config 并启用 DeepSpeed 加速。训练时使用的配置文件为 deepspeed/zero3.json。

# ./train.sh

分词后的训练数据可在 🤗Hugging Face 上获取，该数据源自 pg19 数据集。对于 Mistral 模型，则使用了 Long-Data-Collections 数据集中预训练和微调分割的混合数据，分词后的数据集同样可在 🤗Hugging Face 上找到。

评估

要复现评估结果，请使用 pip install git+https://github.com/EleutherAI/lm-evaluation-harness 安装 lm-evaluation-harness，然后运行提供的两个脚本。

# ./eval.sh
# ./eval-harness.sh

引用

@inproceedings{
      peng2024yarn,
      title={Ya{RN}: Efficient Context Window Extension of Large Language Models},
      author={Bowen Peng and Jeffrey Quesnelle and Honglu Fan and Enrico Shippole},
      booktitle={The Twelfth International Conference on Learning Representations},
      year={2024},
      url={https://openreview.net/forum?id=wHBfxhZu1u}
}

YaRN 快速上手指南

YaRN 是一种高效的大语言模型上下文窗口扩展方法，支持将 LLaMA、Mistral 等模型的上下文长度扩展至 32K、64K 甚至 128K。本指南帮助开发者快速复现论文结果或使用已发布的微调模型。

环境准备

• 系统要求：Linux 或 macOS（推荐 Linux），具备 CUDA 支持的 GPU（用于训练或推理长上下文模型）
• 前置依赖：
  • Python 3.8+
  • git
  • pip
  • PyTorch（建议 2.0+，需与 CUDA 版本匹配）
  • Hugging Face Transformers
  • DeepSpeed（用于训练）
  • accelerate

💡 国内用户可使用清华或中科大镜像加速 Python 包安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package>

安装步骤

克隆仓库并进行本地可编辑安装：

git clone https://github.com/jquesnelle/yarn
cd yarn
pip install -e .

如需训练模型，还需配置 DeepSpeed：

accelerate config

在交互配置中启用 DeepSpeed，并选择 deepspeed/zero3.json 作为配置文件（该文件已包含在仓库中）。

基本使用

方式一：直接使用已发布模型（推荐）

通过 Hugging Face Transformers 直接加载 YaRN 微调后的模型，例如 7B 64K 版本的 LLaMA 2：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "NousResearch/Yarn-Llama-2-7b-64k"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

input_text = "请总结以下长文档的核心内容：" + "..." * 10000  # 模拟长输入
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=512)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

注意：确保你的显存足够加载模型并处理长上下文（64K 上下文可能需要 40GB+ 显存，建议使用多卡或量化方案）。

方式二：复现训练流程（高级用户）

1. 下载分词后的训练数据（以 LLaMA 为例）：

   • 数据集地址：https://huggingface.co/datasets/emozilla/pg_books-tokenized-bos-eos-chunked-65536

2. 运行训练脚本：

   ./train.sh
   

3. 评估模型性能（需先安装 lm-evaluation-harness）：

   pip install git+https://github.com/EleutherAI/lm-evaluation-harness
   ./eval.sh
   ./eval-harness.sh
   

国内用户若访问 Hugging Face 受阻，可使用镜像站如 https://hf-mirror.com 替代 huggingface.co，或在代码中设置环境变量：

export HF_ENDPOINT=https://hf-mirror.com