从零手写现代 Transformer 架构：构建 0.04B 的思考中枢

2.2 积木一：RMSNorm（音量自动平衡器）

2.2.1 为什么深层神经网络必须归一化？

想象一个生活中的真实场景： 有 12 个扩音器串联在一起传递一句话。如果每个扩音器都稍稍放大了一点声音（增益 1.2 倍），传到第 12 个扩音器时，放大倍数就是 $1.2^{12} \approx 8.9$ 倍，音箱会瞬间爆音失真（在神经网络中这叫数值爆炸 / 梯度爆炸）！ 反过来，如果每个扩音器都衰减一点点（增益 0.8 倍），传到最后的声音就会变成 $0.8^{12} \approx 0.068$，微弱得几乎听不见（在神经网络中这叫梯度消失）。

归一化（Normalization）的本质，就是在每一层信号输入前后，安装一个灵敏的“自动增益控制器 (AGC)”： 不管前面的信号波动多剧烈，进来后都把它的总体能量拉回到统一的标准分贝！

2.2.2 RMSNorm vs 传统 LayerNorm：甩掉无用的均值包袱

传统的 LayerNorm 公式包含两步：

1. 先求平均值 $\mu = \frac{1}{d}\sum_{i=1}^d x_i$；
2. 再把每个元素中心化 $(x_i - \mu)$ 并除以标准差 $\sigma$： $\text{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta$

但深入研究发现：在深度神经网络中，真正保命的是“除以方差尺度进行能量缩放”，而“减去平均值”在数学上几乎毫无贡献！ 每次减均值不仅要让 GPU 在显存和计算核心之间来回搬运两次数据做跨核心求和（大幅拖慢速度），还平添了冗余计算。

RMSNorm（Root Mean Square Normalization）做出了天才般的精简：直接算均方根（RMS），然后直接除！ $\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$ $\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \odot \gamma$ 其中：

• $\epsilon$（Epsilon）是一个极小的保护常数（如 $10^{-6}$），作用是防止分母为 0 导致计算溢出；
• $\gamma$（Gamma）是一组和特征维度等长的可学习权重向量（初始全为 1.0），允许模型在后续训练中微调某些特征维度的“个性化音量”。

📐 数学深潜：RMSNorm 的尺度不变性证明

为什么直接除以均方根就能彻底稳定住网络？我们从数学上考察它的尺度不变性（Scale Invariance）：

假设前面的网络由于权重过大，将输入信号放大了 $\alpha$ 倍（变为 $\alpha x$）： $\text{RMS}(\alpha x) = \sqrt{\frac{1}{d} \sum_{i=1}^d (\alpha x_i)^2 + \epsilon} \approx \sqrt{\alpha^2 \cdot \frac{1}{d}\sum_{i=1}^d x_i^2} = |\alpha| \cdot \text{RMS}(x)$ 将其代入归一化公式： $\text{RMSNorm}(\alpha x) = \frac{\alpha x}{\text{RMS}(\alpha x)} \odot \gamma = \frac{\alpha x}{|\alpha|\text{RMS}(x)} \odot \gamma = \text{sign}(\alpha) \cdot \text{RMSNorm}(x)$ 数学结论： 无论前向特征的数值尺度被无序放大多大幅度，经过 RMSNorm 处理后，缩放系数 $\alpha$ 瞬间在分子分母中被精准相消抵消！ 这一性质从根本上锁死了激活值爆炸的可能性，赋予了模型极高的数值鲁棒性。

2.2.3 拿具体数字动手算一遍

为了彻底打消你的疑惑，我们拿 4 个具体的数字现场推演： 假设某一个词的特征向量为：$x = [2.0, -2.0, 2.0, -2.0]$

1. 第一步（求平方）： $[2.0^2, (-2.0)^2, 2.0^2, (-2.0)^2] = [4.0, 4.0, 4.0, 4.0]$
2. 第二步（求均值）： $\text{Mean} = (4.0 + 4.0 + 4.0 + 4.0) / 4 = 4.0$
3. 第三步（开根号求 RMS）： $\text{RMS} = \sqrt{4.0} = 2.0$（这就是它的均方根能量！）
4. 第四步（能量规整）： 每一个数字除以 2.0： $x_{\text{norm}} = [2.0/2.0, -2.0/2.0, 2.0/2.0, -2.0/2.0] = [1.0, -1.0, 1.0, -1.0]$

你看！原本可能很大的数值，瞬间被驯服到了标准的单位能量范围之内。

2.2.4 逐行代码实现与避坑提示

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps # 防止除以 0 的极小保护值
        # 可学习参数 gamma: 长度为 dim 的一维向量，初始全为 1.0
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # 1. x.pow(2).mean(-1, keepdim=True): 沿最后一个维度求平方的均值
        # 2. torch.rsqrt: 即 1.0 / sqrt(...)，求倒数平方根，GPU 上有专用硬件指令，一步完成！
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 💡 导师经验提示: 均方根计算很容易在半精度 (float16) 下发生微小数值下溢。
        # 因此业界标准实践是: 强转为 float32 计算均方根，算完后再转回张量原本的数据类型！
        output = self._norm(x.float()).type_as(x)
        # 乘以每个通道的可学习缩放系数 weight
        return output * self.weight

2.3 积木二：RoPE 旋转位置编码（时钟表针法）

这是现代大模型中最精妙、最具几何美感的设计，也是面试和深入理解模型时最高频考查的知识点。

2.3.1 为什么大模型需要“位置感”？

人类语言有一个基本铁律：文字出现的顺序至关重要！

• 句子 A：“我不喜欢你”
• 句子 B：“你喜欢我不” 两句话使用的字完全一样，但因为顺序不同，含义差了十万八千里。 然而，自注意力机制（Self-Attention）在计算各个词之间的相关性时，其底层的矩阵乘法是无序的（集合对称的）。如果不给词注入位置信息，在模型眼里这两句话没有任何区别！

2.3.2 传统绝对位置编码的死穴

在早期的模型中，做法非常直接生硬： “你是第 1 个词，就给你发一个刻着‘1’的绝对工牌；你是第 2 个词，就发一个刻着‘2’的工牌……” 这种做法被称为绝对位置编码（Absolute Position Embedding），但它有两大致命缺陷：

1. 无法外推（无法看长文）：如果你训练模型时最大长度只见过 1024 个词，模型就只拥有 1024 块工牌。遇到 2048 字的长文章时，后面的词根本没有工牌可发，直接瘫痪！
2. 违背人类理解直觉：人类读文章时，根本不关心一个动词究竟是整本书的“第 3421 个字”，人类只关心**“它和主语之间相对隔了 2 个字”！也就是相对位置（Relative Distance）才是语义关联的核心**。

2.3.3 时钟表针的顿悟：旋转角度衡量相对距离

怎么才能让模型只关注“相对距离”，同时又能自然延伸到无限长呢？ RoPE（Rotary Position Embedding，旋转位置编码）的作者灵机一动，想到了挂在墙上的时钟表针！

想象一下时钟表盘：

• 假设每个词是一根表针，它在句子里的位置 $m$ 决定了表针旋转的角度 $m \cdot \theta$；
• 比如第 1 个词旋转到 1 点钟方向，第 4 个词旋转到 4 点钟方向；
• 此时，第 1 个词和第 4 个词之间的夹角，就是： $\Delta = 4\theta - 1\theta = 3\theta$
• 如果整句话向后平移 10 个位置（从第 11 个词到第 14 个词）：
  • 第 11 个词转到 $11\theta$，第 14 个词转到 $14\theta$；
  • 它们之间的夹角依然是：$14\theta - 11\theta = 3\theta$！

===================================================================================
                       RoPE 时钟旋转相对位置直观图解
===================================================================================

       位置 m = 1 的词 (1点钟)                   位置 m = 4 的词 (4点钟)
              12                                        12
           11    1                                   11    1
        10    ▲     2                             10          2
       9      │      3                           9             3 ──►
        8     │     4                             8           4
           7     5                                   7     5
              6                                         6
              
       【两根表针之间的相对夹角】: Δ = 4θ - 1θ = 3θ (只取决于它们之间相隔了 3 个字！)
===================================================================================

📐 数学深潜：欧拉公式与相对位置内积的严格证明

为什么两根表针一转，内积打分就只与相对距离有关？我们用纯正的高等数学来证明这个优美的结论。

1. 欧拉公式与复数乘法

著名数学家欧拉提出了人类最伟大的公式之一： $e^{i\theta} = \cos\theta + i\sin\theta$ 在复平面上，任何一个二维向量 $(x_1, x_2)$ 都可以视作一个复数 $z = x_1 + i x_2$。 当这个复数乘以 $e^{i\theta}$ 时，几何意义就是逆时针旋转 $\theta$ 角： $z' = z \cdot e^{i\theta} = (x_1 + i x_2)(\cos\theta + i\sin\theta) = (x_1\cos\theta - x_2\sin\theta) + i(x_1\sin\theta + x_2\cos\theta)$ 写成实数坐标的形式，旋转后的新坐标 $(x_1', x_2')$ 为： $x_1' = x_1\cos\theta - x_2\sin\theta$ $x_2' = x_1\sin\theta + x_2\cos\theta$ 写成矩阵形式，这就是大家熟悉的二维旋转矩阵： $\begin{pmatrix} x_1' \\ x_2' \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix}$

2. 内积中的相对距离涌现

现在，我们在位置 $m$ 有一个查询向量 $q$，在位置 $n$ 有一个键向量 $k$。 我们将它们分别赋予位置旋转： $q_m = q \cdot e^{im\theta}, \quad k_n = k \cdot e^{in\theta}$ 在复空间中，两个向量的内积定义为：$\langle u, v \rangle_{\mathbb{C}} = \text{Re}(u \cdot v^*)$，其中 $v^*$ 是 $v$ 的共轭复数（$\overline{a + bi} = a - bi$）。 根据共轭性质，$(e^{in\theta})^* = e^{-in\theta}$。我们直接代入计算两者的内积： $\langle q_m, k_n \rangle_{\mathbb{C}} = \text{Re}\left( (q e^{im\theta}) \cdot (k e^{in\theta})^* \right) = \text{Re}\left( q \cdot k^* \cdot e^{im\theta} \cdot e^{-in\theta} \right)$ 利用同底数幂相乘指数相加法则： $\langle q_m, k_n \rangle_{\mathbb{C}} = \text{Re}\left( q k^* \cdot e^{i(m - n)\theta} \right)$

请停下来凝视这个惊艳的等式！ 原来位置下标 $m$ 和 $n$ 经过复数内积之后，赫然变成了相减的形式：$(m - n)$！ 结论：打分值的大小只取决于相对距离差值 $(m - n)$，与绝对位置 $m$ 或 $n$ 处于句首还是句末毫无关系！

3. 多频波长谱分析

在我们的模型中，单头维度 head_dim = 64，划分为 32 组二维旋转平面。 第 $i$ 组平面的旋转角频率定义为： $\theta_i = b^{-2(i-1)/d} = 10000^{-2(i-1)/64}, \quad i \in \{1, 2, \dots, 32\}$ 其对应的波长（旋转整整一周 $2\pi$ 所需跨越的 Token 数量）为： $\lambda_i = \frac{2\pi}{\theta_i} = 2\pi \cdot 10000^{2(i-1)/64}$

• 高频第一组（$i=1$）：$\theta_1 = 10000^0 = 1.0$，波长 $\lambda_1 = 2\pi \approx \mathbf{6.28\text{ 个 Token}}$。表针转得极快，专攻相邻几个词的语法修饰；
• 低频最后一组（$i=32$）：$\theta_{32} = 10000^{-62/64} \approx 1.25 \times 10^{-4}$，波长 $\lambda_{32} \approx \mathbf{50,000\text{ 个 Token}}$！表针走得极慢，在超长上下文中依然能敏锐定位开头与结尾的逻辑对应！

2.3.4 拆解关键算子：rotate_half

在 PyTorch 实际编程中，我们如何高效地计算旋转变换呢？ 回顾前面的代数形式： $(x_1', x_2') = (x_1, x_2)\cos\theta + (-x_2, x_1)\sin\theta$ 请看公式里的第二项：(-x2, x1)。 怎么把向量后半截取负号并挪到前半截？这就是大模型源码里最经典的算子 rotate_half：

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """
    将向量的前后两半拆开，后半截取负号并与前半截调换位置:
    输入: [x1, x2]
    输出: [-x2, x1]
    """
    half_dim = x.shape[-1] // 2
    x1 = x[..., :half_dim] # 前半截
    x2 = x[..., half_dim:] # 后半截
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """
    按照公式: x * cos(θ) + rotate_half(x) * sin(θ) 执行高速并行旋转
    """
    cos = cos.unsqueeze(0).unsqueeze(2) # 调整广播维度以匹配 [Batch, Seq_Len, Heads, Head_Dim]
    sin = sin.unsqueeze(0).unsqueeze(2)
    return (x * cos) + (rotate_half(x) * sin)

2.4 积木三：SwiGLU 门控前馈网络（私人小书房与智能阀门）

2.4.1 研讨会与私人书房的分工

如果一个大模型内部只有注意力层（Attention），它就会像一群只会交头接耳讨论、但自己肚子里毫无墨水的人。 在大模型里：

• 自注意力层（Attention）负责“信息交互”：每个词去打听“别人身上带了什么情报”；
• 前馈网络层（FeedForward，FFN）负责“知识存储与消化”：大家各自回到**“私人小书房”**，闭门调取自己脑海中沉淀的常识与逻辑记忆。 在大模型中，超过 60% 的模型参数都聚集在前馈网络（FFN）中，它是名副其实的模型“知识储备库”！

2.4.2 传统 FFN 与现代 SwiGLU 的质变

在传统的 Transformer 中，前馈网络结构非常单调： $\text{FFN}(x) = \text{ReLU}(x W_1) W_2$ 它只有一条单向直行管道，不管什么信息来了，全都走同一套非线性激活函数。

而现代开源顶尖模型（LLaMA-3、Mistral、DeepSeek）全部转向了 SwiGLU（Swish Gated Linear Unit）： 它像铺设了两条并行的自来水管道：

1. 内容管道（$W_{up}$）：承载输入的原始内容特征，将其升维展开；
2. 门控管道（$W_{gate} + \text{SiLU}$）：扮演**“智能电磁阀门”，自动计算出 $0.0 \sim 1.0$ 之间的开度比例！ 两条管道流出的数据逐元素相乘，相当于让模型学会了：“遇到噪音直接关闭阀门（乘 0），遇到关键线索全开通过（乘 1）”**！

📐 数学深潜：SiLU 激活函数的导数与自门控特性

为什么这里的门控必须选 SiLU（也称 Swish）而不是经典的 ReLU？我们对比其数学性质： $\text{SiLU}(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}}$ 对 $\text{SiLU}(x)$ 进行微积分求导： $\text{SiLU}'(x) = \sigma(x) + x \cdot \sigma'(x)$ 由于 Sigmoid 导数满足 $\sigma'(x) = \sigma(x)(1 - \sigma(x))$，代入整理得： $\text{SiLU}'(x) = \sigma(x) + x\sigma(x)(1 - \sigma(x)) = \sigma(x)\Big( 1 + x(1 - \sigma(x)) \Big)$

