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 freqsrope_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源码
![[嵌入式 C 语言] int main(int argc, char *argv[])](https://i-blog.csdnimg.cn/direct/2d5cab83e4fd48f2b33a31ca7ec83c3a.png)
