长上下文扩展Skill long-context

---

名称: 长上下文 描述: 使用RoPE、YaRN、ALiBi和位置插值技术扩展transformer模型的上下文窗口。在处理长文档（32k-128k+令牌）、扩展预训练模型超出原始上下文限制或实现高效位置编码时使用。涵盖旋转嵌入、注意力偏置、插值方法和LLM的外推策略。 版本: 1.0.0 作者: Orchestra Research 许可证: MIT 标签: [新兴技术, 长上下文, RoPE, YaRN, ALiBi, 位置插值, 扩展上下文, 旋转嵌入, 注意力偏置, 上下文扩展, 位置编码] 依赖项: [transformers, torch, flash-attn]

长上下文：扩展Transformer上下文窗口

何时使用此技能

当您需要时使用长上下文技术：

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

关键技术: RoPE（旋转位置嵌入）、YaRN、ALiBi（带线性偏置的注意力）、位置插值

论文: RoFormer (arXiv 2104.09864), YaRN (arXiv 2309.00071), ALiBi (arXiv 2108.12409), 位置插值 (arXiv 2306.15595)

安装

# 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（旋转位置嵌入）

import torch
import torch.nn as nn

class RotaryEmbedding(nn.Module):
    """旋转位置嵌入（RoPE）。"""

    def __init__(self, dim, max_seq_len=8192, base=10000):
        super().__init__()
        # 计算逆频率
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len = max_seq_len

    def forward(self, seq_len, device):
        # 位置索引
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)

        # 计算频率
        freqs = torch.outer(t, self.inv_freq)  # (seq_len, dim/2)

        # 计算sin和cos
        emb = torch.cat((freqs, freqs), dim=-1)  # (seq_len, dim)
        return emb.cos(), emb.sin()