这一平滑可导性使得梯度在深度反向传播时极其顺畅，彻底根除了深层神经元坏死的问题。

2.4.3 逐行代码实现

import torch.nn.functional as F

class SwiGLUFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 升维与门控投影矩阵 (无 bias，追求纯粹的高速线性特征)
        self.w_gate = nn.Linear(config.d_model, config.d_ffn, bias=False) # 512 -> 1408
        self.w_up = nn.Linear(config.d_model, config.d_ffn, bias=False)   # 512 -> 1408
        # 降维输出矩阵: 将知识浓缩回原本的 512 维空间
        self.w_down = nn.Linear(config.d_ffn, config.d_model, bias=False) # 1408 -> 512

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. gate 管道: 计算 0~1 的软阀门开度
        gate_out = F.silu(self.w_gate(x))
        # 2. up 管道: 抽取宽阔的知识特征
        content = self.w_up(x)
        # 3. 逐元素相乘 (阀门放行) 并降维输出
        return self.w_down(gate_out * content)

2.5 积木四：GQA 分组查询注意力机制（侦探办案与共享秘书）

自注意力机制是整个模型能够理解语境、产生“智能感”的最核心引擎。

2.5.1 到底什么是 $Q, K, V$？一个侦探寻人比喻

很多教材一上来就列出大串矩阵乘法公式，初学者根本体会不到 $Q, K, V$ 的物理意义。 我们用一个极度生动的“侦探寻人”比喻来彻底弄懂它：

