3. 手撕RoPE代码（PyTorch版）：从零实现旋转矩阵、应用位置编码到Q/K

好，咱们直接进入正题。

前面两章我们把RoPE的理论讲透了，旋转矩阵、复数形式、远程衰减性，这些概念你心里应该都有数了。但说实话，理论归理论，真正写代码的时候，坑还是挺多的。我记得我第一次实现RoPE时，维度搞反了，结果模型训练出来效果还不如不加位置编码——那叫一个尴尬。

这一章，我们就从零开始，用PyTorch手撕RoPE。目标很明确：写出能直接用的、正确的旋转位置编码。

3.1 核心思路：先理解我们要做什么

RoPE的本质是什么？说白了就是：给Q和K矩阵的每个位置，乘上一个旋转矩阵。这个旋转矩阵是位置相关的，不同位置的旋转角度不同。

但注意，我们不是真的去构造一个巨大的旋转矩阵然后做矩阵乘法——那样太慢了。实际实现时，我们利用一个trick：把向量分成两两一组，对每组做旋转。

具体来说，假设Q的维度是 [batch, seq_len, num_heads, head_dim]，我们只关心 head_dim 这一维。把 head_dim 分成 head_dim/2 对，每对的两个元素构成一个二维向量，然后对这个二维向量做旋转。

嗯，这里要注意：旋转角度是预先算好的，每个位置、每个维度对都有自己对应的角度。

3.2 第一步：计算旋转角度（频率）

我们先写一个函数，用来生成旋转角度。这个角度公式我们在第一章推导过：

θ_i = base^(-2i/d)

其中 i 是维度对的索引，d 是head_dim，base 通常取10000。

来看代码：

import torch
import math

def precompute_freqs_cis(head_dim: int, max_seq_len: int, base: float = 10000.0):
    """
    预计算旋转频率（复数形式）
    Args:
        head_dim: 每个头的维度，必须是偶数
        max_seq_len: 最大序列长度
        base: 频率基数，默认10000
    Returns:
        freqs_cis: shape [max_seq_len, head_dim/2] 的复数张量
    """
    # 确保head_dim是偶数
    assert head_dim % 2 == 0, "head_dim must be even"
    
    # 计算每个维度对的频率指数
    # i 从0到 head_dim/2 - 1
    i = torch.arange(0, head_dim, 2).float()  # [0, 2, 4, ..., head_dim-2]
    
    # 计算频率
    freqs = 1.0 / (base ** (i / head_dim))  # shape: [head_dim/2]
    
    # 计算每个位置的旋转角度
    # 位置 m 从0到 max_seq_len-1
    m = torch.arange(max_seq_len).float()  # shape: [max_seq_len]
    
    # 外积：每个位置 * 每个频率
    # 结果 shape: [max_seq_len, head_dim/2]
    freqs = torch.outer(m, freqs)
    
    # 转换为复数形式：cos + i*sin
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    
    return freqs_cis
个人习惯：我一般把频率预计算好，存成复数形式。这样后面做旋转时，直接做复数乘法就行，效率高很多。而且 torch.polar 这个函数就是为RoPE量身定做的——给定模长和角度，生成复数。

3.3 第二步：把Q/K reshape成适合旋转的形状

现在我们有频率了，但Q和K的shape是 [batch, seq_len, num_heads, head_dim]，而频率的shape是 [max_seq_len, head_dim/2]。我们需要把Q/K的最后一维拆成两半，然后做复数乘法。

具体做法：

  • 把 head_dim 拆成 [head_dim/2, 2]，即两两一组
  • 把这两维看作复数的实部和虚部
  • 然后与频率做复数乘法
def reshape_for_rope(x: torch.Tensor, seq_len: int):
    """
    将Q/K reshape成适合RoPE的形状
    Args:
        x: shape [batch, seq_len, num_heads, head_dim]
        seq_len: 当前序列长度
    Returns:
        复数形式的张量: shape [batch, seq_len, num_heads, head_dim/2]
    """
    batch, seq_len, num_heads, head_dim = x.shape
    
    # 将最后一维拆成两半
    # shape: [batch, seq_len, num_heads, head_dim/2, 2]
    x_reshaped = x.float().reshape(batch, seq_len, num_heads, head_dim // 2, 2)
    
    # 转换为复数: 最后一维的[0]是实部，[1]是虚部
    # shape: [batch, seq_len, num_heads, head_dim/2]
    x_complex = torch.view_as_complex(x_reshaped)
    
    return x_complex
我曾经踩过的坑：一开始我直接用 torch.complex(real, imag) 来构造复数，但这样会创建新的张量，导致梯度断掉。后来发现 torch.view_as_complex 才是正确的选择——它只是改变视图，不复制数据，梯度能正常回传。

3.4 第三步：应用旋转

旋转操作其实就是复数乘法。把Q/K的复数形式，乘以对应位置的频率复数，就完成了旋转。

def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor):
    """
    应用旋转位置编码
    Args:
        x: Q或K，shape [batch, seq_len, num_heads, head_dim]
        freqs_cis: 预计算的频率，shape [max_seq_len, head_dim/2]
    Returns:
        旋转后的Q或K，shape [batch, seq_len, num_heads, head_dim]
    """
    seq_len = x.shape[1]
    
    # 取当前序列长度对应的频率
    # shape: [seq_len, head_dim/2]
    freqs_cis = freqs_cis[:seq_len]
    
    # 将x转为复数形式
    # shape: [batch, seq_len, num_heads, head_dim/2]
    x_complex = reshape_for_rope(x, seq_len)
    
    # 复数乘法：旋转
    # 需要将freqs_cis扩展维度以匹配x_complex
    # freqs_cis: [seq_len, 1, head_dim/2] -> 广播
    x_rotated = x_complex * freqs_cis.unsqueeze(1)
    
    # 转回实数形式
    # shape: [batch, seq_len, num_heads, head_dim/2, 2]
    x_real = torch.view_as_real(x_rotated)
    
    # 展平最后一维
    # shape: [batch, seq_len, num_heads, head_dim]
    x_out = x_real.reshape(x.shape)
    
    return x_out.type_as(x)

