手撕 GPT#03：GPT 的核心代码只有 100 行，并且还支持注意力、权重、loss可视化哦

上篇文章我们用 CPU 训练了一个能回答中文问题的 GPT。

你可能好奇：模型里面到底在干什么？注意力机制是什么？位置编码是什么？因果掩码又是什么？

这些概念听起来很吓人，但 GPT 的核心逻辑只有约 100 行代码。

今天我们拆开看。配上真实注意力可视化，你会看到模型"在读你的问题时，眼睛在看哪里"。

一、GPT 在做什么？一句话说清

GPT 就是一个"猜下一个字"的概率预测机。

给它"太阳系有"，它猜"八大"。给它"八大"，它猜"行"。给它"行星"，它猜"："。一路猜下去，就变成了一句话。

训练的过程就是：给它看大量"问题-答案"对，让它学会在什么上下文后面应该猜什么字。

所有的复杂技术——注意力、位置编码、归一化——都是为了把这个"猜字"做得更好。

二、模型的流水线

数据从输入到输出，经过这几步：

输入文字 → 分词（切成 token）→ 词嵌入（变成数字向量） → 位置编码（告诉模型顺序） → ×4 层 Transformer Block（核心计算） → 归一化 → 预测下一个字

每一步都有清晰的代码。我们一个一个看。

三、拆解每一步

▪ 第 1 步：分词——把文字切成小块

计算机不认识"注意力机制"这几个字。分词器把它切成小块：

# "注意力机制通过计算相关性分配权重" # → 切成： # ["注意力", "机制", "通过", "计算", "相关性", "分配", "权重"] # → 变成数字： # [42, 187, 23, 89, 156, 201, 78]

我们用的 BPE（Byte Pair Encoding）算法会自动把经常一起出现的字组合成词。"注意力"出现多了就变成一个整体，不会被切碎成"注"、"意"、"力"。

这是项目的分词器可视化，你可以直观看到 BPE 的切分结果：

┌─────────┐┌────┐┌────┐┌────┐┌─────┐┌────┐┌────┐ │ 注意力 ││机制 ││通过 ││计算 ││相关性││分配 ││权重 │ │ ID:42 ││187 ││ 23 ││ 89 ││ 156 ││ 201││ 78 │ └─────────┘└────┘└────┘└────┘└─────┘└────┘└────┘

▪ 第 2 步：词嵌入——把数字变成"意思向量"

分词后每个 token 只是一个 ID（数字）。词嵌入把这些 ID 变成有意义的向量：

# 词嵌入：把每个 token ID 映射成一个 256 维的向量 self.tok_emb = nn.Embedding(vocab_size, n_embd) # 256 维

为什么需要向量？因为 ID 42 和 ID 43 之间没有关系。但向量 [0.8, 0.2, ...] 和 [0.7, 0.3, ...] 可以表示"意思相近"。

词嵌入让模型理解"注意力"和"机制"是有关系的，而"注意力"和"蒸馏水"没有关系。

▪ 第 3 步：位置编码——告诉模型"谁在前面"

词嵌入不知道顺序。"我爱你"和"你爱我"的词一样，但意思不同。

我们用 RoPE（旋转位置编码）解决这个问题：

def rope(q, k, seq_len, head_dim, device): # 给每个位置算一个旋转角度 pos = torch.arange(seq_len) # 位置 0, 1, 2, ... rates = torch.pow(10000, -2 * idx / head_dim) #旋转频率 theta = pos * rates # 每个位置的旋转角 # 把 Q 和 K 按位置旋转 # 位置 0 的词不旋转，位置 1 的词转一点，位置 2 转更多 cos = torch.cos(theta) sin = torch.sin(theta) return apply_rotation(q, k, cos, sin)

RoPE 的巧妙之处：它让模型通过 Q 和 K 的相对旋转角度就能知道"这两个词隔了多远"，不需要额外的位置参数。

▪ 第 4 步：Transformer Block——核心计算（~50 行）

这是 GPT 的心脏。一个 Block 包含两部分：

注意力 + 前馈网络，各带残差连接：

class Block(nn.Module): def forward(self, x): # 1. 注意力：让每个词"看"其他词 attn_out = self.attn(self.norm1(x))#先归一化，再算注意力 x = x + attn_out # 残差连接：加上原来的 x # 2. 前馈网络：让模型有"思考"的能力 ff_out = self.mlp(self.norm2(x)) # 先归一化，再过前馈 x = x + ff_out # 残差连接 return x