• $Q$（Query，侦探手中的寻人启事）：当前这个词“想打听什么、想找谁”；
• $K$（Key，房间里每个人脖子上挂的名片）：每个历史词“我具备什么特征、我是谁”；
• $V$（Value，每个人背包里真正装的情报内容）：如果我看中了你，你具体能提供给我什么知识。

自注意力的计算过程，就是一场有条不紊的侦探比对过程：

1. 打分（$Q \cdot K^T$）：侦探拿着寻人启事 $Q$，依次去和在场每个人的名片 $K$ 进行点乘比对，计算契合度分数；
2. 归一化（Softmax）：把分数换算成百分比概率（比如：80% 注意力看词 A，20% 注意力看词 B）；
3. 抽取情报（$\dots \times V$）：按照百分比权重，把词 A 和词 B 背包里的情报 $V$ 混合打包带走，形成最新、最丰富的情报综述！

📐 数学深潜：为什么 Attention 必须除以 $\sqrt{d_k}$？独立同分布方差推导

这是所有深度学习面试中最经典的灵魂拷问之一： $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$ 公式里那个不起眼的 $\sqrt{d_k}$ 到底在防范什么？如果不除它会发生什么灾难？

1. 内积方差的数学推演

假设查询向量 $q$ 和键向量 $k$ 均为 $d_k$ 维向量（在我们的模型中 $d_k = 64$）。 假设它们的各个分量 $q_i, k_i$ 是互相独立的随机变量，且均已通过前置归一化满足标准正态分布： $E[q_i] = 0, \quad \text{Var}(q_i) = 1$ $E[k_i] = 0, \quad \text{Var}(k_i) = 1$ 此时，它们的内积为： $S = q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 我们来求内积 $S$ 的期望与方差：

