Nano-VLLM全代码解析笔记(10)-GemmaRMSNorm和MRoPE-平芜编程栈

前言

本节参考资料：

VLLM源码

Transformers源码

(49 封私信) Qwen3.5 架构最全拆解：Linear Attention 源码配图解析、Gated DeltaRule 公式源码逻辑介绍、Full Attention与 MoE 模块算子流程解析 - 知乎

我的代码仓库(迭代更新ing)

WilliamPockey/Nano-Vllm-Qwen-Fit

TODO LIST

（已完成）Qwen3.5架构介绍

（已完成）GemmaRMSNorm和MRoPE

（正在完成）GDN线性注意力层

（待完成）多模态支持（视觉塔/多模态预处理/多模态条件生成类/权重加载）

（待完成）引擎（LLMEngine/Sequence/Scheduler/ModelRunner）修改

（待完成）MTP支持

（待完成）流式输出与多请求支持

Qwen3.5架构图

本节修改的代码

layernorm.py

更改说明：增加了GemmaRMSNorm类和RMSNormGated类，前者是Qwen3.5-0.8B使用的对修改后的RMSNorm，主要是把（y=x/RMS(x)×γ变为了（y=x/RMS(x)×(γ+1)）），简单修改即可。而后者是Qwen3.5-0.8B的Gated DeltaNet所需的模块，跟正常的RMSNorm区别不大，会在介绍这个网络的时候用到，他这里的silu区别于activation.py中的silu，因为这里的gate和hidden_state不是像之前可以一起得到的

