LLaMA中ROPE位置编码实现源码解析

1、Attention中q,经下式,生成新的q。m为句长length,d为embedding_dim/head
θ i = 1 1000 0 2 i d \theta_i=\frac{1}{10000^\frac{2i}{d}} θi=10000d2i1

2、LLaMA中RoPE源码

python 复制代码
import torch

def precompute_freqs_cis(dim: int, end: int, constant: float = 10000.0):
    '''
    计算cos和sin的值,cos值在实部,sin值在虚部,类似于 cosx+j*sinx
    :param dim: q,k,v的最后一维,一般为emb_dim/head_num
    :param end: 句长length
    :param constant: 这里指10000
    :return:
    复数计算 torch.polar(a, t)输出, a*(cos(t)+j*sin(t))
    '''
    # freqs: 计算 1/(10000^(2i/d) ),将结果作为参数theta
    # 形式化为 [theta_0, theta_1, ..., theta_d/2]
    freqs = 1.0 / (constant ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) # [d/2]

    # 计算m
    t = torch.arange(end, device=freqs.device)  # [length]
    # 计算m*theta
    freqs = torch.outer(t, freqs).float()  # [length, d]
    # freqs形式化为 [m*theta_0, m*theta_1, ..., m*theta_d/2],其中 m=0,1,...,length-1

    # 计算cos(m*theta)+j*sin(m*theta)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    # freqs_cis: [cos(m*theta_0)+j*sin(m*theta_0),  cos(m*theta_1)+j*sin(m*theta_1),), ..., cos(m*theta_d/2-1)+j*sin(m*theta_d/2-1)]
    # 其中j为虚数单位, m=0,1,...,length-1
    return freqs_cis # [length, d]

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] # (1, length, 1, d/2)
    return freqs_cis.view(*shape) # [1, length, 1, d/2]

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor,):
    # 先将xq维度变为[bs, length, head,  d/2, 2], 利用torch.view_as_complex转变为复数
    # xq:[q0, q1, .., q(d-1)] 转变为 xq_: [q0+j*q1, q2+j*q3, ..., q(d-2)+j*q(d-1)]
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [bs, length, head, d/2]
    # 同样的,xk_:[k0+j*k1, k2+j*k3, ..., k(d-2)+j*k(d-1)]
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    freqs_cis = reshape_for_broadcast(freqs_cis, xq_) # [1, length, 1, d/2]
    # 下式xq_ * freqs_cis形式化输出,以第一个为例, 如下
    # (q0+j*q1)(cos(m*theta_0)+j*sin(m*theta_0)) = q0*cos(m*theta_0)-q1*sin(m*theta_0) + j*(q1*cos(m*theta_0)+q0*sin(m*theta_0))
    # 上式的实部为q0*cos(m*theta_0)-q1*sin(m*theta_0),虚部为q1*cos(m*theta_0)+q0*sin(m*theta_0)
    # 然后通过torch.view_as_real函数,取出实部和虚部,维度由[bs, length, head, d/2]变为[bs, length, head, d/2, 2],最后一维放实部与虚部
    # 最后经flatten函数将维度拉平,即[bs, length, head, d]
    # 此时xq_out形式化为 [实部0,虚部0,实部1,虚部1,..., 实部(d/2-1), 虚部(d/2-1)]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # [bs, length, head, d]
    # 即为新生成的q

    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

if __name__=='__main__':
    # (bs, length, head, d)
    q = torch.randn((2, 10, 12, 32))  # q=[q0, q1, .., qd-1]
    k = torch.randn((2, 10, 12, 32))
    v = torch.randn((2, 10, 12, 32))
    freqs_cis= precompute_freqs_cis(dim=32, end=10, constant= 10000.0)
    # print(freqs_cis.detach().numpy())
    
    q_new, k_new = apply_rotary_emb(xq=q, xk=k, freqs_cis=freqs_cis)
    print()
相关推荐
Ven%3 分钟前
如何让后台运行llamafactory-cli webui 即使关掉了ssh远程连接 也在运行
运维·人工智能·chrome·python·ssh·aigc
Jeo_dmy8 分钟前
(七)人工智能进阶之人脸识别:从刷脸支付到智能安防的奥秘,小白都可以入手的MTCNN+Arcface网络
人工智能·计算机视觉·人脸识别·猪脸识别
睡觉狂魔er2 小时前
自动驾驶控制与规划——Project 5: Lattice Planner
人工智能·机器学习·自动驾驶
xm一点不soso2 小时前
ROS2+OpenCV综合应用--11. AprilTag标签码跟随
人工智能·opencv·计算机视觉
云卓SKYDROID3 小时前
无人机+Ai应用场景!
人工智能·无人机·科普·高科技·云卓科技
是十一月末3 小时前
机器学习之过采样和下采样调整不均衡样本的逻辑回归模型
人工智能·python·算法·机器学习·逻辑回归
小禾家的3 小时前
.NET AI 开发人员库 --AI Dev Gallery简单示例--问答机器人
人工智能·c#·.net
生信碱移3 小时前
万字长文:机器学习的数学基础(易读)
大数据·人工智能·深度学习·线性代数·算法·数学建模·数据分析
KeyPan3 小时前
【机器学习:四、多输入变量的回归问题】
人工智能·数码相机·算法·机器学习·计算机视觉·数据挖掘·回归
码力全開3 小时前
C 语言奇幻之旅 - 第14篇:C 语言高级主题
服务器·c语言·开发语言·人工智能·算法