• 期望： $E[S] = \sum_{i=1}^{d_k} E[q_i k_i] = \sum_{i=1}^{d_k} E[q_i]E[k_i] = 0$
• 单项方差： $\text{Var}(q_i k_i) = E[(q_i k_i)^2] - (E[q_i k_i])^2 = E[q_i^2] E[k_i^2] - 0 = \text{Var}(q_i) \cdot \text{Var}(k_i) = 1 \times 1 = 1$
• 总和方差（独立变量和的方差等于方差的和）： $\text{Var}(S) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = \sum_{i=1}^{d_k} 1 = \mathbf{d_k}$
• 标准差： $\sigma_S = \sqrt{\text{Var}(S)} = \mathbf{\sqrt{d_k}}$

2. 不除以 $\sqrt{d_k}$ 的致命灾难：Softmax 梯度饱和

当单头维度 $d_k = 64$ 时，内积结果的标准差高达 $\sigma_S = \sqrt{64} = 8$！ 这意味着计算出来的注意力分数可能经常出现 $+16$ 甚至 $+24$ 这样极端的悬殊大数。 当这些极端数字送入 Softmax 函数时： $p_i = \frac{e^{S_i}}{\sum_j e^{S_j}}$ 最大值对应的分子 $e^{24}$ 会以指数级绝对优势碾压其他所有分母项，导致最大项的概率 $p_{\max} \approx 1.0$，而其他所有项的概率几乎全等于 $0.0$！

