llama源码学习·model.py[3]ROPE旋转位置编码(4)ROPE的应用

一、源码注释

python 复制代码
def apply_rotary_emb(
    xq: torch.Tensor, # 查询矩阵
    xk: torch.Tensor, # 键矩阵
    freqs_cis: torch.Tensor, # 旋转嵌入
) -> Tuple[torch.Tensor, torch.Tensor]:
    
    # 首先将xq和xk张量转换为浮点数
    # 然后使用reshape将最后一个维度拆分为两个维度,每个维度都有大小为2,这样做是为了为复数张量提供实部和虚部。
    # 然后,torch.view_as_complex用于从实部和虚部创建复数张量
    
    # *xq.shape[:-1] 是保留原始形状的所有维度,除了最后一个维度。
    # -1 是一个占位符,它告诉PyTorch自动计算这个维度,以保持元素总数不变。
    # 2 是最后一个维度,这是为了为接下来的复数转换做准备。每个复数由两个浮点数表示(实部和虚部),所以最后一个维度是2。
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    
    # 将freqs_cis重新reshape以匹配xq_的形状,以便进行广播运算。
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    
    # 这两行代码将查询和键张量与旋转嵌入相乘,应用位置嵌入。
    # 函数计算xq_和xk_与freqs_cis的元素乘积(这是一个复数乘法),
    # 在复数乘法中,xq_和xk_的实部和虚部会分别与freqs_cis的实部和虚部进行乘法运算。
    # flatten(3) 将两个最后的维度合并回一个维度
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    
    # 函数返回经过旋转嵌入处理的查询和键张量,同时确保它们的数据类型与原始输入相匹配。
    return xq_out.type_as(xq), xk_out.type_as(xk)

二、举例说明

python 复制代码
# query矩阵
xq = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])  
# key矩阵
xk = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
# 频率张量
freqs_cis = torch.tensor([[1.0000+0.0000j], [1.0000+0.0000j]])  

*** xq.shape: *** torch.Size([2, 2, 2])

*** xk.shape: *** torch.Size([2, 2, 2])

freqs_cis.shape: torch.Size([2, 1])

python 复制代码
# 首先,apply_rotary_emb函数会将query和key矩阵reshape并转换为复数张量。
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) 
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

xq.float().reshape(*xq.shape[:-1], *-*1, 2).shape: torch.Size([2, 2, 1, 2])

xk.float().reshape(*xk.shape[:-1], *-*1, 2).shape: torch.Size([2, 2, 1, 2])

xq_.shape: torch.Size([2, 2, 1])

xk_.shape: torch.Size([2, 2, 1])

python 复制代码
# freqs_cis 的形状是 (2, 1),xq_ 的形状是(2, 2, 1), 所以我们需要将freqs_cis形状调整为 (1, 2, 2, 1)
freqs_cis_new = reshape_for_broadcast(freqs_cis, xq_)

freqs_cis_new.shape : freqs_cis_new.shape

python 复制代码
# 函数会将输入复数张量与频率张量相乘。
xq_out_complex = xq_ * freqs_cis_new
xk_out_complex = xk_ * freqs_cis_new

xq_out_complex: tensor([[[1.+2.j], [3.+4.j]], [[5.+6.j], [7.+8.j]]])

python 复制代码
# 将结果重塑并转换回实数张量。
xq_out = torch.view_as_real(xq_out_complex).flatten(3)
xk_out = torch.view_as_real(xk_out_complex).flatten(3)

xq_out: tensor([[[[ 1., 2.], [ 6., 8.]], [[15., 18.], [28., 32.]]]])

相关推荐
try2find2 天前
安装llama-cpp-python踩坑记
开发语言·python·llama
西西弗Sisyphus3 天前
LLaMA-Factory 单卡后训练微调Qwen3完整脚本
微调·llama·llama-factory·后训练
顾道长生'3 天前
(Arxiv-2024)自回归模型优于扩散:Llama用于可扩展的图像生成
计算机视觉·数据挖掘·llama·自回归模型·多模态生成与理解
Zhijun.li@Studio13 天前
【LLaMA-Factory 实战系列】二、WebUI 篇 - Qwen2.5-VL 多模态模型 LoRA 微调保姆级教程
人工智能·自然语言处理·llama·多模态大模型
1213413 天前
LLM:重构数字世界的“智能操作系统”
gpt·aigc·ai编程·llama·gpu算力
冷雨夜中漫步21 天前
Java中如何使用lambda表达式分类groupby
java·开发语言·windows·llama
扫地的小何尚22 天前
全新NVIDIA Llama Nemotron Nano视觉语言模型在OCR基准测试中准确率夺冠
c++·人工智能·语言模型·机器人·ocr·llama·gpu
CFAteam23 天前
DeepSeek AI功能演示:如何生成Verilog脚本
人工智能·ai·fpga开发·llama
Tadas-Gao25 天前
从碳基羊驼到硅基LLaMA:开源大模型家族的生物隐喻与技术进化全景
人工智能·机器学习·大模型·llm·llama
Run_Clover25 天前
llama-factory微调大模型环境配置避坑总结
llama