03-Transformer架构从零搭建

阶段 3：Transformer 架构从零搭建

教学目标

• 深入理解 Transformer Decoder-Only 架构的每个组件
• 掌握 RMSNorm、RoPE、GQA、SwiGLU、MoE 的数学原理与实现
• 能够从零手写一个简化版的 Transformer Block
• 理解 MiniMind 代码中各组件的实现细节

课时安排

• 理论讲解：4 课时（各组件原理 + 数学推导）
• 代码研读：2 课时（源码分析 + 手写实现）

3.1 Transformer 架构总览

3.1.1 Decoder-Only vs Encoder-Decoder

MiniMind 采用 Decoder-Only 架构，与 GPT 系列和 Qwen 系列一致。这种架构通过"因果注意力掩码"（Causal Mask）确保模型只能看到当前及之前的位置，天然适合自回归文本生成。

3.1.2 MiniMind-3 完整架构图

输入 Token IDs: [t₁, t₂, t₃, …, tₙ]
│
▼
┌─────────────────────────────────┐
│ Token Embedding │ vocab_size(6400) × dim(768)
│ 将 token ID 映射为向量 │
└─────────────┬───────────────────┘
│
▼
┌─────────────────────────────────┐
│ RoPE 位置编码 │ 旋转位置编码（非学习式）
│ 注入位置信息 │ 支持 YaRN 长度外推
└─────────────┬───────────────────┘
│
▼
┌─────────────────────────────────┐
│ Transformer Block × 8 │
│ ┌───────────────────────────┐ │
│ │ RMSNorm │ │ Pre-Norm
│ │ └──> Attention (GQA) │ │ q=8头, kv=4头
│ │ └──> 残差连接 │ │
│ │ RMSNorm │ │
│ │ └──> MLP (SwiGLU/MoE) │ │ Dense 或 MoE
│ │ └──> 残差连接 │ │
│ └───────────────────────────┘ │
└─────────────┬───────────────────┘
│
▼
┌─────────────────────────────────┐
│ Final RMSNorm │ 输出归一化
└─────────────┬───────────────────┘
│
▼
┌─────────────────────────────────┐
│ LM Head │ dim(768) → vocab_size(6400)
│ 投射到词表空间 │ (通常与 Embedding 共享权重)
└─────────────┬───────────────────┘
│
▼
输出 Logits: [logits₁, logits₂, …, logitsₙ]
│
▼ Cross-Entropy Loss
预测下一个 Token

3.2 嵌入层（Embedding）

3.2.1 Token Embedding

将每个 token ID 映射为一个固定维度的向量：

数学定义：
E(x) = W_E[x] ∈ ℝ^{d_model}

其中：
W_E ∈ ℝ^{vocab_size × d_model} （词嵌入矩阵）
x ∈ {0, 1, …, vocab_size-1} （token ID）
d_model = 768 （MiniMind-3 的隐藏维度）

MiniMind 中的实现（简化版）：

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
def __init__(self, vocab_size, dim):
super().__init__()
self.embedding = nn.Embedding(vocab_size, dim)

def forward(self, tokens):
# tokens: (batch, seq_len) → (batch, seq_len, dim)
return self.embedding(tokens)

3.2.2 参数量计算

Token Embedding 参数 = vocab_size × d_model
= 6400 × 768
= 4,915,200
≈ 4.9M（占总参数的 7.7%）

3.2.3 权重共享（Weight Tying）

MiniMind 的 Token Embedding 和 LM Head 通常共享权重：

LM Head: logits = x · W_E^T （直接用 Embedding 矩阵的转置）

优点：

• 减少参数量（省去一个 vocab_size × dim 的矩阵）
• 隐式正则化（Embedding 空间和输出空间对齐）

3.3 位置编码：RoPE（Rotary Position Embedding）

3.3.1 为什么需要位置编码？

Transformer 的核心——自注意力机制——本身是位置无关的（置换不变性）。如果不注入位置信息，模型无法区分 “我爱你” 和 “你爱我”。

3.3.2 传统位置编码的局限

3.3.3 RoPE 核心思想

RoPE（Rotary Position Embedding）通过旋转矩阵将位置信息编码到 Query 和 Key 中，而非作为额外的位置向量相加。

核心直觉：将每个位置的 Query/Key 向量"旋转"一个与位置相关的角度。

3.3.4 RoPE 数学推导