让我们求 Softmax 的反向传播梯度： $\frac{\partial p_i}{\partial S_i} = p_i (1 - p_i)$

• 当 $p_i \to 1.0$ 时：$p_i(1 - p_i) \to 1.0 \times 0 = \mathbf{0}$！
• 当 $p_i \to 0.0$ 时：$p_i(1 - p_i) \to 0.0 \times 1 = \mathbf{0}$！ 灾难发生了：整个注意力机制的梯度彻底归零，反向传播网络瞬间陷入脑死亡平台期！

3. 救命解药：缩放因子归一化方差

因此，我们必须在点积后强制除以 $\sqrt{d_k}$： $\text{Var}\left( \frac{q \cdot k}{\sqrt{d_k}} \right) = \left(\frac{1}{\sqrt{d_k}}\right)^2 \cdot \text{Var}(q \cdot k) = \frac{1}{d_k} \cdot d_k = \mathbf{1.0}$ 数学结论： 除以 $\sqrt{d_k}$ 之后，打分方差被精准重置回了 $1.0$！这使得输入始终落在 Softmax 导数最敏感、非饱和的黄金响应区间内！

2.5.2 因果掩码（Causal Mask）的数学极限

大模型是自回归模型，在计算第 $t$ 个词时严禁偷看未来第 $t+1$ 个词。 在注意力得分矩阵中，我们将所有未来位置加上 $-\infty$： $M_{ij} = \begin{cases} 0, & i \ge j \text{（过去与当前，保留）} \\ -\infty, & i < j \text{（未来，遮蔽）} \end{cases}$ 根据高等数学中的指数函数极限： $\lim_{x \to -\infty} e^x = 0$ 在经过 Softmax 之后，未来词的注意力权重自然精准衰减为 $0.0\%$，彻底杜绝了信息穿越偷看！

2.5.3 为什么需要 GQA？侦探团与共享秘书

在第 05 章我们会详细剖析自回归生成：每生成一个词，历史所有词的 $K$ 和 $V$ 都必须常驻在显存里（KV-Cache）。

• 全量多头 (MHA)：如果我们有 12 个 Query 侦探头，就需要配 12 对 KV 秘书。当序列很长时，显存被大批量的 $K$ 和 $V$ 瞬间撑爆，显存带宽严重告急！
• 单查询 (MQA)：为了省钱，让全部 12 个侦探头共用 1 个秘书。显存是省了，但所有侦探被迫看完全一样的信息，模型智商大打折扣。
• 分组查询 (GQA，本教程黄金选型)：让 12 个 Query 头配 4 对 KV 秘书（3 分组）！每 3 个侦探共享一个专属秘书：
  • 显存直接立省 66.7%！
  • 模型的多样性表达能力几乎毫发无损！

