位置编码与ROPE

本文按四部分展开：位置编码 → RoPE → TP 切分影响 → SP 切分影响。

本文由AI生成

位置编码

是什么

位置编码（Positional Encoding）是给序列中每个 token 额外注入”我在第几个位置”信息的机制。
它可以是加在输入 embedding 上的向量（正弦编码、可学习编码），也可以像 RoPE 那样
在 attention 内部对 q、k 做变换。目的只有一个：让 attention 知道 token 的顺序。

为什么需要：attention 本身是”无序”的

attention 的计算公式：

$$\text{Attention}(Q,K,V)=\text{softmax}!\left(\frac{QK^{\top}}{\sqrt{D}}\right)V$$

注意一个关键事实：公式里只有”第 i 个 token 的 q”和”第 j 个 token 的 k”的内积，
完全没有位置信息。分数 $S_{ij}=q_i\cdot k_j$ 只由两个 token 的内容决定。

详细例子：交换两个 token 之后，注意力分数不变

设 2 个 token，head 维度 D=2，为了极简，令每个 token 的 q、k、v 只由内容决定：

token 内容	q	k	v
A（如”猫”）	[1, 0]	[1, 0]	[1, 0]
B（如”吃”）	[1, 1]	[0, 1]	[0, 1]

（特意选了不对称的 q、k：下面会看到 $q_A\cdot k_B = 0$ 而 $q_B\cdot k_A = 1$。）

排列一：序列 [A, B]

$$S = QK^{\top}=\begin{bmatrix}q_A\cdot k_A & q_A\cdot k_B\ q_B\cdot k_A & q_B\cdot k_B\end{bmatrix}=\begin{bmatrix}1 & 0\ 1 & 1\end{bmatrix}$$

逐格展开：

S[0][0] = q_A·k_A = [1,0]·[1,0] = 1×1 + 0×0 = 1    （A 看自己）
S[0][1] = q_A·k_B = [1,0]·[0,1] = 1×0 + 0×1 = 0    （A 看 B）
S[1][0] = q_B·k_A = [1,1]·[1,0] = 1×1 + 1×0 = 1    （B 看 A）
S[1][1] = q_B·k_B = [1,1]·[0,1] = 1×0 + 1×1 = 1    （B 看自己）

对每一行做 softmax（$\frac{e^1}{e^1+e^0}=0.731$，$\frac{e^0}{e^1+e^0}=0.269$）：

• 内容 A 那一行：分数 [1, 0] → 分布 **[0.731, 0.269]**（按内容顺序 A、B）
• 内容 B 那一行：分数 [1, 1] → 分布 [0.5, 0.5]

输出：$o_A = 0.731,v_A + 0.269,v_B$（内容 A 的输出），$o_B = 0.5,v_A + 0.5,v_B$。

排列二：交换成 [B, A]

交换后，位置 0 上站的是内容 B，位置 1 上站的是内容 A。逐格展开：

S'[0][0] = q_B · k_B = [1,1]·[0,1] = 1×0 + 1×1 = 1    （B 看自己）
S'[0][1] = q_B · k_A = [1,1]·[1,0] = 1×1 + 1×0 = 1    （B 看 A）
S'[1][0] = q_A · k_B = [1,0]·[0,1] = 1×0 + 0×1 = 0    （A 看 B）
S'[1][1] = q_A · k_A = [1,0]·[1,0] = 1×1 + 0×0 = 1    （A 看自己）

$$S’ = \begin{bmatrix}1 & 1\ 0 & 1\end{bmatrix}\ \ne\ S = \begin{bmatrix}1 & 0\ 1 & 1\end{bmatrix}$$

注意 $S’ \neq S$——这正是普遍情形（$q_A\cdot k_B = 0$ 而 $q_B\cdot k_A = 1$，
两者本来就不必相等）。交换位置做的是同时交换行和列（$S’ = P,S,P^{\top}$，
P 为置换矩阵）：四个分数 {1, 0, 1, 1} 一个不差，只是换了格子落座——
比如 S’[1][0] 里的 $q_A\cdot k_B = 0$ 就是排列一里的 S[0][1]，以此类推。

而”对顺序无感”依赖的不是 $S’=S$，而是：每一行 softmax 面对的分数集合没变——

• 内容 B 无论站在位置 0 还是 1，都是对 {A, B} 打分，分数都是 {1, 1} → 分布 [0.5, 0.5] 不变；
• 内容 A 的分数都是 {1, 0} → 分布 [0.731, 0.269] 不变（只是写出来的顺序随内容排列颠倒）。

softmax 是逐元素作用在”对每个内容的分数集合”上的，集合相同则分布相同。于是：

• 位置 0（现在是 B）的输出 $= 0.5,v_B + 0.5,v_A = o_B$
• 位置 1（现在是 A）的输出 $= 0.269,v_B + 0.731,v_A = o_A$

整个序列的输出恰好是原输出的逐元素换位：[o_A, o_B] → [o_B, o_A]。