def rotate_half(x):
    """旋转一半的隐藏维度。"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    """将旋转嵌入应用到查询和键。"""
    # q, k 形状: (batch, heads, seq_len, dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

# 使用
rope = RotaryEmbedding(dim=64, max_seq_len=8192)
cos, sin = rope(seq_len=2048, device='cuda')

# 在注意力层中
q_rotated, k_rotated = apply_rotary_pos_emb(query, key, cos, sin)

ALiBi（带线性偏置的注意力）

def get_alibi_slopes(num_heads):
    """获取每个注意力头的ALiBi斜率值。"""
    def get_slopes_power_of_2(n):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        ratio = start
        return [start * (ratio ** i) for i in range(n)]

    if math.log2(num_heads).is_integer():
        return get_slopes_power_of_2(num_heads)
    else:
        # 最接近的2的幂
        closest_power = 2 ** math.floor(math.log2(num_heads))
        slopes = get_slopes_power_of_2(closest_power)
        # 添加额外斜率
        extra = get_slopes_power_of_2(2 * closest_power)
        slopes.extend(extra[0::2][:num_heads - closest_power])
        return slopes

def create_alibi_bias(seq_len, num_heads):
    """创建ALiBi注意力偏置。"""
    # 距离矩阵
    context_position = torch.arange(seq_len)
    memory_position = torch.arange(seq_len)
    relative_position = memory_position[None, :] - context_position[:, None]

    # 获取斜率
    slopes = torch.tensor(get_alibi_slopes(num_heads))

    # 将斜率应用到距离
    alibi = slopes[:, None, None] * relative_position[None, :, :]
    return alibi  # (num_heads, seq_len, seq_len)

# 在注意力中使用
num_heads = 8
seq_len = 2048
alibi_bias = create_alibi_bias(seq_len, num_heads).to('cuda')

# 将偏置添加到注意力分数
# attn_scores 形状: (batch, num_heads, seq_len, seq_len)
attn_scores = attn_scores + alibi_bias
attn_weights = torch.softmax(attn_scores, dim=-1)

LLaMA的位置插值

from transformers import LlamaForCausalLM, LlamaTokenizer

# 原始上下文：2048令牌
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 使用位置插值扩展到32k
# 修改RoPE基础频率
model.config.rope_scaling = {
    "type": "linear",
    "factor": 16.0  # 2048 * 16 = 32768
}

# 或使用动态缩放
model.config.rope_scaling = {
    "type": "dynamic",
    "factor": 16.0
}

# 用长文档微调（需要最小步骤）
# 在此配置更改后，位置插值即可工作

核心概念

1. RoPE（旋转位置嵌入）

工作原理：

• 通过旋转矩阵编码绝对位置
• 在注意力中提供相对位置依赖
• 支持长度外推

数学公式：

q_m = (W_q * x_m) * e^(imθ)
k_n = (W_k * x_n) * e^(inθ)

其中 θ_j = base^(-2j/d) for j ∈ [0, d/2)

优点：

• 距离越远，令牌间依赖性衰减
• 兼容线性注意力
• 比绝对位置编码有更好的外推能力

2. YaRN（又一个RoPE扩展）

关键创新：

• NTK感知插值（神经切线核）
• 注意力温度缩放
• 高效上下文扩展（比基线少10×令牌）

参数：

# YaRN配置
yarn_config = {
    "scale": 16,                    # 扩展因子
    "original_max_position": 2048,  # 基础上下文
    "extrapolation_factor": 1.0,    # NTK参数
    "attn_factor": 1.0,             # 注意力缩放
    "beta_fast": 32,                # 高频缩放
    "beta_slow": 1,                 # 低频缩放
}

性能：

• 将LLaMA扩展到128k令牌
• 训练步骤比基线少2.5×
• 最先进的上下文窗口扩展

3. ALiBi（带线性偏置的注意力）

核心思想：

• 不向令牌添加位置嵌入
• 直接将距离惩罚应用到注意力分数
• 偏置与键-查询距离成比例

公式：

attention_bias[i, j] = -m * |i - j|

其中 m = 每个注意力头的斜率

优点：

• 训练比正弦嵌入快11%
• 内存使用少11%
• 强大的长度外推（训练1k，测试2k+）
• 偏向近期性的归纳偏置

4. 位置插值

技术：

• 线性下调位置索引
• 在训练范围内插值（而非外推超出）
• 需要最小微调

公式：

# 原始：位置索引 [0, 1, 2, ..., L]
# 扩展：位置索引 [0, 0.5, 1.0, ..., L/2]
# （对于2×扩展）

scaled_position[i] = i / extension_factor

结果：

• LLaMA 7B-65B扩展到32k令牌
• 1000微调步骤足够
• 稳定性比外推好600×

方法比较

方法	最大上下文	需要训练	内存	外推能力	最佳用于
RoPE	8k-32k	完整预训练	中等	良好	新模型
YaRN	32k-128k	最小（10×高效）	中等	优秀	扩展现有模型
ALiBi	无限	完整预训练	低（-11%）	优秀	从零开始训练
位置插值	32k+	最小（1k步骤）	中等	差（设计如此）	快速扩展

实现模式

HuggingFace Transformers集成

from transformers import AutoModelForCausalLM, AutoConfig

# 带YaRN缩放的RoPE
config = AutoConfig.from_pretrained("mistralai/Mistral-7B-v0.1")
config.rope_scaling = {
    "type": "yarn",
    "factor": 8.0,
    "original_max_position_embeddings": 8192,
    "attention_factor": 1.0
}

model = AutoModelForCausalLM.from_config(config)

# 位置插值（更简单）
config.rope_scaling = {
    "type": "linear",
    "factor": 4.0
}

# 动态缩放（基于输入长度调整）
config.rope_scaling = {
    "type": "dynamic",
    "factor": 8.0
}

自定义RoPE实现

class LongContextAttention(nn.Module):
    """带RoPE的多头注意力。"""

    def __init__(self, hidden_size, num_heads, max_seq_len=32768):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        # Q, K, V投影
        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, hidden_size)
        self.v_proj = nn.Linear(hidden_size, hidden_size)
        self.o_proj = nn.Linear(hidden_size, hidden_size)

        # RoPE
        self.rotary_emb = RotaryEmbedding(
            dim=self.head_dim,
            max_seq_len=max_seq_len
        )

    def forward(self, hidden_states):
        batch_size, seq_len, _ = hidden_states.shape

        # 投影到Q, K, V
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        # 重塑为多头
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 应用RoPE
        cos, sin = self.rotary_emb(seq_len, device=hidden_states.device)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # 标准注意力
        attn_output = F.scaled_dot_product_attention(q, k, v)

        # 重塑和投影
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)
        output = self.o_proj(attn_output)

        return output

微调以支持长上下文

最小微调（位置插值）

from transformers import Trainer, TrainingArguments

# 扩展模型配置
model.config.max_position_embeddings = 32768
model.config.rope_scaling = {"type": "linear", "factor": 16.0}

# 训练参数（需要最小步骤）
training_args = TrainingArguments(
    output_dir="./llama-32k",
    num_train_epochs=1,
    max_steps=1000,           # 仅1000步骤！
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-5,
    warmup_steps=100,
    logging_steps=10,
    save_steps=500,
)

# 在长文档上训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=long_document_dataset,  # 32k令牌序列
)

trainer.train()

YaRN微调

# 克隆YaRN实现
git clone https://github.com/jquesnelle/yarn
cd yarn

# 用YaRN微调LLaMA
python scripts/train.py \
    --model meta-llama/Llama-2-7b-hf \
    --scale 16 \
    --rope_theta 10000 \
    --max_length 32768 \
    --batch_size 1 \
    --gradient_accumulation 16 \
    --steps 400 \
    --learning_rate 2e-5

最佳实践

1. 选择正确的方法

# 对于新模型（从零开始训练）
use_method = "ALiBi"  # 最佳外推，最低内存

# 对于扩展现有RoPE模型
use_method = "YaRN"  # 最有效的扩展（少10×数据）

# 对于快速扩展，计算量最小
use_method = "Position Interpolation"  # 1000步骤

# 对于中等扩展，效率良好
use_method = "Linear RoPE Scaling"  # 内置，简单

2. 缩放因子选择

# 保守（更安全，质量更好）
scaling_factor = 2.0  # 8k → 16k

# 中等（良好平衡）
scaling_factor = 4.0  # 8k → 32k

# 激进（需要更多微调）
scaling_factor = 8.0  # 8k → 64k
scaling_factor = 16.0  # 8k → 128k

# 规则：更大的因子需要更多微调步骤
steps_needed = 100 * scaling_factor  # 粗略估计

3. 微调数据

# ✅ 好：匹配目标长度的长文档
train_data = [
    {"text": long_doc_32k_tokens},  # 完整32k
    {"text": long_doc_24k_tokens},  # 变长
    {"text": long_doc_16k_tokens},
]

# ❌ 坏：短文档（不会学习长上下文）
train_data = [
    {"text": short_doc_2k_tokens},
]

# 使用数据集如：
# - PG-19（书籍，长文本）
# - arXiv论文
# - 长形式对话
# - GitHub仓库（连接文件）

4. 避免常见陷阱

# ❌ 坏：应用位置插值而不微调
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
# 模型在没有微调的情况下性能会很差！

# ✅ 好：缩放后微调
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
fine_tune(model, long_documents, steps=1000)

# ❌ 坏：没有数据的过于激进缩放
scale_to_1M_tokens()  # 没有大量微调不会工作

# ✅ 好：增量缩放
# 8k → 16k → 32k → 64k（每一步微调）

生产部署

长上下文推理

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载长上下文模型
model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/LLaMA-2-7B-32K",  # 32k上下文
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K")

# 处理长文档
long_text = "..." * 30000  # 30k令牌
inputs = tokenizer(long_text, return_tensors="pt", truncation=False).to('cuda')

# 生成
outputs = model.generate(
    **inputs,
    max_new_tokens=512,
    temperature=0.7,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)

内存优化

# 使用梯度检查点进行微调
model.gradient_checkpointing_enable()

# 使用Flash Attention 2
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    attn_implementation="flash_attention_2",  # 2-3×更快
    torch_dtype=torch.float16
)

# 使用分页注意力（vLLM）
from vllm import LLM

llm = LLM(
    model="togethercomputer/LLaMA-2-7B-32K",
    max_model_len=32768,  # 32k上下文
    gpu_memory_utilization=0.9
)

资源

• RoPE论文: https://arxiv.org/abs/2104.09864 (RoFormer)
• YaRN论文: https://arxiv.org/abs/2309.00071
• ALiBi论文: https://arxiv.org/abs/2108.12409 (Train Short, Test Long)
• 位置插值论文: 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

另请参见

• references/rope.md - 详细RoPE实现和理论
• references/extension_methods.md - YaRN、ALiBi、位置插值比较
• references/fine_tuning.md - 上下文扩展的完整微调指南