2.5.4 逐行实现 Attention 模块

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """将 [Batch, Seq_Len, 4, 64] 沿着头维度重复复制 3 次，扩展为 [Batch, Seq_Len, 12, 64]"""
    if n_rep == 1:
        return x
    batch, slen, n_kv, head_dim = x.shape
    return (
        x[:, :, :, None, :]
        .expand(batch, slen, n_kv, n_rep, head_dim)
        .reshape(batch, slen, n_kv * n_rep, head_dim)
    )

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_q_heads = config.n_q_heads   # 8 个 Query 侦探头
        self.n_kv_heads = config.n_kv_heads # 2 个 KV 秘书头
        self.n_rep = config.n_q_heads // config.n_kv_heads # 组大小 = 4
        self.head_dim = config.head_dim     # 每个头的维度 = 64 (8 * 64 = 512)

        # 线性映射矩阵 (全无 bias，保持干净纯粹)
        self.wq = nn.Linear(config.d_model, config.n_q_heads * config.head_dim, bias=False)
        self.wk = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.wv = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.wo = nn.Linear(config.n_q_heads * config.head_dim, config.d_model, bias=False)

    def forward(self, x, cos, sin, mask=None, kv_cache=None):
        b_sz, seq_len, _ = x.shape
        # 1. 将 512 维特征分别投影出 Q, K, V，并拆分为多头
        xq = self.wq(x).view(b_sz, seq_len, self.n_q_heads, self.head_dim)
        xk = self.wk(x).view(b_sz, seq_len, self.n_kv_heads, self.head_dim)
        xv = self.wv(x).view(b_sz, seq_len, self.n_kv_heads, self.head_dim)

        # 2. 为 Q 和 K 注入时钟旋转位置编码 (RoPE)
        xq = apply_rotary_emb(xq, cos, sin)
        xk = apply_rotary_emb(xk, cos, sin)

        # 3. 维护增量推理 KV-Cache (第 05 章详细讲解)
        new_kv = None
        if kv_cache is not None:
            prev_k, prev_v = kv_cache
            xk = torch.cat([prev_k, xk], dim=1) if prev_k is not None else xk
            xv = torch.cat([prev_v, xv], dim=1) if prev_v is not None else xv
            new_kv = (xk, xv)

        # 4. GQA 广播: 把 2 个 KV 头复制 4 遍，对齐 8 个 Query 头
        xk_expanded = repeat_kv(xk, self.n_rep)
        xv_expanded = repeat_kv(xv, self.n_rep)

        # 5. 调整维度以符合 PyTorch 官方底层硬件加速格式: [Batch, Heads, Seq_Len, Head_Dim]
        xq = xq.transpose(1, 2)
        xk_expanded = xk_expanded.transpose(1, 2)
        xv_expanded = xv_expanded.transpose(1, 2)

        # 6. 调用超强底层硬件融合内核 (FlashAttention / SDPA)
        is_causal = (mask is None and seq_len > 1 and kv_cache is None)
        output = F.scaled_dot_product_attention(
            xq, xk_expanded, xv_expanded,
            attn_mask=mask,
            dropout_p=0.0,
            is_causal=is_causal # 自动构建严密的下三角因果防偷看掩码！
        )

        # 7. 多头汇报汇总: 还原回 [Batch, Seq_Len, 512]
        output = output.transpose(1, 2).contiguous().view(b_sz, seq_len, -1)
        return self.wo(output), new_kv

2.6 两大点睛之笔：残差直连与权重绑定

在把积木拼成完整大厦前，我们还必须了解现代深度学习里最神奇的两个天才设计。

2.6.1 残差直连（Residual Connection）：免检直达电梯

在每一层研讨中，你都会看到类似这行代码： h = x + attn_out

• 为什么不直接使用 h = attn_out，非要加上原本的 x 呢？
• 传话失真比喻：想象你找了 12 个人传话，第 1 个人告诉第 2 个人，第 2 个人再修饰一下告诉第 3 个人……传到第 12 个人手里时，原话往往已经被改得面目全非！在反向传播计算梯度时，12 层的连乘也会导致底层梯度彻底衰减为 0。
• 免检直达电梯：残差连接 x + ... 就像在 12 层大厦旁边架设了一部无阻碍直达电梯。原始的输入信号 $x$ 可以原封不动坐电梯直通顶层，而注意力层和前馈层只需要负责学习“在原基础上的轻微增量修改（Residual 残差）”！这让 100 层的深层网络也能健步如飞。

2.6.2 权重绑定（Weight Tying）：输入字典与输出字典的镜像契约

在我们的 0.04B 模型中：

• 输入层有一个把 4096 个词映射到 512 维空间的 token_embeddings 矩阵（参数量：$4096 \times 512 = 2,097,152$）；
• 输出层同样有一个把 512 维向量重新打分回 4096 个词的 lm_head 矩阵（参数量同样是 $2,097,152$）。

这两组参数加起来高达 419 万，占了整个 0.04B 模型的整整 11.0%！ 既然它们一个负责把词变成向量，另一个负责把向量翻译回词，它们本质上就是同一本中英对照字典的正反两面查阅！ 通过在代码中加上：

self.lm_head.weight = self.token_embeddings.weight

我们让它们共享同一个物理显存矩阵！模型瞬间省下了整整 210 万个冗余参数，不仅显存更小，而且让输入与输出的词语义表征紧密对齐，收敛更快！

2.7 终极大总装：构建完整的 MiniLLaMAForCausalLM

现在，所有的地基、梁柱、门窗都已经准备就绪，我们把它们正式组装成一个拥有 8180 万参数的完整大语言模型！