对于维度 d 的向量中的每一对元素 (q₁, q₂)，定义旋转操作：

对于第 m 个位置，第 i 对维度：

旋转角度: θ_m,i = m · θ_i

其中 θ_i = base^(-2i/d)，base 通常为 10000（MiniMind 使用 1e6）

旋转矩阵:
R(θ_m,i) = | cos(θ_m,i) -sin(θ_m,i) |
| sin(θ_m,i) cos(θ_m,i) |

旋转后的 Query/Key:
q'_m = R(θ_m) ⊗ q_m
k'_n = R(θ_n) ⊗ k_n

关键性质：旋转后，Q 和 K 的点积自然包含相对位置信息：

q'_m · k'_n = q_m · k_{n-m} （仅依赖相对位置 m-n）

3.3.5 RoPE 代码实现

import torch
import torch.nn as nn
import math

class RotaryPositionEmbedding(nn.Module):
def __init__(self, dim, max_seq_len=32768, base=1000000.0):
super().__init__()
# 计算频率
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

# 预计算所有位置的旋转角度
self.register_buffer(
"inv_freq",
inv_freq,
)

# 预计算 cos 和 sin（用于缓存）
t = torch.arange(max_seq_len).float()
freqs = torch.outer(t, inv_freq) # (max_seq_len, dim/2)
self.register_buffer("cos_cached", freqs.cos())
self.register_buffer("sin_cached", freqs.sin())

def forward(self, x, seq_len):
"""
x: (batch, n_heads, seq_len, head_dim)
返回旋转后的 x
"""
cos = self.cos_cached[:seq_len] # (seq_len, dim/2)
sin = self.sin_cached[:seq_len]

# 将 x 拆分为相邻对
x1, x2 = x[..., ::2], x[..., 1::2]

# 应用旋转
# [cos, -sin] [x1]
# [sin, cos] [x2]
x_rotated = torch.stack([
x1 * cos – x2 * sin,
x1 * sin + x2 * cos
], dim=–1)

# 重新交织
return x_rotated.flatten(–2)

3.3.6 RoPE 长度外推：YaRN

当推理长度超过训练时的最大位置编码时，RoPE 的外推性能会急剧下降。YaRN（Yet another RoPE extensioN）通过以下方式解决：

核心方法：动态调整 RoPE 的 base 值

原始: θ_i = base^(-2i/d)
外推: θ_i = (base × factor)^(-2i/d)

其中 factor > 1，使得旋转角度变化更缓慢，从而覆盖更长的序列。

MiniMind 支持通过配置开启 YaRN 外推：

# 在推理时开启
python eval_llm.py ––weight full_sft ––inference_rope_scaling

3.4 归一化：RMSNorm（Root Mean Square Normalization）

3.4.1 从 LayerNorm 到 RMSNorm

3.4.2 RMSNorm 数学定义

RMSNorm(x) = x / RMS(x) · γ

其中:
RMS(x) = √(mean(x²) + ε) （均方根，ε=1e-6 防除零）
γ ∈ ℝ^{d_model} （可学习的缩放参数）

3.4.3 RMSNorm 代码实现

class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
# 可学习的缩放参数，初始化为 1
self.gamma = nn.Parameter(torch.ones(dim))

def forward(self, x):
# x: (batch, seq_len, dim)
# 计算均方根
rms = torch.sqrt(torch.mean(x ** 2, dim=–1, keepdim=True) + self.eps)
return x / rms * self.gamma

3.4.4 Pre-Norm vs Post-Norm

MiniMind 使用 Pre-Norm（预归一化）：

Post-Norm（原始 Transformer）：
x' = LayerNorm(x + SubLayer(x))

Pre-Norm（MiniMind 使用）：
x' = x + SubLayer(RMSNorm(x))

Pre-Norm 的优势：

• 训练更稳定（梯度流更顺畅）
• 不需要学习率预热（或只需极少预热）
• 深层网络训练效果更好

3.5 注意力机制：GQA（Grouped-Query Attention）

3.5.1 注意力机制回顾

标准 Multi-Head Attention（MHA）：

Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V

其中:
Q = X · W_Q X ∈ ℝ^{n×d}, W_Q ∈ ℝ^{d×(h×d_k)}
K = X · W_K W_K ∈ ℝ^{d×(h×d_k)}
V = X · W_V W_V ∈ ℝ^{d×(h×d_k)}
h = 注意力头数
d_k = d_model / h