结论（置换等变性）：不带位置信息的 attention 是一个”集合运算”——
把序列任意打乱，输出只是同样地跟着打乱，模型**完全无法区分
“猫吃鱼” 和 “鱼吃猫”**。语言顺序是有语义的，所以必须显式注入位置信息，
这就是位置编码存在的理由。

---

RoPE

RoPE 是什么

RoPE（Rotary Position Embedding，旋转位置编码）出自 RoFormer（Su et al., 2021）。
核心思想：把位置变成”角度”，在 attention 之前对 q、k 做一次与位置有关的旋转，
使得旋转后的内积只依赖两个 token 的相对距离：

$$\langle R_m q,\ R_n k\rangle = q^{\top} R_{n-m}, k$$

形式上是绝对位置编码（每个位置有确定的旋转角），效果上是相对位置编码
（分数只看 m−n），且不引入任何可学习参数、不改任何形状、天然支持 KV cache。
LLaMA、Qwen 等主流模型均采用。

针对 [B, N, S, D] 形状的 qkv

attention 输入的 q、k、v 形状为 [B, N, S, D]：
B = batch size，N = head 数量，S = 序列长度，D = head_dim（每个头的维度，
注意不是隐藏层总维度——隐藏层 H 投影后切成 N×D）。

RoPE 作用在哪个维度？

严格分三层说：

• D 维（变换发生的地方）：旋转发生在每个 D 维向量内部——
  两两配对做二维旋转，不跨 token、不跨头、不跨 batch。
• S 维（每个位置的值都变，但各位置互相独立）：位置 m 决定该行旋转的角度
  θ = m·ω。不同位置角度不同，所以每行数值都会被改，但行与行之间不混合。
• B、N 维（完全不动）：所有 batch、所有 head 共用同一张 cos/sin 表，
  实现上靠广播 (1, 1, S, D/2) 作用到全部切片。

一句话：RoPE 是一个”按位置取参数”的 D 维旋转变换。

频率表是什么，怎么算？

频率表就是每个 (位置, 维度对) 的旋转角 θ，分两步：

第一步：每个维度对 i 的基础转速（频率）

$$\omega_i = \text{base}^{-2i/D},\quad \text{base}=10000$$

• D=4 时：ω = [10000⁰, 10000⁻⁰·⁵] = [1, 0.01]
• 低维对转得快（每前进一个 token 转 1 弧度），高维对转得慢——
  快钟分辨近距离，慢钟分辨远距离，类似秒针与时针。

第二步：乘上位置得到角度

$$\theta_{m,i} = m \cdot \omega_i$$

以 D=4、S=3 为例，角度表 shape 为 (S, D/2) = (3, 2)：

m	θ₀ = m×1	θ₁ = m×0.01
0	0	0
1	1	0.01
2	2	0.02

随后查 cos/sin，拼成 (S, D) 的表（实现上常直接保存 (S, D/2) 再广播）。
表只依赖 S 和 D，与 batch、head 无关，可启动时预计算缓存到 max_len，
decode 时按绝对位置偏移取行即可。

两两配对的方式？有哪几种？

把 D 维向量的相邻/对应两维看作复平面上的一个点 (实部, 虚部)，D 维共有 D/2 对。
主流实现有两种配对方式，数学上完全等价，只是维度排列约定不同：

方式一：相邻配对（GPT-J / 复数原始形式，interleaved）

第 i 对取 (x_{2i}, x_{2i+1})，即 (x0,x1)、(x2,x3)、…
cos/sin 在维度方向按 [c0, s0, c1, s1, …] 交错排列。

方式二：对半配对（GPT-NeoX / HuggingFace rotate_half 形式）

第 i 对取 (x_i, x_{i+D/2})，即 (x0,x_{D/2})、(x1,x_{D/2+1})、…
cos/sin 按 [c0, c1, …, c0, c1, …] 前后两半拼接（cat(angles, angles)），
配合 rotate_half（把向量后半取负交换）实现同一旋转。

两者的关系：同一个旋转矩阵作用在不同的维度排列上，
从一种格式加载另一种格式的 checkpoint 时必须做维度 permute
（如 GPT-J → LLaMA 转换），否则 RoPE 角度全部错位，模型表现诡异且难排查。
本文示例采用方式二（HuggingFace rotate_half 即此形式）。

补充：2D RoPE（视觉/DiT 模型）：图像 patch 有 (h, w) 两个坐标，
常见做法是各坐标独立做一维 RoPE 再拼接（或交错），Qwen-Image、
各类 DiT 的 latent 序列多用此变体；原理不变，只是”位置”从一维变成多维。

算完 RoPE 后的形状？

完全不变，仍是 [B, N, S, D]。 旋转是逐元素的乘加（x·cos ± rotate_half(x)·sin），
不增删维度、不引入参数、不需要 padding。形状链路：

q, k: [B,N,S,D] --RoPE(乘 cos/sin + rotate_half)--> [B,N,S,D] --QKᵀ--> [B,N,S,S] --softmax·V--> [B,N,S,D]

这也是 RoPE 工程上讨喜的原因：插入点只有 attention 里的一行，其他一切照旧。

