超长上下文预训练：位置编码、效率与稳定性
说明：本篇为综述/无第一手实验数据，采用公开文献与工程实践路线。
实验目标
"把上下文从 4K 拉到 128K"不是把 max_seq_len 改大那么简单——注意力从 O(n) 变成 O(n²)、位置编码外推失效、长序列训练不稳定，三个问题环环相扣。本篇目标：拆解超长上下文预训练的工程方法——长序列的位置编码与长度外推（RoPE/ALiBi）→ 注意力效率（FlashAttention/稀疏注意力/序列并行）→ 训练稳定性与数据配比（长短样本混合/续训策略）。读完能回答：从短上下文升到长上下文，先改什么、怎么控制成本、稳定性怎么保。
代码逐段拆解
第一段：位置编码与长度外推（RoPE 的伸缩）
python
# RoPE 位置编码的"长度外推"核心：对旋转角度做插值/缩放
# 线性缩放（Linear Scaling）：把超出训练长度的位置映射回训练范围
def rope_linear_scale(pos, dim, base=10000, scale=8.0):
theta = pos / scale # 除以缩放因子
freqs = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
return torch.outer(theta, freqs) # 位置 → 旋转角
# 更优的 NTK/YaRN：按"高频保持、低频缩放"逐维插值，外推质量更好
# 工程要点：RoPE 用 base 缩放 + 训练时覆盖目标长度（如先 8K 再 128K 两阶段）要点：
- 外推的三种姿势：直接外推（超出训练长度，质量急剧下降）、线性缩放（全维度均匀压，短距离退化）、NTK/YaRN（高频不变低频缩放，实践最佳）。"位置编码外推"是超长上下文的第一个工程决策。
- 两阶段训练是主流：先短上下文（8K）训到收敛，再长上下文（128K）续训——长度外推 + 续训配合，成本远低于从头长上下文训。
base是现成旋钮：RoPE 的 base（10000）调大即实现低频拉伸（NTK 思想的手动版），是"不改模型只改配置"的外推手段。
第二段：注意力效率——FlashAttention 与序列并行
python
# 注意力是 O(n²) 的墙：128K 上下文 = 每层 16384² 的注意力矩阵
# 手段 1：FlashAttention（内核级，不写大矩阵）
# 手段 2：序列并行（Sequence Parallelism）——把序列维切到多卡
# 手段 3：稀疏/分块注意力（局部窗口 + 全局 token，如 Longformer 思路）要点：
- FlashAttention 是"必装"不是"可选"：它把 O(n²) 显存开销降到 O(n)（分块计算+重算），128K 上下文没有 FlashAttention 基本不可能训——内核实现是 08/09 章 Triton 思路的典范。
- 序列并行是扩展手段：单卡放不下整个序列的激活时，把序列维切分到多卡（配合注意力分块）——这是"显存装不下"的并行答案，与 10 章的数据/张量并行正交。
- 稀疏注意力是精度取舍：局部窗口 + 全局 token 把每 token 的注意力限制在窗口内，效率高但表达能力受限——适合"长文档检索类"任务，通用训练不首选。
原理讲解
长上下文的成本账本
| 组件 | 短上下文（4K） | 长上下文（128K） |
|---|---|---|
| 注意力计算 | 每层 ~4K² | 每层 ~16M（32 倍） |
| KV Cache（推理） | 小 | 显存大头（12 章 KV 账本） |
| 激活显存 | 可控 | 需激活重算/序列并行 |
| 数据成本 | 普通语料 | 需长文档语料，获取难 |
核心认知：长上下文的成本爆炸在"注意力平方"与"长文档数据稀缺"，两个都要工程化解决——效率靠 FlashAttention/并行，数据靠"长文档采样 + 长短混合配比"。
训练稳定性与数据配比
混合配比（经验口径）：
短样本（<4K） 60% —— 保训练稳定性与泛化
中样本（4K-32K）30% —— 衔接
长样本（>32K） 10% —— 长上下文能力稳定性要点：全喂长样本会让 loss 剧烈波动（batch 内样本长度方差大、梯度噪声大）——长短混合 + 按长度分桶采样（同 batch 内长度相近）是稳定训练的关键；loss spike 治理接 10 章稳定性篇的方法。长上下文续训的学习率也要降（短→长阶段 lr 一般再降 5-10 倍）。
实验结果
（估算口径，非实测）
| 方案组合 | 相对成本 | 外推/长序列质量 |
|---|---|---|
| 直接外推（不处理位置编码） | 低 | 差（超训练长度急剧退化） |
| RoPE base 缩放 | 低 | 中（短距离略退化） |
| NTK/YaRN + 两阶段续训 | 中 | 好 |
| 从头训 128K + FlashAttention + 序列并行 | 高 | 最好但成本最高 |
工程启示
- 先做"两阶段 + 外推"再考虑从头训：短→长续训是成本最优路径，外推质量用 NTK/YaRN 类方法；从头训长上下文是最后的选择。
- FlashAttention 是前提：128K 上下文没它无法落地，先确认训练框架支持（10 章框架篇衔接）。
- 数据配比决定稳定性：长短混合 + 长度分桶采样，长样本占比 10% 起步逐步调——别一上来全长文档。
- 续训降 lr、监控 loss spike：短→长阶段 lr 降 5-10 倍，配合 10 章 spike 检测与断点重启，长上下文训练才能跑稳。
下一篇预告：长上下文把序列维打穿了，下一步看调度与资源编排——"Kubernetes 上 GPU 作业编排：Kueue/Volcano 队列与配额"。