【声明】本博客所有内容均为个人业余时间创作所述技术案例均来自公开开源项目如GithubApache基金会不涉及任何企业机密或未公开技术如有侵权请联系删除标题247、【AI】【模型部署】基座模型研究RoPE 用在前向还是反向——训练、推理与 KV cache背景上篇 blog【AI】【模型部署】基座模型研究RoPE 的维度与频率——32 对、多频率与 rope_theta把频率讲透了第i ii对的角速度是inv_freq [ i ] 1 / θ 2 i / d \text{inv\_freq}[i]1/\theta^{2i/d}inv_freq[i]1/θ2i/dhead_dim64得到 32 个频率它们按几何级数排开、低维转得快近处、高维转得慢远处波长2 π θ 2 i / d 2\pi\,\theta^{2i/d}2πθ2i/d随i ii指数增长最慢的波长决定能分辨多远rope_theta越大波长越长、量程越宽Qwen2 取10 6 10^6106是为长上下文还提到 partial rotary 与位置 0 的 cos 全是 1。到这里RoPE 的数学245点积只依赖n − m n-mn−m和频率设计24632 对、多尺度都清楚了。本篇从数学转到工程回答一个很实际的问题旋转到底发生在哪个阶段——前向还是反向训练还是推理它和243篇的 KV cache 又是怎么配合的模型部署先给结论旋转只发生在每一层注意力的前向里训练和推理都做而且它没有任何可训练参数。这四个关键词——前向、训练、推理、无参数——每个都值得展开尤其是和 KV cache 的配合因为那正是243篇留下的接口。理解了这一篇227篇里那些apply_rotary_pos_emb的调用时机、以及为什么缓存的是旋转后的 K就都串起来了。旋转发生在前向RoPE 不是训练前预处理也不是训练后补丁而是每一层注意力前向计算的一部分拿到q , k q,kq,k之后、算q ⋅ k q\cdot kq⋅k之前按位置把它们旋转一下。计算图上是这样一条链图 1 上半部分是前向q , k q,kq,k先乘上R ( m θ ) R(m\theta)R(mθ)完成旋转再送入注意力、算出损失。没有位置参数、没有查表、没有预处理开关——旋转就是一次固定的乘加藏在注意力里面。它有没有可训练参数没有。245篇的旋转角由位置和频率算出来频率又是inv_freq 1/θ^{2i/d}这条公式现算的在实现里作为 buffer 存在不是nn.Parametercos/sin 也可以预先算好缓存。所以优化器不会、也不该去学习旋转角反向传播里没有RoPE 的权重梯度这一项训练前后RoPE 的这部分计算完全不变——它是固定的几何操作不是学出来的层。这一点常被误解成RoPE 是某种可学习的位置嵌入。不是它是公式驱动的旋转参数只有配置层面的rope_theta超参数不参与梯度更新。反向传播时发生什么虽然 RoPE 没有参数但它是可微的梯度会穿过它。旋转是线性正交变换它的转置就是反向旋转如果前向是乘R ( m θ ) R(m\theta)R(mθ)那么反向回传的梯度就乘R ( m θ ) ⊤ R ( − m θ ) R(m\theta)^\topR(-m\theta)R(mθ)⊤R(−mθ)。回到图 1 的下半部分前向按位置把q , k q,kq,k转过去反向只是把梯度按同样角度反向转回来继续传给前面的层。因为旋转保长度245 245245篇这一步既不会放大也不会缩小梯度非常干净。所以RoPE 用在反向这个说法要修正反向里发生的是梯度的反向旋转不是位置编码本身被更新。训练整段并行训练时序列是完整已知的位置编号就是0 , 1 , … , L − 1 0,1,\dots,L-10,1,…,L−1。注意力对所有位置并行计算RoPE 也一次性把整段的q , k q,kq,k都旋转好——每个位置用它自己的角。这和在序列上逐个处理在数学上等价只是训练可以并行、效率更高。图 2 左边画的就是训练以及推理里的 prefill所有位置一起处理位置编号从 0 排到L − 1 L-1L−1。推理prefill 与 decode推理分成两段243篇讲过prefill一次并行处理 prompt 的全部位置编号0 00到L − 1 L-1L−1decode之后逐个生成 token。RoPE 在两段都做区别只是位置编号prefill和训练一样一次性旋转整段decode每生成一个 token就用它在整段里的位置编号等于已经有多少个 token旋转这个新 token 的q , k q,kq,k。图 2 右边是 decode位置编号继续L , L 1 , L 2 , … L, L1, L2,\dotsL,L1,L2,…递增。关键是无论训练还是推理同一个位置永远用同一个角度、同一个公式——这一点后面会再强调。和 KV cache 怎么配合这是本篇的核心。243篇说 KV cache 缓存的是k , v k,vk,v那么缓存里的k kk到底是旋转前还是旋转后的答案是旋转后的。原因藏在相对位置里。设缓存里第n nn个位置的键已经旋转成R ( n θ ) k R(n\theta)kR(nθ)k当前新 token 在位置t tt它的查询旋转成R ( t θ ) q R(t\theta)qR(tθ)q。两者做点积( R ( t θ ) q ) ⊤ ( R ( n θ ) k ) q ⊤ R ( ( t − n ) θ ) k \big(R(t\theta)q\big)^\top\big(R(n\theta)k\big)q^\top R\big((t-n)\theta\big)k(R(tθ)q)⊤(R(nθ)k)q⊤R((t−n)θ)k只有( t − n ) (t-n)(t−n)正是想要的相对距离。如果缓存里存的是没旋转的原始k kk那每生成一步就得把缓存里所有历史K KK各自按它们的位置重新旋转一遍——缓存的意义就没了。存旋转后的 K才能做到新查询转一次、旧键直接拿来用。图 3 画了这个流程新 token 的k kk旋转后用当前位置追加进缓存它的q qq也旋转后去和缓存里全部已旋转的 K做点积。V VV不参与旋转245篇讲过原因原样缓存。243那句缓存k / v k/vk/v补全起来是缓存旋转后的k kk、未旋转的v vv。一次 decode 的完整数据流把上面几步串成一个 decode 步的清单就非常清楚了新 tokenx t x_txt进入这一层投影出q t , k t , v t q_t,k_t,v_tqt,kt,vt用位置编号t tt把q t , k t q_t,k_tqt,kt旋转v t v_tvt不动把旋转后的k t k_tkt和原样的v t v_tvt追加进缓存q t q_tqt与缓存里全部K KK做点积、softmax、对全部V VV加权求和得到输出。第 4、5 步就是243篇的缓存流程而本篇补上了缓存前先旋转、V VV不旋转这两个细节。每一层都重复这套动作位置编号在各层之间是一致的。训练与推理必须一致由此能推出一个实践上很重要的规则训练用的旋转方式推理时必须一模一样。具体来说位置编号要对齐训练0.. L − 1 0..L-10..L−1推理也从0 00开始连续递增不跳号、不重置频率、rope_theta、配对方式、要旋转的维度训练和推理必须相同不能在推理时把旋转关掉。一旦关掉模型看到的是没有位置的q , k q,kq,k与训练分布严重不符输出会明显退化。常见的坑也来自这里列几个错误做法后果推理时关闭旋转模型看不到位置输出严重退化缓存里存未旋转的 K只转新q qq新旧键不在同一坐标相对偏移错乱把已旋转的 K 再转一次位置角翻倍点积含义被破坏位置编号每步从 0 重来相对偏移整体错位长文本尤其明显训练与推理的rope_theta/维度不一致频率对不上等于换了套位置编码推理位置远超训练长度落入未见过的角度区间效果下降外推问题见#249/#250这些错误的共同点都是破坏了同一个位置用同一个角度这条约定。只要训练和推理严格共用同一套旋转参数与编号规则RoPE 就不会出问题。这套设计换来了什么回到最初的三个问题可以给出一个统一的回答前向还是反向——旋转本身在前向反向只是把梯度反向旋转一次、继续回传没有 RoPE 权重更新训练还是推理——两者都做且必须一致和 KV cache 的关系——缓存旋转后的 K、未旋转的 V使新查询一次旋转 旧键直接复用成为可能。好处也很清楚无参数不占权重、不增加训练变量、相对位置点积只依赖n − m n-mn−m、和缓存天然兼容旋转后的 K 可以直接存。这三条合起来就是 RoPE 成为主流位置编码的工程原因。三个常见疑问“缓存的 K 是旋转前还是后”后。这是为了让新查询旋转一次就能和历史键直接点积、自动得到相对距离存旋转前的 K 会逼着每步重算违背缓存初衷。“推理时新 token 的位置编号取多少”取它前面已有多少个 token。prompt 有L LL个 token那第一个生成 token 的位置是L LL依此类推。“RoPE 有参数是不是也算训练的一部分”它参与前向和反向的计算但没有可训练参数——rope_theta、频率都是配置或公式优化器不更新它们。训练到底改的还是W q , W k , W v W_q,W_k,W_vWq,Wk,Wv这些投影权重。一句话记忆RoPE 只发生在每层注意力的前向里训练和推理都做它没有可训练参数反向传播只是把梯度反向旋转乘R ⊤ R^\topR⊤后继续回传。训练整段并行、位置0.. L − 1 0..L-10..L−1推理按 prefill并行 decode逐 token给位置编号。和243的 KV cache 配合的关键是缓存里存的是旋转后的 K、未旋转的 V这样新查询转一次就能和历史键点积、自动得到相对偏移( t − n ) (t-n)(t−n)。训练与推理的旋转方式必须严格一致推理时不能关闭旋转。下一篇换一个方向不旋转还有哪些位置编码方案各自适合什么场景。OK本篇先到这里如有疑问欢迎评论区留言讨论祝各位功力大涨技术更上一层楼更多内容见下篇 blog