rope编码代码分享

python 复制代码
from typing import Tuple
import torch

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    """
    Helper function to reshape frequency tensor to have the same shape as the target tensor 'x'
    for the purpose of broadcasting the frequency tensor during element-wise operations.

    Args:
        freqs_cis (torch.Tensor): Frequency tensor to be reshaped.
        x (torch.Tensor): Target tensor for broadcasting compatibility.

    Returns:
        torch.Tensor: Reshaped frequency tensor.

    Raises:
        AssertionError: If the frequency tensor doesn't match the expected shape.
        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
    """
    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)]
    return freqs_cis.view(shape)
#########填充维度,方便计算

def apply_rotary_emb(
        query: torch.Tensor,
        key: torch.Tensor,
        head_dim: int,
        max_seq_len: int,
        theta: float = 10000.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Apply rotary embeddings to input tensors using the given frequency tensor.

    Args:
        query (torch.Tensor): Query tensor to apply rotary embeddings. Shape: (batch_size, seqlen, n_local_heads, head_dim)
        key (torch.Tensor): Key tensor to apply rotary embeddings. Shape: (batch_size, seqlen, n_local_kv_heads, head_dim)
        head_dim (int): Dimension of each attention head.
        max_seq_len (int): Maximum sequence length supported by model.
    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
    """

    _, seqlen, _, _ = query.shape  # 获取查询张量的形状参数
    device = query.device  # 获取查询张量的设备信息(如在 CPU 或 GPU 上)
    seq_len, batch_size, num_heads = query.size(1), query.size(0), query.size(2)  # 获取序列长度、批次大小和头部数量

    # reshape xq and xk to match the complex representation
    query_real, query_imag = query.float().reshape(query.shape[:-1] + (-1, 2)).unbind(-1)  # 将查询张量重塑并分为实部和虚部
    key_real, key_imag = key.float().reshape(key.shape[:-1] + (-1, 2)).unbind(-1)  # 将键张量重塑并分为实部和虚部

    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2.0, device=device) / head_dim))
    pos_seq = torch.arange(0, seqlen, device=device)
    sinusoid_inp = torch.einsum("i,j->ij", pos_seq, inv_freq)  #
    sin = torch.sin(sinusoid_inp)
    cos = torch.cos(sinusoid_inp)

    # Use the reshape_for_broadcast function to reshape cos and sin terms for broadcasting
    cos_rotations = reshape_for_broadcast(cos, query_real)  # 调整余弦值张量的形状以进行广播
    sin_rotations = reshape_for_broadcast(sin, query_imag)  # 调整正弦值张量的形状以进行广播

    # Apply the rotations to the real and imaginary parts
    query_rot_real = cos_rotations * query_real - sin_rotations * query_imag  # 应用旋转到查询张量的实部
    query_rot_imag = sin_rotations * query_real + cos_rotations * query_imag  # 应用旋转到查询张量的虚部
    key_rot_real = cos_rotations * key_real - sin_rotations * key_imag  # 应用旋转到键张量的实部
    key_rot_imag = sin_rotations * key_real + cos_rotations * key_imag  # 应用旋转到键张量的虚部

    # Reassemble the real and imaginary parts back into the original format
    # query_out = torch.cat([query_rot_real, query_rot_imag], dim=-1).view_as(query)  # 重新组合并调整查询张量的形状
    # key_out = torch.cat([key_rot_real, key_rot_imag], dim=-1).view_as(key)  # 重新组合并调整键张量的形状

    query_out = torch.stack((query_rot_real, query_rot_imag), dim=-1).flatten(-2)
    key_out = torch.stack((key_rot_real, key_rot_imag), dim=-1).flatten(-2)


    return query_out, key_out  # 返回包含旋转位置嵌入的查询和键张量

上述代码和b站这个up讲的,或者一般的rope代码有两点不同

1,q0,q1,q2,q3...和用两个相同cos,sin张量堆叠起来的新张量点乘的操作,变成先将张量q分离成q0,q2,q4...和q1,q3,q5...两个张量去和相同的cos,sin张量点乘

2,补全张量维度由代码

python 复制代码
cos_cached idx_theta2.cos()[:,None,None,:]
sin_cached idx_theta2.sin()[:,None,None,:]

变成

python 复制代码
cos_rotations = reshape_for_broadcast(cos, query_real)  # 调整余弦值张量的形状以进行广播
sin_rotations = reshape_for_broadcast(sin, query_imag)  # 调整正弦值张量的形状以进行广播
相关推荐
橘子编程26 分钟前
Python-Word文档、PPT、PDF以及Pillow处理图像详解
开发语言·python
蓝婷儿41 分钟前
Python 机器学习核心入门与实战进阶 Day 2 - KNN(K-近邻算法)分类实战与调参
python·机器学习·近邻算法
之歆1 小时前
Python-封装和解构-set及操作-字典及操作-解析式生成器-内建函数迭代器-学习笔记
笔记·python·学习
埃菲尔铁塔_CV算法2 小时前
基于 TOF 图像高频信息恢复 RGB 图像的原理、应用与实现
人工智能·深度学习·数码相机·算法·目标检测·计算机视觉
天天爱吃肉82182 小时前
ZigBee通信技术全解析:从协议栈到底层实现,全方位解读物联网核心无线技术
python·嵌入式硬件·物联网·servlet
Allen_LVyingbo3 小时前
Python常用医疗AI库以及案例解析(2025年版、上)
开发语言·人工智能·python·学习·健康医疗
智能砖头3 小时前
LangChain 与 LlamaIndex 深度对比与选型指南
人工智能·python
中杯可乐多加冰4 小时前
【AI落地应用实战】AIGC赋能职场PPT汇报:从效率工具到辅助优化
人工智能·深度学习·神经网络·aigc·powerpoint·ai赋能
烟锁池塘柳04 小时前
【大模型】解码策略:Greedy Search、Beam Search、Top-k/Top-p、Temperature Sampling等
人工智能·深度学习·机器学习
风逸hhh4 小时前
python打卡day58@浙大疏锦行
开发语言·python