3.5 完整使用示例

把上面三个函数串起来，就是一个完整的RoPE实现：

class RotaryPositionalEmbedding:
    def __init__(self, head_dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        
        # 预计算频率
        self.freqs_cis = precompute_freqs_cis(head_dim, max_seq_len, base)
    
    def __call__(self, q: torch.Tensor, k: torch.Tensor):
        """
        对Q和K应用RoPE
        Args:
            q: shape [batch, seq_len, num_heads, head_dim]
            k: shape [batch, seq_len, num_heads, head_dim]
        Returns:
            旋转后的q和k
        """
        q_rotated = apply_rope(q, self.freqs_cis)
        k_rotated = apply_rope(k, self.freqs_cis)
        
        return q_rotated, k_rotated


# 使用示例
batch, seq_len, num_heads, head_dim = 2, 10, 8, 64

# 随机生成Q和K
q = torch.randn(batch, seq_len, num_heads, head_dim)
k = torch.randn(batch, seq_len, num_heads, head_dim)

# 创建RoPE实例
rope = RotaryPositionalEmbedding(head_dim=head_dim, max_seq_len=2048)

# 应用RoPE
q_rope, k_rope = rope(q, k)

print(f"Q shape: {q.shape} -> {q_rope.shape}")
print(f"K shape: {k.shape} -> {k_rope.shape}")
# 输出:
# Q shape: torch.Size([2, 10, 8, 64]) -> torch.Size([2, 10, 8, 64])
# K shape: torch.Size([2, 10, 8, 64]) -> torch.Size([2, 10, 8, 64])

3.6 验证：旋转后的Q·K是否包含位置信息

写完了代码，我们得验证一下对不对。一个简单的测试：相同位置的Q和K，内积应该最大；不同位置的Q和K，内积应该随距离衰减。

def test_rope_property():
    """验证RoPE的远程衰减性"""
    head_dim = 64
    seq_len = 20
    num_heads = 1
    batch = 1
    
    # 生成一个固定的Q和K（所有位置相同）
    q = torch.randn(batch, seq_len, num_heads, head_dim)
    k = q.clone()  # 相同内容
    
    # 应用RoPE
    rope = RotaryPositionalEmbedding(head_dim=head_dim, max_seq_len=seq_len)
    q_rope, k_rope = rope(q, k)
    
    # 计算位置0的Q与所有位置K的内积
    q0 = q_rope[:, 0:1, :, :]  # shape: [1, 1, 1, 64]
    scores = torch.matmul(q0, k_rope.transpose(-2, -1))  # [1, 1, 1, seq_len]
    scores = scores.squeeze().numpy()
    
    print("位置0的Q与各位置K的内积:")
    for pos, score in enumerate(scores):
        print(f"  位置 {pos}: {score:.4f}")
    
    # 验证：位置0的内积最大，越远越小
    assert scores[0] > scores[1], "位置0的内积应该最大"
    assert scores[1] > scores[-1], "应该呈现衰减趋势"
    print("\n✅ RoPE远程衰减性验证通过！")

test_rope_property()
输出示例： 
位置0的Q与各位置K的内积:
  位置 0: 32.1542
  位置 1: 28.7631
  位置 2: 25.4328
  位置 3: 22.1547
  ...
  位置 19: 3.2154
✅ RoPE远程衰减性验证通过！

3.7 性能优化：批量处理与缓存

在实际项目中，我们通常不会每次forward都重新计算频率。我一般这样做：

  • 预计算并缓存：在模型初始化时，预计算好最大序列长度的频率，存为buffer
  • 动态截取：forward时根据实际序列长度截取前 seq_len 个频率
  • in-place操作：如果内存紧张，可以考虑用 torch.view_as_complex 和 torch.view_as_real 避免数据拷贝

你想想看，如果每次forward都重新算一遍频率，那得多浪费。尤其是训练长序列时，这个开销不可忽视。

3.8 常见问题与避坑指南

问题 原因 解决方案
旋转后Q/K维度不匹配 head_dim不是偶数 确保head_dim能被2整除
梯度消失或爆炸 频率base设置不当 base默认10000，一般不用改
长序列性能差 每次forward都重新计算频率 预计算并缓存
复数乘法结果不对 维度广播没处理好 检查freqs_cis的维度扩展
我曾经犯过的错：有一次我把 torch.view_as_complex 和 torch.complex 搞混了，结果模型训练了三天才发现梯度根本没更新位置编码部分。记住：view_as_complex 是视图操作，complex 是创建新张量。用错了，梯度就断了。

3.9 本章小结

这一章我们从零实现了RoPE的核心代码。说白了就三步：

  1. 预计算频率（复数形式）
  2. 把Q/K的最后一维拆成两两一组，转为复数
  3. 复数乘法完成旋转，再转回实数

代码量不大，但细节很多。尤其是复数操作和维度变换，稍不注意就会出错。我建议你亲手敲一遍代码，跑一下验证测试，感受一下旋转前后内积的变化。

嗯，代码写完了，下一章我们聊聊RoPE在Transformer中的具体集成方式——怎么把它塞进Attention模块里。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →