【声明】本博客所有内容均为个人业余时间创作所述技术案例均来自公开开源项目如GithubApache基金会不涉及任何企业机密或未公开技术如有侵权请联系删除标题226、【AI】【模型部署】基座模型研究RMSNorm 与残差归一化为什么必要背景上篇 blog【AI】【模型部署】基座模型研究从文字到向量分词器与嵌入层讲清了基座的入口Qwen2Tokenizer用 BPE 把文本切成子词hello是 1 个 token、模型部署是 2 个vocab.json151643 条加特殊 token 凑成vocab_size151936嵌入层是一张(151936, 896)的查表把 id 变成 896 维向量并与lm_head共享权重。向量从这里进入 24 层Qwen2DecoderLayer。224 拆出的骨架里每层第一件事是归一化 → 子层 → 残差本篇就从这个骨架的第一个零件读起Qwen2RMSNorm模型部署224 的Qwen2DecoderLayer.forward里self.input_layernorm出现在注意力之前、self.post_attention_layernorm出现在 MLP 之前——归一化是每层的开场动作。它到底做了什么看源码最直接。源码RMSNormRoot Mean Square Normalization均方根归一化 只有几行Qwen2RMSNormmodeling_qwen2.py:238-255的核心就三步实测打印源码classQwen2RMSNorm(nn.Module):def__init__(self,hidden_size,eps:float1e-6):self.weightnn.Parameter(torch.ones(hidden_size))self.variance_epsilonepsdefforward(self,hidden_states):variancehidden_states.to(torch.float32).pow(2).mean(-1,keepdimTrue)hidden_stateshidden_states*torch.rsqrt(varianceself.variance_epsilon)returnself.weight*hidden_states拆开看算平方均值 → 开根号取倒数 → 乘回去最后再乘一个可学习权重。没有减均值这一步——这正是它和LayerNormLayer Normalization层归一化 的最大区别。公式与手写验证把源码翻译成公式对向量x先算variance mean(x²)输出weight · x / sqrt(variance eps)。用一句 Python 手写一遍和模型里的层对比实测lnmodel.model.layers[0].input_layernorm# Qwen2RMSNorm(896)xtorch.randn(2,5,896)refln(x)defmy_rms_norm(x,weight,eps1e-6):varx.float().pow(2).mean(-1,keepdimTrue)returnweight*(x.float()*torch.rsqrt(vareps)).to(x.dtype)minemy_rms_norm(x,ln.weight,ln.variance_epsilon)print(最大误差:,(ref-mine).abs().max().item())实测最大误差 0.0——手写实现与 transformers 完全一致。eps1e-6是防止方差为 0 时除零的保险项来自 config 的rms_norm_eps。和 LayerNorm 差在哪实测对比RMSNorm只按平方均值缩放LayerNorm还会先减去均值再缩放。差异用数据说话实测lnormnn.LayerNorm(896)a,bln(x),lnorm(x)print(RMSNorm 输出均值:,a.mean().item())print(LayerNorm 输出均值:,b.mean().item())print(两者差异最大:,(a-b).abs().max().item())实测RMSNorm输出均值约0.000298没有归零LayerNorm约3.4e-9被减均值拉到 0 附近两者逐元素最大差异3.748。可见 RMSNorm 保留了向量的整体偏移只控制尺度。为什么这样反而好省掉减均值就省了跨特征的均值计算与一次减法在超大模型上更省算力而实践表明只做尺度归一化已经足够稳定训练这就是它被 Qwen、Llama 等采用的原因。维度RMSNormLayerNorm减均值否只按平方均值缩放是先中心化再缩放计算量更少多一步均值与减法可学习参数一组weightweightbias代表采用Qwen、Llama早期 Transformer、GPT-2两个实现细节为什么要转 float32、eps 防什么源码里有一行容易被略过hidden_states.to(torch.float32)——先升到 fp32 算方差与开方再转回原精度input_dtypehidden_states.dtype hidden_stateshidden_states.to(torch.float32)# 升精度再算variancehidden_states.pow(2).mean(-1,keepdimTrue)hidden_stateshidden_states*torch.rsqrt(varianceself.variance_epsilon)returnself.weight*hidden_states.to(input_dtype)# 转回原精度原因模型权重是 bf16config 里torch_dtypebfloat16而rsqrt对精度敏感用 bf16 直接算容易放大误差放进 fp32 里算、再转回去是数值稳定的常规做法。至于epsvariance_epsilon1e-6——它加在方差里防止某段向量恰好全 0 → 方差为 0 → 除零加一个极小数让rsqrt永远有定义。在层里的位置pre-norm 残差224 读过的Qwen2DecoderLayermodeling_qwen2.py:279-297把归一化放在子层之前pre-norm子层输出再与输入相加残差x ──input_layernorm── 注意力 ──┐ │ │ └────────── 残差相加 ─────────┘ ──post_attention_layernorm── MLP ──┐ │ │ └──────────── 残差相加 ─────────────────┘归一化让每层输入保持稳定尺度残差让梯度能直通深网络——两者配合24 层才堆得起来。为什么是 pre-norm先归一化再进子层如果反过来post-norm先子层再归一化深层网络里梯度要穿过一次次归一化训练早期容易不稳、需要小心调 warmuppre-norm 让残差那条直通路径始终干净梯度可以无衰减地回传这是深层 Transformer 能稳定训练的关键改动之一。Qwen2 的Qwen2DecoderLayer用的正是 pre-norm。顺带算一笔归一化层有多少参数每层有两个 RMSNorminput_layernorm与post_attention_layernorm每个只有一组 896 维weight24 层共 48 个参数合计48 × 896 43,008——相比 MLP 的千万级几乎可以忽略。但它不是常量weight是可学习参数训练时会调整这个缩放系数让网络自己决定每层的归一化强度。组件单层参数量占比两个 RMSNorm2 × 896 1,792极小注意力含 GQA1,836,160中等SwiGLU MLP13,074,432大头这也解释了 224 打印结构时看到的顺序input_layernorm、self_attn、post_attention_layernorm、mlp——归一化夹在两个重子层前后本身很轻却决定了整个网络能不能训得动。两个归一化的weight初始化为全 1训练前等价于不做额外缩放训练中再由梯度学习调整。下一篇从注意力的位置信息讲起RoPE。一句话记忆Qwen2RMSNorm只做三件事算mean(x²)、乘rsqrt(varianceeps)、再乘可学习权重不减均值手写实现与源码输出误差为 0与LayerNorm的差异输出均值 0.0003 vs 3.4e-9正体现只归尺度、不归中心它被放在每层子层之前pre-norm并与残差配合是 24 层能稳定堆叠的基础。OK本篇先到这里如有疑问欢迎评论区留言讨论祝各位功力大涨技术更上一层楼更多内容见下篇 blog【AI】【模型部署】基座模型研究RoPE 旋转位置编码