2.7.1 单层圆桌会议：TransformerBlock

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 注意力层及其门前音量平衡器
        self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.attn = Attention(config)
        # 前馈网络层及其门前音量平衡器
        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.ffn = SwiGLUFeedForward(config)

    def forward(self, x, cos, sin, mask=None, kv_cache=None):
        # 1. 第一阶段: 前置归一化 + 注意力情报交互 + 残差直连
        attn_out, new_kv = self.attn(self.attn_norm(x), cos, sin, mask=mask, kv_cache=kv_cache)
        h = x + attn_out

        # 2. 第二阶段: 前置归一化 + 前馈书房独立思考 + 残差直连
        out = h + self.ffn(self.ffn_norm(h))
        return out, new_kv

2.7.2 整体因果语言模型：MiniLLaMAForCausalLM

class MiniLLaMAForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        # 1. 词嵌入映射
        self.token_embeddings = nn.Embedding(config.vocab_size, config.d_model)

        # 2. 预先计算 RoPE 旋转角度查找表，注册为非求导缓存 (Buffer)
        cos, sin = precompute_freqs_cis(config.head_dim, config.max_seq_len, config.rope_theta)
        self.register_buffer("cos_cached", cos, persistent=False)
        self.register_buffer("sin_cached", sin, persistent=False)

        # 3. 堆叠 12 层 TransformerBlock 圆桌研讨层
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])

        # 4. 终层归一化平滑
        self.norm = RMSNorm(config.d_model, eps=config.norm_eps)

        # 5. 语言模型预测输出头
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # 6. 开启权重绑定: 输入输出共享词典矩阵
        if config.tie_word_embeddings:
            self.lm_head.weight = self.token_embeddings.weight

    def forward(self, input_ids, labels=None, past_key_values=None, start_pos=0):
        b_sz, seq_len = input_ids.shape

        # 第一步: 查表获取初始 512 维特工档案袋
        h = self.token_embeddings(input_ids)

        # 第二步: 切出当前窗口对应的 RoPE 旋转角度
        cos = self.cos_cached[start_pos : start_pos + seq_len]
        sin = self.sin_cached[start_pos : start_pos + seq_len]

        new_past_key_values = [] if past_key_values is not None else None

        # 第三步: 顺次走过 12 层圆桌会议的深度研讨
        for idx, layer in enumerate(self.layers):
            layer_kv = past_key_values[idx] if past_key_values is not None else None
            h, updated_kv = layer(h, cos, sin, kv_cache=layer_kv)
            if new_past_key_values is not None:
                new_past_key_values.append(updated_kv)

        # 第四步: 终层音量稳定
        h = self.norm(h)
        # 第五步: 投影回 4096 个备选词的未归一化打分
        logits = self.lm_head(h)

        # 第六步: 如果给定了真实标签 labels，计算自回归交叉熵损失
        loss = None
        if labels is not None:
            # 错位 1 个 Token 计算自回归目标: 用当前词的打分去预测下一个真实的词！
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1)
            )

        return logits, loss, new_past_key_values

📐 数学深潜：从极大似然估计到自回归交叉熵损失

为什么代码第 385~390 行要将 logits 和 labels 错开 1 位计算交叉熵？我们从概率论与信息论根源揭秘：

1. 语言模型的联合概率分布

给定一段长度为 $T$ 的文本序列 $W = (w_1, w_2, \dots, w_T)$，大语言模型的核心目标是学习整个句子的联合概率分布 $P(w_1, w_2, \dots, w_T)$。 根据概率论的乘法法则（链式法则），联合概率可以严格因式分解为单步条件概率的连乘积： $P(w_1, w_2, \dots, w_T) = P(w_1) \cdot P(w_2 \mid w_1) \cdot P(w_3 \mid w_1, w_2) \cdots P(w_T \mid w_1, \dots, w_{T-1}) = \prod_{t=1}^T P(w_t \mid w_{<t})$

2. 极大似然估计（MLE）与负对数似然（NLL）

在统计学中，我们希望模型的参数 $\theta$ 使得训练集中的真实文本出现概率最大化（极大似然估计）： $\theta^* = \arg\max_\theta \prod_{t=1}^T P_\theta(w_t \mid w_{<t})$ 为了防止数以万计的微小概率连乘引发计算机浮点数下溢（Underflow），同时把乘法转换为更易求导的加法，我们对两侧取自然对数 $\ln$，并取负号将“最大化问题”转化为深度学习标准的“损失最小化问题”： $\mathcal{L}_{\text{NLL}} = -\ln \prod_{t=1}^T P_\theta(w_t \mid w_{<t}) = -\sum_{t=1}^T \ln P_\theta(w_t \mid w_{<t})$

3. 为什么它等于交叉熵？

