long-context

用途

Long Context 技能用于扩展 Transformer 模型的上下文窗口，使模型能够处理 32k、64k、128k+ token 的长文档。它适用于以下场景：

使用 Transformer 模型处理长文档（32k、64k、128k+ token）
扩展预训练模型（如 LLaMA、Mistral 等）的上下文窗口
实现高效的位置编码（RoPE、ALiBi）
训练具备长度外推能力的模型
部署能高效处理可变长度输入的模型
以最小算力微调现有模型以支持更长上下文

主要能力

技能涵盖四种关键技术：

RoPE（旋转位置编码）：通过旋转矩阵编码绝对位置，在注意力中提供相对位置依赖，支持长度外推。
YaRN（Yet another RoPE extensioN）：采用 NTK 感知插值和注意力温度缩放，实现高效上下文扩展（相比基线节省 10 倍 token），可将 LLaMA 扩展至 128k token。
ALiBi（带线性偏置的注意力）：不向 token 添加位置嵌入，而是直接对注意力分数施加距离惩罚，训练速度比正弦嵌入快 11%，内存占用少 11%，具备强长度外推能力。
位置插值（Position Interpolation）：线性缩小位置索引，在已训练范围内插值而非外推，仅需极少微调（约 1000 步）即可将 LLaMA 7B-65B 扩展至 32k token。

技能还提供方法对比表，涵盖最大上下文、所需训练、内存、外推能力和最佳适用场景。

使用方式

安装

# HuggingFace Transformers（包含 RoPE、YaRN 支持）
pip install transformers torch

# 自定义实现
pip install einops  # 张量操作
pip install rotary-embedding-torch  # 独立 RoPE

# 可选：FlashAttention 提升效率
pip install flash-attn --no-build-isolation

快速开始

技能提供 RoPE、ALiBi 的完整 PyTorch 实现代码，包括 RotaryEmbedding 类、apply_rotary_pos_emb 函数、get_alibi_slopes 和 create_alibi_bias 函数。

HuggingFace 集成

通过修改模型配置启用 RoPE 缩放：

config.rope_scaling = {
    "type": "yarn",
    "factor": 8.0,
    "original_max_position_embeddings": 8192,
    "attention_factor": 1.0
}

也支持 linear、dynamic 等缩放类型。

微调

位置插值微调：修改配置后仅需约 1000 步微调。
YaRN 微调：克隆 YaRN 仓库后运行训练脚本，指定 scale、max_length、steps 等参数。

生产部署

技能提供长上下文推理示例、内存优化方法（梯度检查点、Flash Attention 2、vLLM 分页注意力）。

适用场景

需要处理超长文档的 NLP 任务
将现有短上下文模型扩展为长上下文模型
从零训练具备长度外推能力的新模型
在有限算力下快速扩展模型上下文

限制

位置插值若不经微调直接应用，模型表现会很差。
过于激进的缩放（如直接扩展到 100 万 token）需要大量微调，否则无法工作。
建议采用增量缩放（8k → 16k → 32k → 64k），每一步都进行微调。
微调数据应使用与目标长度匹配的长文档，短文档无法教会模型长上下文能力。
缩放因子越大，所需微调步数越多（粗略估计为 100 × 缩放因子）。

参考资源

RoPE 论文：https://arxiv.org/abs/2104.09864
YaRN 论文：https://arxiv.org/abs/2309.00071
ALiBi 论文：https://arxiv.org/abs/2108.12409
位置插值论文：https://arxiv.org/abs/2306.15595
HuggingFace RoPE 工具：https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py
YaRN 实现：https://github.com/jquesnelle/yarn
Together AI 博客：https://www.together.ai/blog/llama-2-7b-32k