内嵌数值例子（自包含，D=4，q=k=[1,2,3,4]，位置 m=0,1,2），完整走一遍验证”只依赖相对距离”：

旋转（以 m=1 为例）：对半拆得前半 [1,2]、后半 [3,4]，跨半配对成 (1,3) 和 (2,4)：

• 第 0 对 (1,3)，θ₀=1：新a = 1×cos1 − 3×sin1 = −1.9841，新b = 1×sin1 + 3×cos1 = 2.4624
• 第 1 对 (2,4)，θ₁=0.01：新a = 1.9599，新b = 4.0198
• 拼回：q_rot = [−1.9841, 1.9599, 2.4624, 4.0198]（实部放前半、虚部放后半），长度不变

三个位置都转完后算分数 S = Q·Kᵀ（每格 = 模长平方 × cos(角度差)）：

S = [[30.000, 25.402, 15.835],
     [25.402, 30.000, 25.402],
     [15.835, 25.402, 30.000]]

每一格严格等于 10·cos(m−n) + 20·cos(0.01·(m−n))
（10 = 1²+3²、20 = 2²+4² 是两对向量的模长平方）。距离为 1 的格子全是 25.402，
距离为 2 的全是 15.835——分数矩阵沿对角线平移不变，只依赖相对位置 m−n；
对角线恒为 30，因为旋转不改变模长。

---

TP 切分对 RoPE 影响

TP（Tensor Parallelism）沿 N（head 数） 切分 attention：
每张卡持有 N/tp 个完整的头（head_dim D 不切，完整地留在单卡上）。

结论：RoPE 与 TP 天然兼容，零改动、零通信。 原因：

1. RoPE 只作用在 D 维内部，且逐头独立（每个头自己转自己的），
   切走一部分头后，剩下的头在该卡上做 RoPE 与单卡完全一致；
2. cos/sin 表对所有 head 相同、只依赖 (S, D/2)，体积极小，
   每张卡冗余保存一份即可，无需分片、无需 all-gather；
3. 旋转发生在 QKᵀ 之前，TP 在 attention 内部本来就持有完整 S 维，
   不涉及任何跨卡位置信息交换。

两个工程注意点（都是坑而非原理问题）：

• head→rank 的划分必须与 checkpoint 的 head 排列一致。TP 按 head 维切块，
  若加载权重时 head 顺序错乱（或”两两配对的方式”一节中 GPT-J/NeoX 配对格式 permute 漏做），
  每张卡上”头”与”权重”错位，RoPE 角度随之错位——训练能跑、loss 正常下降一点点后停滞，
  极难排查。
• 长度外推改频率（NTK-aware、YaRN 等）时，所有 TP rank 必须用同一份缩放后的频率表，
  保证各卡角度一致（通常配置在模型 config 里天然一致，但动态改 base 的实现要注意广播）。

---

SP 切分对 RoPE 影响

SP（Sequence Parallelism）在 TP 基础上进一步沿 S（序列维） 切分 activation，
每张卡只持有 S/tp 段 token。分两种情况讨论：

情况一：Megatron 式 SP（attention 核心计算前把 S 拼回完整）

Megatron-LM 的 SP 在 LayerNorm/Dropout 等 region 按 S 切，
进 attention 前通过 all-gather 把序列拼回完整 S（TP 维度上仍是各卡不同的 head），
算完再 reduce-scatter 回去。因此 RoPE 看到的仍是完整 S 维，
计算与纯 TP 完全相同，无需任何修改。

情况二：attention 内部也切 S（Context Parallelism，CP）

每张卡只持有自己那 S/cp 段 token，attention 时通过 CP 通信换取其他卡的 KV。
此时 RoPE 的正确性依赖一个不变量：

每个 token 的全局绝对位置 id 不因切分而改变。

做法：cos/sin 预计算完整表 (max_S, D/2)，每张卡按自己 token 的全局位置偏移
（如 rank * local_S + local_idx）取出对应行做旋转。由于 RoPE 的分数只依赖
角度差 m−n，只要 q、k 都按各自真实全局位置旋转，切分前后数学上严格等价，
不需要任何补偿项——这正是相对位置编码在分布式下的优势
（可学习绝对位置编码按位置查表，表长受限且与切分强耦合，就没有这么省心）。

CP 特有的注意点：

• 位置映射要跟负载均衡方案走：为均衡 causal attention 的计算量，
  Megatron CP 默认用 zigzag（前后段交错配对）切分，此时 rank 上第 i 个 token
  的全局位置不是简单的 rank*local_S + i，必须按 zigzag 映射查 cos/sin 行；
• KV 通信的是旋转后的 k（RoPE 在通信前本地完成，k 缓存/传输都带位置信息，
  不需要额外传位置）；
• 推理侧同理：KV cache 按卡存各自 token 段的旋转后 k，
  新 token 用自己的绝对位置取 cos/sin 行。

一句话总结：RoPE = “按全局位置取角度的 D 维旋转”。
TP 切 N——完全无感；SP/CP 切 S——只要保证全局位置 id 正确、
按偏移取 cos/sin 行即可，相对位置特性使它天然适配各种并行切分。