《Born》第15章:Transformer 架构深度解析
Transformer 是现代 NLP 和视觉模型的基石。这一章拆解它的每一个组件,对照 Born 的源码实现。
---
Transformer Block 结构
Input
├──► RMSNorm ──► GQA ────────────► Add ──┐
│ │
└──────────────────────────────────────────┘
├──► RMSNorm ──► SwiGLU FFN ─────► Add ──┐
│ │
└──────────────────────────────────────────┘
Output
---
多头注意力(MHA)
mha := nn.NewMHA[float32](dModel, nHead, backend)
// Q, K, V 投影
Q := mha.Wq.Forward(x) // [batch, seq, dModel]
K := mha.Wk.Forward(x)
V := mha.Wv.Forward(x)
// 分头
Q = Q.Reshape([batch, seq, nHead, headDim])
// 注意力得分
scores := Q.MatMul(K.T()) / sqrt(headDim)
scores = scores.Softmax(-1)
output := scores.MatMul(V)
---
RoPE:旋转位置编码
// RoPE 把位置信息编码进 Q/K 向量
func applyRoPE(x Tensor, pos int) Tensor {
// x: [seq, headDim]
// 对每个位置 m,应用旋转矩阵
for i := 0; i < headDim/2; i++ {
theta := math.Pow(10000, -2*float64(i)/float64(headDim))
angle := float64(pos) * theta
cosVal := math.Cos(angle)
sinVal := math.Sin(angle)
// 旋转二维子空间
x1, x2 := x[i], x[i+headDim/2]
x[i] = x1*cosVal - x2*sinVal
x[i+headDim/2] = x1*sinVal + x2*cosVal
}
}
RoPE 的优势:外推性好,长序列表现稳定。
---
GQA:分组查询注意力
// MHA: 每个头有自己的 K/V
// GQA: nHead 个 Q 共享 nKVHead 个 K/V
gqa := nn.NewGQA[float32](dModel, nHead, nKVHead, backend)
// nKVHead < nHead,减少 KV-Cache 内存
LLaMA 2/3 使用 GQA,内存占用减少 ~30%。
---
📘 《Born》连载技术书,第 15/22 章。
💬 讨论回复 (0)
推荐
🌟 智谱 GLM-5 已上线
我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用,智谱新一代旗舰模型 GLM-5 已上线,在推理、代码、智能体综合能力达到开源模型 SOTA 水平。
🎁 领取 2000万 Tokens