【大模型实战篇】RoPE旋转位置编码PyTorch代码分析

1. 背景介绍   

         之前我们通过两篇技术文章《LLaMA3结构关键模块分析》和《RoPE旋转位置编码底层数学原理分析》对旋转位置编码RoPE的原理进行了必要的讲解。接下来，我们将针对来自rotary-embedding-torch【1】开源库的实现，对RoPE的PyTorch代码实现进行分析，主要对关键代码进行讲解，帮助自己更好理解相应的实现原理。该项目除了实现常规的旋转位置编码【2】，另外还支持轴向旋转嵌入、长度可外推的旋转嵌入【3】、插值序列位置【4】。其中【3】通过给旋转位置编码（rotary embeddings）添加类似于ALiBi的衰减机制，解决了长度外推问题，该技术命名为XPos。【4】通过在插值后的序列位置上进行微调，以便将预训练模型扩展到更长的上下文长度。

2. torch代码实现

        以下是具体的代码实现和用法。

rotary_embedding_torch.py

from __future__ import annotations
from math import pi, log

import torch
from torch.nn import Module, ModuleList
from torch.cuda.amp import autocast
from torch import nn, einsum, broadcast_tensors, Tensor

from einops import rearrange, repeat

from typing import Literal


# 辅助函数

def exists(val):
    return val is not None


def default(val, d):
    return val if exists(val) else d


# 广播张量并沿给定维度连接

def broadcat(tensors, dim=-1):
    broadcasted_tensors = broadcast_tensors(*tensors)
    return torch.cat(broadcasted_tensors, dim=dim)


# 旋转位置编码辅助函数

