总目录 › Unsloth › 第 7 章
Chapter 07

长上下文：RoPE scaling + Unsloth 的长序列优化

📌 commit 1cf145c070ea 4090 单卡训 128K 上下文

Unsloth 的另一个营销点是"单卡训长上下文"。本章讲怎么把 Llama-3-8B 训到 32K / 128K， 涉及 RoPE scaling、chunked CE 的长 seq 受益、flex attention、和一些工程实践。

7.1长上下文训练的两个问题

问题本质解决方向
显存爆炸 attention 矩阵 $O(S^2)$；激活也线性增长 FlashAttention + chunked CE + gradient ckpt
位置编码外推 训练用 8K，想推 128K，RoPE 没见过那么长 RoPE scaling（YaRN / NTK / Linear）

7.2RoPE scaling 的数学

RoPE 给每个位置 $m$ 和每个维度 $i$ 计算频率：

$$ \theta_i = 10000^{-2i/d} $$

位置 $m$ 应用旋转 $m\theta_i$。问题：训练时只见过 $m \in [0, S_\text{train}]$，超出范围模型表现极差。

几种 scaling 方法

方法公式特点
Linear（PI） $m \to m / \alpha$ 简单，但所有频率一起压；高频损失
NTK-aware $\theta_i \to \theta_i \cdot \alpha^{-2i/d}$（base 变大） 低频压更狠、高频几乎不动
YaRN NTK + smoothing + temperature 当前最优；Llama-3 长上下文默认
LongRoPE 每个频率单独学缩放 Phi-3 用

7.3在 Unsloth 里启用 RoPE scaling

从 v2024.5 起 Unsloth 自动处理 RoPE scaling，模型 config 里有就用：

model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = 32768,        # ★ 直接设到目标长度
    dtype = None,
    load_in_4bit = True,
)

Unsloth 检测到 max_seq_length 超过模型原始训练长度（如 Llama-3 是 8K），会自动注入 YaRN scaling 配置。

如果想显式控制：

model, tokenizer = FastLanguageModel.from_pretrained(
    ...,
    rope_scaling = {
        "type": "yarn",
        "factor": 4.0,                  # 8K → 32K = 4× scaling
        "original_max_position_embeddings": 8192,
    },
)

7.4显存对照（长上下文 LoRA）

Llama-3-8B QLoRA + bs=1 + r=16 在不同 seq 下的峰值显存（4090 24G）：

seq_length纯 HF + PEFT+ Unsloth
2K 14 GB 8 GB
4K 18 GB 10 GB
8K 26 GB（OOM）13 GB
16K — 17 GB
32K — 21 GB
64K — 23 GB（极限）
128K — 需多卡或更省的设置

4090 单卡训到 64K 是 Unsloth 的甜点。再长需要多卡（Pro 版）或砍 batch / 砍 r。

7.5长上下文训练的关键开关

model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = 32768,
    dtype = None,
    load_in_4bit = True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    use_gradient_checkpointing = "unsloth",   # ★ 长上下文必开
    use_rslora = False,
)

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    max_seq_length = 32768,
    packing = False,                            # ★ 长上下文建议关 packing
    args = TrainingArguments(
        per_device_train_batch_size = 1,         # ★ 必须 1
        gradient_accumulation_steps = 4,
        bf16 = True,
        optim = "adamw_8bit",
        ...
    ),
)

关键点：

  1. max_seq_length 在 from_pretrained 和 SFTTrainer 都要设；
  2. 必须用 "unsloth" 模式 gradient checkpointing；
  3. batch size = 1；
  4. 关掉 packing（长上下文每条样本本来就长）。

7.6flex attention：自定义 attention mask

Unsloth 集成了 PyTorch FlexAttention（看 unsloth/kernels/flex_attention.py），让你能写复杂 mask 又不损失性能：

mask 类型用途
causal 标准 causal LM（默认）
sliding window 只看最近 N token，限制 attention 复杂度
document mask packing 时不让样本互看
prefix LM 前缀双向，后缀单向
retrieval-augmented RAG 时只看检索段落

使用：

from unsloth.kernels.flex_attention import sliding_window_mask

model.config.attention_mask_type = sliding_window_mask(window_size=4096)

7.7chunked CE 在长上下文里的作用

seq=32K + vocab=128K 时未 chunk 的 logits：

$$ 1 \times 32768 \times 128000 \times 4 \approx 16.8\ \text{GB} $$

单这一项就把 4090 压爆。chunked CE 把它降到几百 MB，没有 chunked CE 就没有 Unsloth 长上下文训练。 v2024.5+ 默认开启，用户无需配置。

7.8长上下文数据怎么准备

问题：常见公开数据集（alpaca / sharegpt）大多 1K-4K，怎么"训长上下文"？

方案 1：packing 短样本

把多条短样本拼接到 max_seq_length：

trainer = SFTTrainer(
    ...,
    packing = True,             # ★ 短数据训长上下文必开
    max_seq_length = 32768,
)

注意 packing 会把不同样本拼一起，建议用 flex attention 的 document mask 避免跨样本 attention。

方案 2：用真实长文档数据

数据集用途
togethercomputer/RedPajama-Data-V2 大规模长文档
LongBench / LongAlign 长上下文 benchmark + 训练数据
m-a-p/LongInstruct 长指令
wikipedia / arxiv 长文本预训练

方案 3：合成长 prompt

对话模型可以 inject "context-heavy" 任务，如多文档问答、长代码理解。让 GPT-4 / Claude 生成长 prompt 数据。

7.9长上下文训练的"实验体感"

seq4090 训练速度说明
2K ~ 60 it/s "日常"
4K ~ 30 it/s —
8K ~ 12 it/s —
16K ~ 4 it/s 开始慢
32K ~ 1.5 it/s 勉强可用
64K ~ 0.5 it/s 过夜跑

seq 翻倍 attention 复杂度翻 4 倍，所以速度衰减明显。训长上下文要算清楚时间预算。

7.10验证：是不是真"会用"长上下文

训完后用"大海捞针"测试：在长上下文中插入一条关键信息，让模型回答。例子（Needle In A Haystack）：

prompt = (
    "下面是一份长文档：\n"
    + ("这是一段无关内容。\n" * 5000)
    + "★ 关键信息：小明的生日是 1990 年 3 月 15 日。\n"
    + ("这是一段无关内容。\n" * 5000)
    + "\n问：小明的生日是？"
)
output = model.generate(...)

模型答不出来 = 长上下文训练失败（位置编码外推没成功）。LongBench 提供了标准化的长上下文评测。

7.11这章带走的