在每一步 $t$，真实的下一个词 $w_t$ 可以表示为一个 one-hot 概率向量 $y_t \in \{0, 1\}^V$（除了真实词位置为 1，其余全为 0）。 交叉熵的通用数学定义为： $H(y_t, \hat{y}_t) = -\sum_{v=1}^V y_{t, v} \ln \hat{y}_{t, v}$ 因为 $y_t$ 只有在真实词 $v = w_t$ 处才为 1，上式直接坍缩为： $H(y_t, \hat{y}_t) = -\ln \hat{y}_{t, w_t} = -\ln P_\theta(w_t \mid w_{<t})$ 数学结论： 对全序列求平均损失，正好就是我们在代码中调用的： F.cross_entropy(shift_logits.view(-1, 4096), shift_labels.view(-1))！ 这就是为什么我们在第 385 行把输入序列截断到倒数第 2 个字（logits[..., :-1, :]），去监督真实序列从第 2 个字到末尾（labels[..., 1:]）的全部下一步！

2.8 导师小结与课后动手试炼

通过本章的循序渐进搭建，原本神秘的大模型中枢已经在你的眼皮底下彻底透明：

1. Pre-RMSNorm：像灵敏的调音台，去除了冗余的均值运算，具有严格的尺度不变性，保障深层信号能量稳定；
2. RoPE：利用欧拉公式复数乘法，以纯几何旋转在内积中天然导出相对距离差 $(m-n)$；
3. SwiGLU：以处处平滑可导的 SiLU 门控点乘赋予模型更高的记忆密度与知识过滤能力；
4. GQA：通过 $\sqrt{d_k}$ 缩放因子将内积方差重置为 1.0 杜绝 Softmax 梯度饱和，以 4 分组共享节约 75% 的 KV-Cache 显存；
5. 残差直连与权重绑定：像直达免检电梯和正反查阅的字典，让模型深而不散、轻盈紧凑。

💡 课后动手小试炼：亲手唤醒你的 0.04B 模型

步骤 1：确认环境依赖（安装 PyTorch）

本章的模型构建依赖纯原生的 PyTorch 基础张量库。如果在终端中运行提示 No module named 'torch'，请在终端（PowerShell / CMD）中先执行安装：

pip install torch

步骤 2：在项目根目录下运行（避坑指南 ⚠️）

在运行任何 Python 代码或测试前，请确保你的当前终端工作目录位于项目根目录（llm-start），而不是在子目录 src 中！

• 错误做法 ❌：在 d:\code\project\llm-start\src 下运行 python，再执行 from src.model import ...。因为当前目录已经是 src，Python 会在 src 内部寻找不存在的子文件夹 src，直接导致报错 ModuleNotFoundError: No module named 'src'！
• 正确做法 ✅：终端始终保持在项目总根目录 d:\code\project\llm-start 下执行命令。

方式 A：推荐一键执行测试脚本（最省心、零失误）

项目已经贴心地为你准备好了单体出厂体检脚本 tests/test_model.py。在项目根目录下直接运行：

python tests/test_model.py

终端预期打印输出：

============================================================
       Mini-LLaMA-0.04B 架构出厂体检与前向验证
============================================================
1. 正在初始化模型 (隐层维度: 512, 层数: 12, 词表: 4096)...
2. 模型总参数量核算: 35,926,528 个参数 (35.93 M)
3. 模拟输入张量形状: torch.Size([1, 5]) (Batch=1, Seq_Len=5)
4. 模型输出 Logits 形状: torch.Size([1, 5, 4096]) (符合预期: [1, 5, 4096])

🎉 恭喜！Mini-LLaMA-0.04B 神经网络架构前向推理与形状校验全部通过！✅
============================================================

方式 B：在 Python 交互式命令行中逐步体验

如果你想进入 Python 交互环境（REPL）体验“逐行调用”，请在项目根目录下运行 python，然后依次输入：

# 1. 导入配置与模型定义
from src.model import MiniLLaMAForCausalLM, MiniLLaMAConfig
import torch

# 2. 实例化配置与模型
config = MiniLLaMAConfig()
model = MiniLLaMAForCausalLM(config)

# 3. 打印真实参数量
total_params = sum(p.numel() for p in model.parameters())
print(f"总参数量: {total_params:,} ({total_params/1e6:.2f} M)")
# 终端打印: 总参数量: 35,926,528 (35.93 M)

# 4. 模拟送入一句话 (Batch=1, 长度为 5 的假 Token 序列)
dummy_input = torch.tensor([[1, 232, 58, 209, 2]])
logits, loss, _ = model(dummy_input)
print(f"输出 Logits 形状: {logits.shape}") 
# 终端打印: 输出 Logits 形状: torch.Size([1, 5, 4096])

试着观察终端打印：

• 参数量是不是恰好为你算出的 35.93 M？
• 输出的形状是不是整整齐齐的 [1, 5, 4096]？

大脑已经完全铸就！在下一章中，我们将编写一条极速运转的数据供油管道，源源不断地给这个大脑灌注海量知识——第 03 章｜语料预处理与高性能二进制数据流水线！