def rotate_half(x):
    x = rearrange(x, '... (d r) -> ... d r', r=2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return rearrange(x, '... d r -> ... (d r)')


@autocast(enabled=False)
def apply_rotary_emb(freqs, t, start_index=0, scale=1., seq_dim=-2):
    dtype = t.dtype

    if t.ndim == 3:
        seq_len = t.shape[seq_dim]
        freqs = freqs[-seq_len:]

    rot_dim = freqs.shape[-1]
    end_index = start_index + rot_dim

    assert rot_dim <= t.shape[
        -1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'

    # 将t分为三部分：左、中（要变换的部分）和右
    t_left = t[..., :start_index]
    t_middle = t[..., start_index:end_index]
    t_right = t[..., end_index:]

    # 应用旋转嵌入
    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)

    out = torch.cat((t_left, t_transformed, t_right), dim=-1)

    return out.type(dtype)


# 学习旋转辅助函数

def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
    if exists(freq_ranges):
        rotations = einsum('..., f -> ... f', rotations, freq_ranges)
        rotations = rearrange(rotations, '... r f -> ... (r f)')

    rotations = repeat(rotations, '... n -> ... (n r)', r=2)
    return apply_rotary_emb(rotations, t, start_index=start_index)


# 旋转嵌入类

class RotaryEmbedding(Module):
    def __init__(
            self,
            dim,
            custom_freqs: Tensor | None = None,
            freqs_for: Literal['lang', 'pixel', 'constant'] = 'lang',
            theta=10000,
            max_freq=10,
            num_freqs=1,
            learned_freq=False,
            use_xpos=False,
            xpos_scale_base=512,
            interpolate_factor=1.,
            theta_rescale_factor=1.,
            seq_before_head_dim=False,
            cache_if_possible=True
    ):
        super().__init__()
        # 来自Reddit用户bloc97的建议，重新调整旋转嵌入以适应更长的序列长度，无需微调
        # 与NTK文献有一定联系
        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/


        theta *= theta_rescale_factor ** (dim / (dim - 2))

        self.freqs_for = freqs_for

        if exists(custom_freqs):
            freqs = custom_freqs
        elif freqs_for == 'lang':
            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
        elif freqs_for == 'pixel':
            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
        elif freqs_for == 'constant':
            freqs = torch.ones(num_freqs).float()

        self.cache_if_possible = cache_if_possible

        self.tmp_store('cached_freqs', None)
        self.tmp_store('cached_scales', None)

        self.freqs = nn.Parameter(freqs, requires_grad=learned_freq)

        self.learned_freq = learned_freq

        # 设备占位符

        self.tmp_store('dummy', torch.tensor(0))

        # 默认序列维度

        self.seq_before_head_dim = seq_before_head_dim
        self.default_seq_dim = -3 if seq_before_head_dim else -2

        # 插值因子

        assert interpolate_factor >= 1.
        self.interpolate_factor = interpolate_factor

        # xpos 相关设置

        self.use_xpos = use_xpos
        if not use_xpos:
            self.tmp_store('scale', None)
            return

        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
        self.scale_base = xpos_scale_base
        self.tmp_store('scale', scale)

        # 添加apply_rotary_emb作为静态方法
        self.apply_rotary_emb = staticmethod(apply_rotary_emb)

    @property
    def device(self):
        return self.dummy.device

    def tmp_store(self, key, value):
        self.register_buffer(key, value, persistent=False)

    def get_seq_pos(self, seq_len, device, dtype, offset=0):
        return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor

    def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None):
        seq_dim = default(seq_dim, self.default_seq_dim)

        assert not self.use_xpos or exists(
            scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings'

        device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]

        seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)

        freqs = self.forward(seq, seq_len=seq_len, offset=offset)

        if seq_dim == -3:
            freqs = rearrange(freqs, 'n d -> n 1 d')

        return apply_rotary_emb(freqs, t, scale=default(scale, 1.), seq_dim=seq_dim)

    def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
        dtype, device, seq_dim = q.dtype, q.device, default(seq_dim, self.default_seq_dim)

        q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
        assert q_len <= k_len

        q_scale = k_scale = 1.

        if self.use_xpos:
            seq = self.get_seq_pos(k_len, dtype=dtype, device=device)

            q_scale = self.get_scale(seq[-q_len:]).type(dtype)
            k_scale = self.get_scale(seq).type(dtype)

        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** -1)

        rotated_q = rotated_q.type(q.dtype)
        rotated_k = rotated_k.type(k.dtype)

        return rotated_q, rotated_k

    def rotate_queries_and_keys(self, q, k, seq_dim=None):
        seq_dim = default(seq_dim, self.default_seq_dim)

        assert self.use_xpos
        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]

        seq = self.get_seq_pos(seq_len, dtype=dtype, device=device)

        freqs = self.forward(seq, seq_len=seq_len)
        scale = self.get_scale(seq, seq_len=seq_len).to(dtype)

        if seq_dim == -3:
            freqs = rearrange(freqs, 'n d -> n 1 d')
            scale = rearrange(scale, 'n d -> n 1 d')

        rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)
        rotated_k = apply_rotary_emb(freqs, k, scale=scale ** -1, seq_dim=seq_dim)

        rotated_q = rotated_q.type(q.dtype)
        rotated_k = rotated_k.type(k.dtype)

        return rotated_q, rotated_k

    def get_scale(
            self,
            t: Tensor,
            seq_len: int | None = None,
            offset=0
    ):
        assert self.use_xpos

        should_cache = (
                self.cache_if_possible and
                exists(seq_len)
        )

        if (
                should_cache and \
                exists(self.cached_scales) and \
                (seq_len + offset) <= self.cached_scales.shape[0]
        ):
            return self.cached_scales[offset:(offset + seq_len)]

        scale = 1.
        if self.use_xpos:
            power = (t - len(t) // 2) / self.scale_base
            scale = self.scale ** rearrange(power, 'n -> n 1')
            scale = torch.cat((scale, scale), dim=-1)

        if should_cache:
            self.tmp_store('cached_scales', scale)

        return scale

    def get_axial_freqs(self, *dims):
        Colon = slice(None)
        all_freqs = []

        for ind, dim in enumerate(dims):
            if self.freqs_for == 'pixel':
                pos = torch.linspace(-1, 1, steps=dim, device=self.device)
            else:
                pos = torch.arange(dim, device=self.device)

            freqs = self.forward(pos, seq_len=dim)

            all_axis = [None] * len(dims)
            all_axis[ind] = Colon

            new_axis_slice = (Ellipsis, *all_axis, Colon)
            all_freqs.append(freqs[new_axis_slice])

        all_freqs = broadcast_tensors(*all_freqs)
        return torch.cat(all_freqs, dim=-1)

    @autocast(enabled=False)
    def forward(
            self,
            t: Tensor,
            seq_len=None,
            offset=0
    ):
        should_cache = (
                self.cache_if_possible and \
                not self.learned_freq and \
                exists(seq_len) and \
                self.freqs_for != 'pixel'
        )

        if (
                should_cache and \
                exists(self.cached_freqs) and \
                (offset + seq_len) <= self.cached_freqs.shape[0]
        ):
            return self.cached_freqs[offset:(offset + seq_len)].detach()

        freqs = self.freqs

        freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs)
        freqs = repeat(freqs, '... n -> ... (n r)', r=2)

        if should_cache:
            self.tmp_store('cached_freqs', freqs.detach())

        return freqs