class GemmaRMSNorm(nn.Module): """RMS normalization for Gemma. difference from the above RMSNorm: 1. x * (1 + w) instead of x * w. """ def __init__( self, hidden_size: int, eps: float = 1e-6, ) -> None: super().__init__() self.weight = nn.Parameter(torch.zeros(hidden_size)) self.eps = eps @torch.compile def rms_forward( self, x: torch.Tensor, weight: torch.Tensor, ) -> torch.Tensor: orig_dtype = x.dtype x = x.float() var = x.pow(2).mean(dim=-1, keepdim=True) x.mul_(torch.rsqrt(var + self.eps)) x = x.to(weight.dtype).mul_(weight) return x.to(orig_dtype) @torch.compile def add_rms_forward( self, x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: orig_dtype = x.dtype x = x.float().add_(residual.float()) residual = x.to(orig_dtype) var = x.pow(2).mean(dim=-1, keepdim=True) x.mul_(torch.rsqrt(var + self.eps)) x = x.to(weight.dtype).mul_(weight) return x.to(orig_dtype), residual def forward( self, x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" weight = self.weight.float() + 1.0 if residual is None: return self.rms_forward(x, weight) return self.add_rms_forward(x, residual, weight) class RMSNormGated(nn.Module): def __init__( self, hidden_size: int, eps: float = 1e-6, **kwargs ) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.activation = "silu" def forward( self, hidden_states: torch.Tensor, gate: torch.Tensor ) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) hidden_states = hidden_states * nn.functional.silu(gate.to(torch.float32)) return hidden_states.to(input_dtype)

rotary_embedding.py

更改说明：Qwen3.5系列采用了部分旋转位置编码(partial RoPE)和MRoPE的策略，因此我们需要修改原来的代码，主要是改动了apply_rotary_emb函数和RotaryEmbedding类以适配部分旋转位置编码，新增了MRotaryEmbedding类以支持MRoPE。这里没有对MRoPE进行缓存，主要是考虑到多维的旋转位置编码会导致位置编码的变化。

MRoPE介绍：【面试高频】M-RoPE 多模态位置编码全解

简单而言，就是对token的最后一维hidden_state，之前每一个地方的向量旋转角度只有位置决定，我们现在要把它拆成3块，每一块的向量旋转角度由对应的3个维度和维度中的位置决定(T时间\H高度\W宽度)

from functools import lru_cache import torch from torch import nn # def apply_rotary_emb( # x: torch.Tensor, # cos: torch.Tensor, # sin: torch.Tensor, # ) -> torch.Tensor: # x1, x2 = torch.chunk(x.float(), 2, dim=-1) # y1 = x1 * cos - x2 * sin # y2 = x2 * cos + x1 * sin # return torch.cat((y1, y2), dim=-1).to(x.dtype) #新版本部分位置旋转兼容旧版本全位置旋转 def apply_rotary_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> torch.Tensor: rotary_dim = cos.size(-1) * 2 # 现有 cache/现算路径都给 32 → 64 x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] x1, x2 = torch.chunk(x_rot.float(), 2, dim=-1) y1 = x1 * cos - x2 * sin y2 = x2 * cos + x1 * sin return torch.cat((y1, y2, x_pass), dim=-1).to(x.dtype) class RotaryEmbedding(nn.Module): def __init__( self, head_size: int, rotary_dim: int, max_position_embeddings: int, base: float, ) -> None: super().__init__() self.head_size = head_size # assert rotary_dim == head_size assert rotary_dim <= head_size and rotary_dim % 2 == 0 inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) t = torch.arange(max_position_embeddings, dtype=torch.float) freqs = torch.einsum("i,j -> ij", t, inv_freq) # 外积：[位置数, 频率数] cos = freqs.cos() sin = freqs.sin() cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1) self.register_buffer("cos_sin_cache", cache, persistent=False) @torch.compile def forward( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: cos_sin = self.cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) query = apply_rotary_emb(query, cos, sin) key = apply_rotary_emb(key, cos, sin) return query, key class MRotaryEmbedding(nn.Module): def __init__( self, head_size: int, rotary_dim: int, base: float, mrope_section=(11, 11, 10) ) -> None: super().__init__() self.head_size = head_size self.rotary_dim = rotary_dim inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) # 预计算每个频率列取自 T/H/W 哪一路：默认 0(T)，H 占 slice(1, s1*3, 3)，W 占 slice(2, s2*3, 3) sel = torch.zeros(rotary_dim // 2, dtype=torch.long) sel[1 : mrope_section[1] * 3 : 3] = 1 sel[2 : mrope_section[2] * 3 : 3] = 2 self.register_buffer("freq_axis_sel", sel, persistent=False) def forward( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: # positions: (3, N) freqs: (3, N, 32) freqs = positions.float().unsqueeze(-1) * self.inv_freq #freqs.transpose(0, 1): (N, 3, 32) #index：sel (32,) → view → (1, 1, 32) → expand → (N, 1, 32) #freqs_t: (N, 32) freqs_t = freqs.transpose(0, 1).gather(1, self.freq_axis_sel.view(1, 1, -1).expand(freqs.size(1), 1, -1)).squeeze(1) #cos: (N, 1, 32) cos = freqs_t.cos().unsqueeze_(1) sin = freqs_t.sin().unsqueeze_(1) return apply_rotary_emb(query, cos, sin), apply_rotary_emb(key, cos, sin)

几个问题

1.self.register_buffer("freq_axis_sel", sel, persistent=False)，这里最后的参数是什么意思 ❌ 保存模型时不会保存这个 buffer ❌ 加载模型时不会恢复这个 buffer ❌ model.state_dict() 不包含这个 buffer ✅ model.to(device) 仍然会移动这个 buffer（因为存在内存中） 但如果设为True，前三个则反过来，最后一个不变 2.为什么foward前不能加@torch.compile 因为@torch.compile对gather不友好 3.为什么不将MRoPE进行缓存 这里主要是考虑到输入的不同，比如我输入文本、输入文本和图片、输入文本和视频都会导致位置编码变化。 而且图片的输入长宽会变，MRoPE就算考虑设置最大长宽，也会导致要保存的cos_sin太大 所以缓存不一定会命中，所以暂时不考虑，等后续优化时再看，这里可能把缓存队列开大一些可以支持特定情况下的缓存 4.mrope_section为什么是(11,11,10) 这里是考虑到rotary_dim是32维，而MRoPE简单来说就是让32维向量拆分成3个维度，每个维度单独使用ROPE 因此这里(11,11,10)是人为设计 5.解释freqs = positions.float().unsqueeze(-1) * self.inv_freq freqs_t = freqs.gather(0, self.freq_axis_sel.unsqueeze(0).expand(freqs[0].shape)) MRoPE的三个维度是（时间，高度，宽度），token也被要求给出自己这三个维度对应的数值 因此freqs -> freqs_t可以理解为，根据一个 token的每一个 hidden_state 中的每一个位置的值选 T/H/C 中的一个频率 gather介绍 # 如果 dim=0 out[i][j][k] = input[ index[i][j][k] ][ j ][ k ] # 如果 dim=1 out[i][j][k] = input[ i ][ index[i][j][k] ][ k ] # 如果 dim=2 out[i][j][k] = input[ i ][ j ][ index[i][j][k] ] 这个公式看起来可能有点抽象，它的具体工作过程是这样的： 确定输出位置：输出张量 out 的形状和 index 一模一样。我们会遍历 out 中的每一个位置（比如 (i, j, k)）。 读取索引值：看 index 在同样位置 (i, j, k) 上的数值是多少。 替换对应维度的索引： 如果 dim=0，意味着我们要替换的是第一个维度的索引。所以，输出位置 (i, j, k) 的值，来自于 input 中位置 ( index[i][j][k], j, k ) 的元素。 如果 dim=1，就替换第二个维度的索引，取值位置变为 ( i, index[i][j][k], k )。以此类推。 gather的算法等价于 # freqs_t: [N,32] 输出 freqs_t = torch.empty(N, 32) for n in range(N): # 遍历每个token for i in range(32): # 遍历每个频率（对应一对hidden） axis = freq_axis_sel[i] # 0=T,1=H,2=W，只看i，不看n freqs_t[n,i] = positions[axis, n] * inv_freq[i] 因此如果是单模态，比如说纯文本，让T=H=C，这样就退化成了ROPE 6.我的疑惑是对于 hidden_state，他选了通道 [0,1,2,0,1,2...,0,1,2]，apply 里面是 chunk 拆成了两份，那就成了 [0,1,2,0,1,2,0,1,2,0,1,2,0,1,2,0] 和 [1,2,0,1,2,0,1,2,0,1,2,0,1,2,0,1]，然后再是旋转，问题是比如说第一个分量，分别是 0 和 1，这都不是一个通道，怎么能进行旋转？我的意思是 x_rot 里面存的是 [0,1,2,0,1,2] 对应的不同通道的频率，但我们 apply_rotary_embed 相当于对不同通道的频率进行计算 1.`i = 0,1,2,3,…31`：**频率索引** - `freq_axis_sel[i] ∈ {0,1,2}`：决定**第 i 号频率**用 T/H/W 哪一套位置算出角度。 - 输出：`cos[...,i]`、`sin[...,i]`：这是**第 i 号频率的旋转角度系数**。 2. `x_rot`（hidden 向量，shape `[…,64]`） 物理排布硬规则（RoPE 原生约定，M‑RoPE 不改）： ``` x_rot[..., 0], x_rot[..., 1] ← 配对，使用 i=0 的cos/sin x_rot[..., 2], x_rot[..., 3] ← 配对，使用 i=1 的cos/sin x_rot[..., 4], x_rot[..., 5] ← 配对，使用 i=2 的cos/sin x_rot[..., 6], x_rot[..., 7] ← 配对，使用 i=3 的cos/sin …… x_rot[...,2*i], x_rot[...,2*i+1] ← 配对，使用 i号频率的cos/sin ``` > > 重点！ > hidden 上**位置 2i、2i+1 这一对，固定绑定第 i 号频率的 cos/sin**。 > 这个绑定关系是按数组下标位置硬绑定，**不会因为 sel [i]=0/1/2 发生任何改变**。 > `sel[i]` 仅仅改变：**i 号频率的 cos/sin 值是拿 T 算出来，还是 H、还是 W 算出来**。它绝不调换 hidden 配对关系。 举极简小例子，F=3（i=0,i=1,i=2） - `freq_axis_sel = [0, 1, 2]` - i=0：选 T 位置算角度 → `cos0, sin0` - i=1：选 H 位置算角度 → `cos1, sin1` - i=2：选 W 位置算角度 → `cos2, sin2` `x_rot = [a0,a1, b0,b1, c0,c1]` - `a0,a1`：hidden 的 0、1 号位，**强制使用 i=0 的 (cos0,sin0)**（T 角度） - `b0,b1`：hidden 的 2、3 号位，**强制使用 i=1 的 (cos1,sin1)**（H 角度） - `c0,c1`：hidden 的 4、5 号位，**强制使用 i=2 的 (cos2,sin2)**（W 角度）