3.5.2 从 MHA 到 GQA 的演进

Multi-Head Attention (MHA) 每个 Q 头有独立的 K, V 头
q_heads=8, kv_heads=8 KV Cache: h × seq_len × d_k

Multi-Query Attention (MQA) 所有 Q 头共享 1 组 K, V
q_heads=8, kv_heads=1 KV Cache: 1 × seq_len × d_k

Grouped-Query Attention (GQA) Q 头分组，每组共享 K, V
q_heads=8, kv_heads=4 KV Cache: (kv_heads) × seq_len × d_k

3.5.3 MiniMind-3 的 GQA 配置

q_heads = 8 # 8 个 Query 头
kv_heads = 4 # 4 个 KV 头
head_dim = 768/8 = 96

分组方式:
Q头 [0,1] → KV组 0
Q头 [2,3] → KV组 1
Q头 [4,5] → KV组 2
Q头 [6,7] → KV组 3

KV Cache 减少: 8/4 = 2 倍

3.5.4 GQA 代码实现

class GroupedQueryAttention(nn.Module):
def __init__(self, dim, n_heads=8, n_kv_heads=4):
super().__init__()
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = dim // n_heads
self.n_groups = n_heads // n_kv_heads # 每个 KV 组服务的 Q 头数

# Q, K, V 投影
self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
self.wv = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False)

def forward(self, x, freqs_cis, mask=None):
"""
x: (batch, seq_len, dim)
freqs_cis: RoPE 频率
mask: 因果注意力掩码
"""
bsz, seqlen, _ = x.shape

# Q, K, V 投影
q = self.wq(x) # (bsz, seqlen, n_heads * head_dim)
k = self.wk(x) # (bsz, seqlen, n_kv_heads * head_dim)
v = self.wv(x) # (bsz, seqlen, n_kv_heads * head_dim)

# reshape 为多头格式
q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
k = k.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
v = v.view(bsz, seqlen, self.n_kv_heads, self.head_dim)

# 应用 RoPE 位置编码
q = apply_rotary_emb(q, freqs_cis)
k = apply_rotary_emb(k, freqs_cis)

# GQA: 扩展 K, V 以匹配 Q 头数
# K/V 的每个头被复制 n_groups 次
k = k.repeat_interleave(self.n_groups, dim=2) # (bsz, seqlen, n_heads, head_dim)
v = v.repeat_interleave(self.n_groups, dim=2)

# 计算注意力分数
q = q.transpose(1, 2) # (bsz, n_heads, seqlen, head_dim)
k = k.transpose(1, 2)
v = v.transpose(1, 2)

scores = torch.matmul(q, k.transpose(–2, –1)) / math.sqrt(self.head_dim)

# 应用因果掩码（下三角矩阵）
if mask is not None:
scores = scores + mask

attn = torch.softmax(scores.float(), dim=–1).type_as(q)
out = torch.matmul(attn, v)

# 输出投影
out = out.transpose(1, 2).contiguous().view(bsz, seqlen, –1)
return self.wo(out)

3.5.5 因果掩码（Causal Mask）

def create_causal_mask(seq_len):
"""
创建因果注意力掩码（上三角为 -inf）
确保位置 i 只能看到位置 ≤i 的内容
"""
mask = torch.full((seq_len, seq_len), float("-inf"))
mask = torch.triu(mask, diagonal=1) # 上三角为 -inf
return mask

# 示例：seq_len = 5 的因果掩码
# 0 -inf -inf -inf -inf
# 0 0 -inf -inf -inf
# 0 0 0 -inf -inf
# 0 0 0 0 -inf
# 0 0 0 0 0

3.6 前馈网络：SwiGLU

3.6.1 从 ReLU 到 SwiGLU 的演进

3.6.2 SwiGLU 数学定义

SwiGLU 是一种门控线性激活函数，结合了 Swish（SiLU）和门控机制：

SwiGLU(x) = SiLU(xW₁) ⊗ (xW₂)

其中:
SiLU(x) = x · σ(x) （Swish 函数）
σ(x) = 1/(1+e^(-x)) （Sigmoid 函数）
W₁ ∈ ℝ^{d_model × hidden_dim}
W₂ ∈ ℝ^{d_model × hidden_dim}
W₃ ∈ ℝ^{hidden_dim × d_model} （输出投影）
⊗ 表示逐元素相乘（Hadamard积）