rope_test.py

import torch
from rotary_embedding_torch import RotaryEmbedding

# 初始化旋转位置嵌入类实例，指定特征维度为32
rotary_emb = RotaryEmbedding(dim=32)

# 模拟查询（queries）和键（keys）张量
# 张量的最后两个维度应该是序列长度（seq_len）和特征维度（feature dimension）
q = torch.randn(1, 8, 1024, 64)  # 查询张量，形状为 (batch, heads, seq len, dimension of head)
k = torch.randn(1, 8, 1024, 64)  # 键张量，形状相同

print("初始查询张量 q:", q)
print("初始键张量 k:", k)

# 将旋转位置编码应用于查询和键
# 在注意力头部分离后，但在点积和后续softmax（注意力机制）之前进行
q_rope = rotary_emb.rotate_queries_or_keys(q)
k_rope = rotary_emb.rotate_queries_or_keys(k)

print("应用旋转位置编码后的查询张量 q_rope:", q_rope)
print("应用旋转位置编码后的键张量 k_rope:", k_rope)

3. 关键代码分析

        代码rotary_embedding_torch.py有点长，但我们可以重点关注以下几个部分：

3.1 旋转角度计算

        旋转角度决定了频率。或者说旋转位置编码中的频率用于定义位置编码的旋转角度。频率通常是由一种函数生成的，并且决定了位置编码如何在不同位置之间变化。在旋转位置编码中，频率通常与模型的维度和序列长度相关。频率的选择会影响旋转的位置编码如何在不同位置之间变化。为了处理更长的序列或优化模型性能，可能会使用动态频率调整技术，例如XPos。这些技术会根据当前处理的序列长度动态调整频率，从而优化模型的表现。

        源码中给出了多种频率计算方式：

if exists(custom_freqs):
    freqs = custom_freqs
elif freqs_for == 'lang':
    freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
elif freqs_for == 'pixel':
    freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
elif freqs_for == 'constant':
    freqs = torch.ones(num_freqs).float()

看下freqs_for == 'lang'，其实现基本和原论文【2】保持一致：

以d=32为例， torch.arange的输出为：

tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18., 20., 22., 24., 26.,28., 30.])

与预期输出保持一致。

3.2 向量旋转

3.2.1 实现分析

        首先还是回顾原论文中给出的旋转矩阵的高效实现方法：

 这个实现方法，涉及到两部分：

(1) 给定向量， 通过旋转得到新的序列

(2)给定向量, 扩展成, 同理对于sin部分也需要做类似翻倍扩展。

3.2.2 针对向量旋转

接下来我们来看看代码如何实现：

