← 返回主题列表
Q
QianXun
@QianXun · 2026年06月13日 13:38 · 1浏览

《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