LLM智能应用开发

第5讲: 大语言模型解析 II
基于HF LlaMA实现的讲解

RMSNorm and FFN

LLM结构的学习路径

  • LLM结构解析(开源LlaMA)
  • 自定义数据集构造
  • 自定义损失函数和模型训练/微调

Decoder-only：归一化与 FFN 的位置

  • 图中为 Llama 的 Decoder-only 结构
  • 输入部分
    • Input embedding
    • Positional embedding
  • Transformer部分
    • Attention module
    • Feed forward network

HF LlaMA模型结构

LlamaForCausalLM(
  (model): LlamaModel(
    ...
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention
        (mlp): LlamaMLP
        (input_layernorm): LlamaRMSNorm
        (post_attention_layernorm): LlamaRMSNorm
    )
    ...
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

LlamaDecoderLayer内部结构

(self_attn): LlamaAttention(
  (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
  (k_proj): Linear(in_features=2048, out_features=512, bias=False)
  (v_proj): Linear(in_features=2048, out_features=512, bias=False)
  (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
  (rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
  (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
  (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
  (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
  (act_fn): SiLU()
)
(input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)

LlamaDecoderLayer内部结构

  • Normalization模块
    • LlamaRMSNorm
  • 两个主要模块
    • LlamaAttention
    • FFN（代码类名 LlamaMLP）

FFN（前馈网络）是 Transformer 中的子层，通常由 MLP（多层感知机）实现。
本讲统一称 FFN，代码中保留 LlamaMLP 和 mlp 的真实名称。

Normalization

(input_layernorm): LlamaRMSNorm
  • Normalization，中文又称“归一化”，“标准化”,...
  • 作用：调整数据的分布(scaling)，亦可称“缩放”
  • 数据包括：输入、输出、中间层表示，...

最经典的Normalization

Standard score: , Minmax feature scaling:

Normalization示例

一维的输入，归一化后的输出
center

归一化如何帮助优化

对输入特征做尺度调整，有助于改善优化条件。

center

图中展示参数空间的优化轨迹；输入缩放的效果仅作示意。

“更容易优化”，就一定“更不容易过拟合”吗？

Normalization 与 Regularization

前一页讨论优化是否容易；还要区分泛化是否可靠。

  • 归一化：调整输入或激活的尺度，帮助稳定计算与优化。
  • 正则化：通过约束、惩罚或随机化等方式，控制过拟合。

例如，损失惩罚是正则化的一种形式：

归一化也可能影响泛化，但不能据此将它等同于正则化。

大语言模型引入Normalization

Normalization：“调整数据分布”

  • 数据长什么样子？都是tensor！维度很多！
    • 原始输入: vocab embedding
      • tensor shape: <batch_size, sequence_length, hidden_dim>
    • 深度学习模型中间层表示(hidden states/representations)
      • tensor shape: <batch_size, sequence_length, hidden_dim>

Normalization案例

center

Normalization 的统计范围

输入 的形状为 ：batch、token 位置、hidden 特征。

  • 固定 ，对向量 的 个特征计算统计量。
  • 不混合其他样本，也不跨整段序列统计。
  • LayerNorm 与 RMSNorm 的计算范围相同，使用的统计量不同。
输入 X       [B, S, H]
统计量       [B, S, 1]   ← 沿 hidden 维，keepdim=True
广播并缩放   [B, S, H]   ← 每个 token 使用自己的统计量

LayerNorm

  • 针对每个token在hidden_dim维度上做标准化，避免依赖batch/序列长度
  • 计算流程：
  • 可学习参数恢复尺度与偏移，保持表达能力
  • 优势：序列自回归场景中激活分布更稳定，梯度更易传播

RMSNorm

固定一个 token，记它的 hidden 向量为 ：

  • 为 hidden 维度，求和只涉及当前 token 的特征。
  • 是可学习的逐维缩放参数。
  • 同一层的所有 batch、token 共享同一个 。
  • 用于数值稳定；RMSNorm 不减去均值。

LLM 中的 Normalization

center

手搓 RMSNorm

输入与输出：。计算分成三步：

  1. x.square().mean(-1, keepdim=True)：平方均值，形状 。
  2. rsqrt(mean_square + eps)：缩放系数，广播乘到 个特征。
  3. 逐维乘以 ，形状 ，输出仍为 。

这里 始终表示 hidden 维度；batch 与 sequence 维保持不变。

编码实现 RMSNorm

x = torch.tensor([[[1., 2., 1., 2.], [2., 4., 2., 4.]]])
norm = nn.RMSNorm(4, eps=1e-6)
x32 = x.float()
mean_square = x32.square().mean(-1, keepdim=True)
y = (x32 * torch.rsqrt(mean_square + norm.eps)).to(x.dtype)
y = y * norm.weight

torch.testing.assert_close(y, norm(x))

keepdim=True 保留 [b, s, 1]，广播到各自的 4 个特征。
验证时统一 和 ； 为可学习参数。

为什么选择 RMSNorm？

两者都用于控制每个 token 的激活尺度，但处理方式不同：

方法 中心化 尺度调整
LayerNorm 减去均值 除以标准差，再学习缩放与偏移
RMSNorm 不减均值 除以均方根，再学习缩放
  • 计算更简单：省去减均值及相关计算。
  • 效果与成本的取舍：原论文在所测任务中取得相近效果，Llama 采用这一设计。
  • 这不意味着 LayerNorm 失效；均值并非永远不重要，速度收益依赖实现和硬件。

Attention 后，为什么还有 FFN？

Attention 让当前位置从其他 token 获取上下文。
FFN 接着对这个位置的表示做非线性特征加工。

  • Attention：跨 token 汇聚与交互。
  • FFN：在每个 token 内组合、调节特征。
  • 同一层的 FFN 参数在所有位置共享。

例如“我不喜欢这部电影”：汇聚“不”的信息后，
当前位置还需要把这些特征组合成有用的更新。

升维：展开特征组合

单个 token 的表示：，映射到 个中间特征，通常 。

  • 每个中间单元都有一组权重，学习一种特征组合。
  • 示意：两维输入 [2,1] → 四维响应 [3,1,2,1]（和、差、原值）。
  • 升维提供更宽的计算空间，并不增加输入的原始信息。

若只做线性升维再降维：，
仍等价于一个线性层。

展开特征之后，怎样根据当前输入进一步加工它们？

门控：按当前输入调节中间特征

Llama 的 SwiGLU 有两条升维分支：

up_proj 产生候选特征，gate_proj 根据当前 产生调节系数。

候选特征 u     [ 2,  -1,  3 ]
门控系数 g     [ 0, 0.5,  2 ]   ← 示意值
逐元素相乘 z   [ 0, -0.5,  6 ]

系数接近 0 时抑制该特征，大于 1 时放大。
SiLU 的输出也可为负；它不是 0/1 开关或概率。

升维展开特征，激活与门控对这些特征进行非线性加工。

降维与残差

将 个中间特征加权组合成 维更新：

降维是可学习的组合；输出回到 hidden 维，才能与主干相加。

是保留的主干表示， 是交给 FFN 的归一化输入。
图中上方 Linear 输出 ，旁路保留 ，顶端相加得到 。

FFN实现，LlamaMLP

(mlp): LlamaMLP(
  (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
  (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
  (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
  (act_fn): SiLU()
)
  • 组件：三个nn.Linear层，一个SiLU激活函数
    • SiLU: torch.nn.functional.silu(x)
    • Linear: torch.nn.Linear(in_features, out_features)

FFN流程

  • 输入tensor: <batch_size, sequence_length, hidden_dim>
  • 第一步：
    • 通过gate_proj获得gate tensor，经过SiLU激活得到gate tensor
    • 通过up_proj获得up tensor
  • 第二步：元素乘(elementwise multiply): gate tensor 和 up tensor
  • 第三步: 通过down_proj获得down tensor

激活函数(引入非线性)

  • 通过引入非线性函数(nonlinear function)可增加模型非线性的能力，在模型中，也称之为激活函数
    • 线性层:
    • 激活:
  • activation类型
    • ReLU、Sigmoid、SiLU 等；SwiGLU 是含 SiLU 的门控模块。

经典激活函数

  • Sigmoid：，输出压缩到，常用于概率建模
  • Tanh：，零均值输出缓解偏置，会出现梯度消失
  • ReLU：，计算高效
  • Leaky ReLU： with ，为负半轴保留小梯度

激活函数(引入非线性)

center

激活函数(引入非线性)

center

SwiGLU FFN：门控与参数量

  • 门控分支：。
  • 候选分支：。
  • 逐元素相乘，再投影回 hidden 维：。
  • SiLU：。
  • 不同于概率门控，SiLU 的输出不限于 。

隐藏维度为 、中间维度为 时（忽略 bias）：
三次线性映射约有 个参数；普通两层 FFN 约为 。

比较参数开销时，需要同时说明中间维度。

SwiGLU FFN 的前向实现

center

上图是 SiLU 激活曲线；下面是完整 SwiGLU FFN 的前向计算。

down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

补充：RoPE 实现

RoPE的2D实现（相邻维度配对）
center

RoPE的n维实现
center

RoPE公式

  • 基础频率：，常用
  • 位置角度：，其中为token位置，为维度对索引
  • 旋转变换：
  • 可写成旋转矩阵：，保持范数不变

RoPE实现

  • 目标：实现RoPE(对应的)

    • 准备一堆，以及对应的cos和sin数值

  • 构建RoPE矩阵: 流程

    • 批量算cos，再批量算sin
    • 涉及torch算子
      • torch.arange, torch.sin, torch.cos, torch.outer

RoPE高效实现

center

RoPE高效实现

Llama 使用前后半维配对：；cos/sin 也按同样布局复制。

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    # q/k: [b, s, h, d]; cos/sin: [s, d]
    cos = cos[None, :, None, :]
    sin = sin[None, :, None, :]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

https://marp.app/

![bg right:40% 100%](images/l4/transformer.png)

此页讲输入特征缩放对参数优化的直觉，不是直接对参数做 norm，也不是 RMSNorm 让损失曲面变圆的证明。原图“较大参数梯度主导”应口头解释为不同方向优化条件不均衡，不能以参数大小推断梯度大小。结尾问题引出下一页的优化与泛化目标区别。

两者的效果可能交叉，此处区分主要作用机制，不应把“优化”与“泛化”讲成互斥。正则化不只限于修改损失函数，例如 dropout 采用随机化。下一页开始讨论模型内部 hidden activations 的归一化。

对比：沿 batch 维归约会将不同样本混在一起；沿完整 sequence 维归约会让当前位置依赖其他位置，可能包含未来 token。当前 LLM LayerNorm(H)/RMSNorm(H) 均只沿 H 统计。参考 PyTorch LayerNorm 与 RMSNorm 官方文档。

PyTorch 实现：torch.nn.RMSNorm(H)，https://pytorch.org/docs/stable/generated/torch.nn.RMSNorm.html

延续原 CV 图的着色方式：横轴 B（文本样本）、斜轴 S（token 位置）、竖轴 H（hidden 特征）。着色竖条是 X[b,s,:]，一次统计只涉及这条向量。LN 与 RMSNorm 范围相同，统计量不同。公式中的分母实际包含 epsilon，gamma/beta 分别表示逐维缩放/偏移。

依据：Root Mean Square Layer Normalization，https://arxiv.org/abs/1910.07467。论文结果不应泛化为所有模型、任务或现代内核的固定加速比例。Llama pre-norm 分支采用 RMSNorm。

语义例子仅用于直觉：不表示某个神经元必然对应“否定”或“喜欢”。Attention 自身也包含特征投影与 softmax 非线性；这里强调分工，不能讲成 attention 只搬运数据。参考 Attention Is All You Need §3.3：https://arxiv.org/html/1706.03762v7#S3.SS3

升维不增加输入的信息量；更长的向量也不保证效果更好。示例中的 a+b、a-b 是手工组合，真实权重通过训练学习。这里的线性合并结论只针对连续线性投影，不包括两个输入相关分支的逐元素乘法。

这是输入相关的逐维调节；不同 token 使用同一组投影参数，但计算出的门值不同。参考 GLU Variants Improve Transformer：https://arxiv.org/html/2002.05202v1

Llama pre-norm decoder block：t 是第一次残差后的表示，h 是 Norm2(t)，z 是门控后的中间特征；Delta_FFN 的命名强调更新量，不意味着 h 与 y 直接相减。全程 batch、sequence 不变。公式使用列向量；PyTorch Linear 对应 z @ W_d.T。

* SiLU：$x\cdot\sigma(x)$，兼具平滑和门控性质，常作为SwiGLU中的激活