完整 MLP:
output = SwiGLU(x) · W₃
= (SiLU(xW₁) ⊗ xW₂) · W₃

3.6.3 SwiGLU 直觉理解

输入 x
│
├──> W₁ → SiLU → ──┐ （门控信号：决定"开放"程度）
│ ⊗ → W₃ → 输出
└──> W₂ ────────────┘ （信息流：决定"传递"什么）

门控机制的作用：让模型学会"选择性地传递信息"，而非简单地全部激活或全部关闭。

3.6.4 SwiGLU 代码实现

class SwiGLUMLP(nn.Module):
def __init__(self, dim, hidden_dim):
super().__init__()
self.w1 = nn.Linear(dim, hidden_dim, bias=False) # 门控分支
self.w2 = nn.Linear(dim, hidden_dim, bias=False) # 信息分支
self.w3 = nn.Linear(hidden_dim, dim, bias=False) # 输出投影

def forward(self, x):
# SiLU(xW₁) ⊗ (xW₂)
return self.w3(torch.nn.functional.silu(self.w1(x)) * self.w2(x))

3.6.5 参数量分析

标准 MLP (ReLU):
参数 = d_model × hidden_dim + hidden_dim × d_model
= 2 × d_model × hidden_dim
= 2 × 768 × 3072 = 4,718,592

SwiGLU MLP:
参数 = 3 × d_model × hidden_dim
= 3 × 768 × 3072 = 7,077,888

SwiGLU 比 ReLU MLP 多 50% 参数，但效果更好。

3.7 Transformer Block 组装

3.7.1 完整 Block 结构

输入 x
│
▼
RMSNorm(x) ──────────────> GQA(q,k,v + RoPE, causal_mask) ──> + x (残差)
│
▼
RMSNorm(x') ──> SwiGLU(x') ──> + x' (残差)
│
▼
输出 x''

3.7.2 代码实现

class TransformerBlock(nn.Module):
def __init__(self, dim, n_heads, n_kv_heads, hidden_dim):
super().__init__()
self.attention = GroupedQueryAttention(dim, n_heads, n_kv_heads)
self.feed_forward = SwiGLUMLP(dim, hidden_dim)
self.attention_norm = RMSNorm(dim)
self.ffn_norm = RMSNorm(dim)

def forward(self, x, freqs_cis, mask):
# Pre-Norm + Attention + 残差
x = x + self.attention(self.attention_norm(x), freqs_cis, mask)
# Pre-Norm + FFN + 残差
x = x + self.feed_forward(self.ffn_norm(x))
return x

3.8 MoE（Mixture of Experts）

3.8.1 MoE 核心思想

MoE 将标准的前馈网络替换为多个专家网络，每次只激活其中少数几个：

标准 FFN:
x → [FFN] → output （所有参数都参与计算）

MoE FFN:
x → Router → 选择 Top-K 个专家
│
├── Expert₁ ──┐
├── Expert₂ ──┤ → 加权求和 → output
├── Expert₃ ──┤
└── Expert₄ ──┘ （仅 K 个专家的参数参与计算）

3.8.2 MiniMind-3 MoE 配置

3.8.3 Router 机制

class Router(nn.Module):
def __init__(self, dim, num_experts):
super().__init__()
self.gate = nn.Linear(dim, num_experts, bias=False)

def forward(self, x):
"""
x: (batch, seq_len, dim)
返回: 路由权重 (batch, seq_len, num_experts)
"""
return torch.softmax(self.gate(x), dim=–1)

3.8.4 Top-K 路由

def top_k_routing(router_output, k=1):
"""
router_output: (batch, seq_len, num_experts)
k: 选择 top-k 个专家
"""
topk_weights, topk_indices = torch.topk(router_output, k, dim=–1)
# 归一化 top-k 权重
topk_weights = topk_weights / topk_weights.sum(dim=–1, keepdim=True)
return topk_weights, topk_indices

3.8.5 MoE vs Dense 参数对比

Dense (minimind-3):
每层 FFN 参数: 3 × 768 × 3072 = 7,077,888
8 层总计: 8 × 7,077,888 ≈ 56.6M

MoE (minimind-3-moe):
每个专家 FFN 参数: 3 × 768 × 3072 = 7,077,888
4 个专家: 4 × 7,077,888 = 28,311,552
8 层总计: 8 × 28,311,552 ≈ 226.5M（实际约 198M，因部分层共享）
每次活跃参数: 8 × 7,077,888 ≈ 56.6M（与 Dense 相当）