注意力解决的问题是：当模型看到"它"这个词时，怎么知道"它"指的是太阳还是行星？

答案是：每个词都去"看"其他所有词（只能看前面的词），计算一个"关注度"分数。

# 注意力核心逻辑（简化版） q = W_q(x) # Query：我在找什么？ k = W_k(x) # Key：我有什么？ v = W_v(x) # Value：我的内容是什么？ score = q @ k.T / sqrt(d) # Q 和 K 的相似度 score = score + causal_mask # 因果掩码：只看前面的词 weight = softmax(score) # 变成概率分布 output = weight @ v # 按权重组合 Value

因果掩码是什么？一个下三角矩阵：

1 0 0 0 ← 位置 1 只能看到自己 1 1 0 0 ← 位置 2 能看到 1 和 2 1 1 1 0 ← 位置 3 能看到 1、2、3 1 1 1 1 ← 位置 4 能看到所有

为什么？因为 GPT 是逐字生成的。写第 3 个字的时候，你还没写第 4 个字呢，当然看不到。

前馈网络（SwiGLU）给模型"思考"能力：

# SwiGLU：门控机制，比普通 ReLU 更强 gate = silu(W_gate(x)) # 门：决定放多少信息通过 up = W_up(x) # 上投影：扩展维度 return W_down(gate * up) # 下投影：压缩回来

▪ 第 5 步：输出预测

经过 4 层 Block 后，最后一个线性层把向量映射回词表：

x = self.norm(x) # 归一化 logits = self.head(x) # 映射回词表大小 # logits[i] = 第 i 个词是下一个词的概率（未归一化）

取概率最高的词，就是模型的预测。

四、真实注意力可视化

上面讲的是原理。下面给你看真实数据——模型在回答问题时，注意力到底在看哪里。

问模型："什么是注意力机制？"

这是第 4 层（最后一层）的注意力热力图：

用:什 是 注 意 力 机 制 ？ 助 手:注 意 力 机 制 通 过 计 算 用: ■ □ 什: ■ ■ □ 是: ■ ■ ■ □ 注: ■ ■ ■ ■ ■ 意: ■ ■ ■ ■ ■ ■ ■ ■ 力: ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ 机: ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ 制: ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ？: ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ 助手: □ □ □ □ □ □ □ □ □ ■ ■ ■ ■ ■ ■ ■ 注: □ □ □ ■ □ □ □ □ □ □ ■ ■ ■ ■ ■ ■ ■ 意: □ □ □ □ ■ □ □ □ □ □ □ ■ ■ ■ ■ ■ ■ ■ ...

你能看到一个清晰的规律：输出部分的"注意力"一词，强烈关注输入部分的"注意力"一词。模型在"回头看"问题中的关键词来生成答案。

这就是注意力机制在做的事：让模型在写每个字的时候，知道该参考前面的哪个字。

你可以用项目自带的可视化工具生成真实的注意力热力图：

uv run python -m src.visualize --only real_attention --prompt "什么是注意力机制？"

五、为什么是 Llama 风格？

我们的模型用了几个 Llama 的技术。每个都解决一个具体问题：

每个技术都是为了解决一个具体问题，不是为了花哨。

六、100 行代码能跑，为什么大模型要几千亿参数？

因为"能用"和"好用"之间差了几个数量级：

• 3M 参数：能记住 4-7 个字的答案，18 个字就开始忘
• 8B 参数（Llama 3 小号）：能写文章、写代码
• 175B 参数（GPT-3）：能做推理、写小说

模型大小决定的是"能装多少知识"和"能做多少推理"，不改变核心原理。

所以用 3M 模型学习原理是最高效的——架构一样，代码看得懂，训练跑得动。

七、自己动手看

# 克隆项目 git clone https://github.com/helloworldtang/GPT_teacher-3.37M-cn.git cd GPT_teacher-3.37M-cn # 看核心代码 cat src/model.py # 模型定义，约 193 行 # 生成可视化 uv sync uv run python -m src.visualize # 全部可视化 uv run python -m src.visualize --only real_attention#真实注意力 uv run python -m src.visualize --only tokenizer # 分词过程 uv run python -m src.visualize --only causal # 因果掩码

生成的图表：

核心代码在 src/model.py，193 行，带注释。

建议配合本文一起读。

这是「手撕 GPT」系列第 3 篇。上一篇：《从乱码到说人话，模型经历了什么》。下一篇：《我训练了一个满分模型，问它一个问题，后悔了》。

项目地址：https://github.com/helloworldtang/GPT_teacher-3.37M-cn