从零手搓大语言模型：模型结构深度解析与实现

首先，我们来深入解析现代大语言模型（LLM）的结构核心：当前主流模型几乎一致采用了Decoder-Only架构。与传统的Encoder-Decoder不同，Decoder-Only只专注于一件事——根据给定的上文，预测下一个词。无论是GPT、Llama、Qwen还是DeepSeek，这些明星模型的核心结构都可以归结为这一设计。

因此，我们从MiniMind出发，把这个结构掰开揉碎地讲透彻。先看看整体骨架：

MiniMindForCausalLM（因果语言模型外壳）
├── embed_tokens：Token Embedding，将token_id映射为768维向量
├── 8 × MiniMindBlock（Transformer层，逐层堆叠）
│   ├── RMSNorm → Attention（多头注意力） → 残差连接
│   └── RMSNorm → FeedForward（SwiGLU） → 残差连接
├── RMSNorm（最终归一化）
└── lm_head：线性投影层，将768维向量映射为6400维词表概率分布

这个结构看起来简洁明了，但每个组件都蕴含着深刻的设计原理。下面我们逐一剖析。

一、RMSNorm：归一化层

1.1 为什么需要归一化

在神经网络逐层计算的过程中，矩阵乘法会持续放大数值的尺度，如果不加以约束，经过多层叠加后，数值要么溢出，要么梯度消失——导致训练无法稳定收敛。归一化的作用就是在每个计算模块的入口将数值拉回标准范围。

1.2 RMSNorm 的实现

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        self.weight = nn.Parameter(torch.ones(dim)) # 可学习的缩放因子

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return self.weight * self.norm(x.float())

RMSNorm和LayerNorm最大的区别在于：它只对均方根进行缩放，不减去均值。少了一步计算，速度更快，而在实际应用中效果相当。这是典型的在速度与性能之间做出的平衡选择。

1.3 归一化不会丢失语义

归一化改变的只是向量的“长度”（尺度），并不改变向量的“方向”（维度间的比例关系）。语义信息主要编码在方向中，因此不会因归一化而丢失。而且那个可学习的weight参数，正是模型用来恢复必要尺度信息的手段。

1.4 三处归一化的职责

• input_layernorm：Attention计算前的尺度校准
• post_attention_layernorm：FeedForward计算前的尺度校准
• final norm（MiniMindModel.norm）：送入lm_head前的最终校准

三者功能完全相同，只是放置位置不同。目的是确保每个计算模块都能接收到尺度正常的输入。

二、RoPE：旋转位置编码

2.1 问题的起点

Transformer的自注意力机制本身并不具备位置感知能力。“我打你”和“你打我”，如果不加入位置信息，对模型来说其实是等价的。因此，我们需要一种方式让模型知道每个词在序列中的具体位置。

2.2 核心思路

RoPE的做法非常巧妙：它把每个位置的Q和K向量按特定角度进行旋转。位置不同，旋转的角度也不同。当两个词做点积时，结果只和它们的相对距离有关——这个特性是RoPE的核心优势。

2.3 多频率设计

768维向量被拆分成384对，每对分配一个不同的旋转频率：

freqs = 1.0 / (1000000 ** (torch.arange(0, dim, 2) / dim))

• 高频维度：每步旋转角度大，擅长捕捉近距离的位置关系
• 低频维度：每步旋转角度极小，适合区分远距离位置

384个频率的组合，确保了任意两个位置的编码都是独一无二的。这就好比钟表上秒针、分针、时针的组合，能唯一标识一个时刻。高频与低频的配合，让模型既能感知短距离，也能精确分辨长距离的位置关系。

2.4 YaRN 长度外推

当推理时遇到的序列长度超过训练时的最大长度时，高频维度的角度会进入模型从未见过的范围。YaRN的思路是：将这些高频维度的频率降低（除以一个factor），把远距离位置的角度“压缩”回模型熟悉的范围。听起来复杂，实际上是一个非常工程化的解决方案。