关键理解：MoE 的优势是总参数量大（模型容量大），但每次推理只激活一小部分（计算量与 Dense 相当）。

3.9 MiniMind 完整模型类

3.9.1 ModelArgs 配置

@dataclass
class ModelArgs:
dim: int = 768 # 隐藏维度
n_layers: int = 8 # Transformer 层数
n_heads: int = 8 # Q 头数
n_kv_heads: int = 4 # KV 头数
vocab_size: int = 6400 # 词表大小
max_seq_len: int = 32768 # 最大序列长度
hidden_dim: int = 3072 # FFN 隐藏维度
rope_base: float = 1000000.0 # RoPE 基数
norm_eps: float = 1e-6 # RMSNorm epsilon
num_experts: int = None # MoE 专家数（None 表示 Dense）

3.9.2 MiniMindModel 代码结构

class MiniMindModel(nn.Module):
def __init__(self, args: ModelArgs):
super().__init__()
self.args = args

# Token Embedding
self.embedding = nn.Embedding(args.vocab_size, args.dim)

# RoPE 位置编码
self.rope = RotaryPositionEmbedding(
args.dim // args.n_heads,
args.max_seq_len,
args.rope_base
)

# Transformer Blocks
self.layers = nn.ModuleList([
TransformerBlock(
dim=args.dim,
n_heads=args.n_heads,
n_kv_heads=args.n_kv_heads,
hidden_dim=args.hidden_dim,
num_experts=args.num_experts
) for _ in range(args.n_layers)
])

# Final RMSNorm
self.norm = RMSNorm(args.dim)

# LM Head（通常与 Embedding 共享权重）
self.lm_head = nn.Linear(args.dim, args.vocab_size, bias=False)

def forward(self, tokens, mask=None):
"""
tokens: (batch, seq_len) – token IDs
返回: (batch, seq_len, vocab_size) – logits
"""
bsz, seqlen = tokens.shape

# Embedding
h = self.embedding(tokens)

# RoPE
freqs_cis = self.rope(h, seqlen)

# 因果掩码
if mask is None:
mask = create_causal_mask(seqlen)

# Transformer Blocks
for layer in self.layers:
h = layer(h, freqs_cis, mask)

# Final Norm
h = self.norm(h)

# LM Head
logits = self.lm_head(h)

return logits

3.10 参数量完整计算

3.10.1 MiniMind-3 Dense 参数分解

实际约 64M，差异来自 hidden_dim 和其他优化细节。

实践任务

任务 3.1：从零手写 Transformer Block（必做）

不依赖任何框架的高级 API（仅使用 torch.nn.Linear、torch.nn.Embedding、torch.nn.Parameter），从零实现一个简化版 Transformer Block：

# 要求：
# 1. 实现 RMSNorm
# 2. 实现标准 Multi-Head Attention（先不做 GQA）
# 3. 实现 SwiGLU MLP
# 4. 组装为 TransformerBlock
# 5. 验证：输入随机 tensor，检查输出 shape 正确

# 配置：
dim = 256
n_heads = 4
hidden_dim = 1024
seq_len = 128
batch_size = 2

任务 3.2：源码研读对比（必做）

阅读 MiniMind 项目中的以下文件，与自己的实现对比：

记录以下差异：

• MiniMind 的 RoPE 实现与你的实现有什么不同？
• GQA 的 repeat_interleave 操作是如何工作的？
• MiniMind 的 RMSNorm 是否有额外的优化？

任务 3.3：参数量验证（必做）

# 编写脚本，计算 MiniMind-3 的实际参数量：
# 1. 加载 MiniMind-3 模型
# 2. 逐模块统计参数量
# 3. 与 3.10.1 的理论计算对比
# 4. 分析差异原因

from model.model import MiniMindModel, ModelArgs

args = ModelArgs()
model = MiniMindModel(args)

# 统计代码…

任务 3.4：MoE 实现分析（选做）

挑战任务（选做）

• 实现 YaRN 长度外推，测试在 2× 训练长度上的推理效果
• 尝试修改模型结构（如增加层数、调整 head 数），训练一个自定义配置的模型

下一阶段预告

在阶段 4 中，我们将正式进入**预训练（Pretrain）**阶段。学习如何使用大规模无标注文本训练语言模型，理解 next-token prediction 的训练目标，并亲手运行一次完整的预训练流程。