def rotate_half(x):
    x = rearrange(x, '... (d r) -> ... d r', r=2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return rearrange(x, '... d r -> ... (d r)')

• rearrange 是 einops 库中的一个函数，用于灵活地重塑、重排和转置张量。它的用法是通过指定模式来表达如何重新排列张量的形状。
• 输入x 是一个张量，其最后一个维度的大小必须是偶数（因为我们需要拆分成两个相等的部分）。
• 模式 '... (d r) -> ... d r' 表示将张量 x 的最后一个维度（即 (d r) 部分）拆分为两个独立的维度 d 和 r，其中 r=2。
• ... 表示保持前面的所有维度不变。
• 假设原始张量 x 的形状为 [..., 2 * n]，在重排后，它的形状将变为 [..., n, 2]，即最后一个维度被拆分为大小为 n 和 2 的两个新维度。
• unbind 函数沿着指定的维度（这里是 dim=-1，即最后一个维度）将张量分解为多个张量。由于前面重排后的张量 x 的最后一个维度是大小为 2，x.unbind(dim=-1) 将张量 x 沿最后一个维度拆分成两个张量。
• 输入 (-x2, x1) 是两个张量的元组，-x2 是对 x2 中的每个元素取负。
• torch.stack 沿着最后一个新维度（dim=-1）将两个张量堆叠，形成一个新张量 x，其形状为 [..., n, 2]。
• 结果是将前半部分（x1）和旋转后的后半部分（-x2）拼接在一起。

举个数值示例，假设对x = torch.tensor([1, 2, 3, 4]) 处理，最后rotary输出的结果为：[-2, 1, -4, 3]，符合我们的预期。

3.2.2 针对向量翻倍扩展

        使用einops的repeat方法能够快速实现

freqs = repeat(freqs, '... n -> ... (n r)', r=2)

举个数值示例，假设freqs = torch.tensor([1,2,3,4])， 输出为tensor([1, 1, 2, 2, 3, 3, 4, 4])。

3.3 缩放因子

        源代码中实现了根据序列输入的长度，动态调整缩放因子。目的是优化模型在处理长序列时的性能，核心思想是动态调整旋转位置编码中的旋转角度。动态调整位置编码，能够根据当前输入的序列长度进行动态调整，使得位置编码适应不同长度的输入。

        有的处理方式是通过调整底数来适应不同长度的序列【5】。先来理解一下这种模式：

位置编码的频率计算：RoPE 中的频率通常是根据以下公式计算的：

其中，i 是维度索引，d 是总的维度，base 是一个常数。

底数 base 的作用：当 base 增大时，freqs 中的频率将减少，角度的旋转变化会变得更缓慢；反之，base 减小时，freqs 中的频率增加，旋转角度的变化变得更快。底数控制了位置编码变化的快慢。

当输入序列长度变得更长时，使用固定的频率（或者说固定的 base）可能不够灵活。动态调整 base，可以让模型适应长序列带来的挑战。

短序列：当序列较短时，位置编码频率需要相对高一些，以便更精细地捕捉相对位置信息。此时，较小的 base 适合，因为它能提供更高的频率变化。

长序列：当序列变长时，使用同样高的频率变化可能导致模型的表示能力不足，因为它对更远距离的位置信息缺乏适应能力。此时，需要较大的 base，以减缓频率变化，让模型在更长的范围内有效利用位置信息。

因此，通过动态调整 base，我们可以使得频率变化适应当前的序列长度。对于更长的序列，可以选择更大的 base，让 RoPE 的角度变化更缓慢，避免位置编码过于快速变化导致模型难以捕捉远距离关系；对于较短的序列，选择较小的 base，以提供足够的频率分辨率。

在源码中，并没有直接修改底数，而是乘上缩放因子。

scale = 1.
if self.use_xpos:
    power = (t - len(t) // 2) / self.scale_base
    scale = self.scale ** rearrange(power, 'n -> n 1')
    scale = torch.cat((scale, scale), dim=-1)


#作用scale
(t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)

在llama的源码【6】中，也可以看到类似的动态调整freq的代码：

3.4 旋转位置编码变换

        有了上述关键模块，就可以实施q、k的旋转变换。

@autocast(enabled=False)
def apply_rotary_emb(freqs, t, start_index=0, scale=1., seq_dim=-2):
    dtype = t.dtype

    if t.ndim == 3:
        seq_len = t.shape[seq_dim]
        freqs = freqs[-seq_len:]

    rot_dim = freqs.shape[-1]
    end_index = start_index + rot_dim

    assert rot_dim <= t.shape[
        -1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'

    # 将t分为三部分：左、中（要变换的部分）和右
    t_left = t[..., :start_index]
    t_middle = t[..., start_index:end_index]
    t_right = t[..., end_index:]

    # 应用旋转嵌入
    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)

    out = torch.cat((t_left, t_transformed, t_right), dim=-1)

    return out.type(dtype)

 对应t_transformed代码来理解下图：

4. 参考材料

【1】lucidrains/rotary-embedding-torch

【2】RoFormer: Enhanced Transformer with Rotary Position Embedding

【3】A Length-Extrapolatable Transformer

【4】Extending Context Window of Large Language Models via Positional Interpolation

【5】llama3代码阅读

【6】llama github源码