三、GQA：分组查询注意力

3.1 标准多头注意力（MHA）的代价

在标准多头注意力中，Q、K、V各有8个头，每个头独立计算注意力。推理时，我们需要缓存所有K、V头的历史状态（即KV Cache），显存开销非常大。当模型规模变大时，这个开销几乎不可忽视。

3.2 GQA 的优化思路

MiniMind采用了一个很实用的设定：Q头数=8，KV头数=4。也就是说，每2个Q头共享1组KV：

self.q_proj = nn.Linear(768, 8 * 96)  # 8 个 Q 头
self.k_proj = nn.Linear(768, 4 * 96)  # 4 个 KV 头
self.v_proj = nn.Linear(768, 4 * 96)

计算时通过repeat_kv把4组KV复制扩展为8组，与Q头数匹配。这样既保留了多头注意力的表达能力，又大幅减少了KV Cache的占用。

3.3 收益：一半成本，几乎无损耗

KV Cache减半，推理速度和显存都明显改善。Google在65B规模上验证过，这种设置带来的效果损失不到0.5%——以忽略不计的代价换来了翻倍的效率提升，这笔账是划算的。

四、SwiGLU：门控前馈网络

4.1 FeedForward 的真实作用

Attention处理的是词与词之间的关系，而FeedForward是对每个位置独立进行知识加工。它的基本结构是：升维 → 激活 → 降维。

• 升维（768 → 2432）：把特征展开到更高维空间，便于分离不同模式
• 激活：引入非线性，让模型能学到复杂规律
• 降维（2432 → 768）：将筛选后的信息压缩回原始维度

4.2 SwiGLU 的独特之处

def forward(self, x):
    gate = F.silu(self.gate_proj(x))  # 门控信号：每个维度的"开关"
    info = self.up_proj(x)            # 信息流：实际内容
    return self.down_proj(gate * info) # 门控 × 信息 → 降维

相比普通的ReLU激活函数，SwiGLU引入了一个门控机制。模型可以精细控制“哪些信息通过、通过多少”，而不是ReLU那种非开即关的二值方式。这让信息筛选变得更加灵活可控。

五、残差连接

每个MiniMindBlock包含两次残差连接：

residual = hidden_states
hidden_states = self.self_attn(self.input_layernorm(hidden_states))
hidden_states += residual  # 第一次残差

hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))  # 第二次残差

每个子模块（Attention、FeedForward）都独立拥有一条残差通道。这保证了两个关键特性：

1. 梯度可以沿直通路径回传，避免深层网络中的梯度消失问题
2. 即使某层学得不够好，原始信息仍能无损地传递到下一层

六、因果遮罩与 KV Cache

6.1 因果遮罩

Decoder-Only架构的核心约束是：生成时只能看到“过去”，不能看到“未来”。实现上，通过一个上三角矩阵将未来位置的注意力分数设为负无穷：

scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf")).triu(1)

6.2 KV Cache

推理时，已经生成好的token的K、V是固定的、不会变化的。把它们缓存起来后，每一步只需要对新token进行计算即可，避免了大量重复运算。这是实际工程中大幅提升生成速度的关键设计。

七、模型参数设计

MiniMind-3的参数配置如下：

总参数量大约64M——只有GPT-3的1/2700。用小模型讲清楚大模型的核心设计，这正是MiniMind的设计初衷。

八、小结

回头来看，当前主流大语言模型的结构已经高度趋同：Decoder-Only + RoPE + GQA + SwiGLU + RMSNorm + Pre-Norm。这些组件的选择并非偶然，每一个都经过了充分的理论验证和实践检验。从根本上说，模型的本质就是一组参数矩阵。训练的过程，无非是通过优化算法将这些矩阵从随机值调整到能够准确预测下一个token的有意义数值。结构设计的核心，就是选择合适的计算组件，并用合理的方式连接它们，确保梯度传播顺畅、训练稳定、推理高效。这正是工程之美所在。