Karpathy 的 243 行纯 Python GPT:人人都能看懂的教程
Karpathy 的 243 行纯 Python GPT:人人都能看懂的教程
Karpathy 只用了 243 行纯 Python(仅依赖 os/math/random),实现了 GPT 的完整训练和推理。本文逐章拆解每一行代码背后的原理。
先来一个总体类比
想象你在教一个外星人学写地球人的名字。这个外星人:
- 完全不认识字母(需要一个"翻译词典"→ Tokenizer)
- 脑子是一张白纸(需要随机初始化参数)
- 通过看大量名字来学习规律(训练循环)
- 每看一个名字,就微调自己的"直觉"(反向传播 + 优化器)
- 学完之后,可以自己编名字(推理)
Karpathy 的这 243 行代码,就是这整个过程的最小完整实现。
全局架构图
┌─────────────────────────────────────────────────────────────────┐
│ 整体流程 (243 行的故事) │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ① 数据准备 "Emma", "Olivia", "Ava" ... │
│ ↓ │
│ ② Tokenizer "Emma" → [0, 5, 13, 13, 1] │
│ ↓ │
│ ③ 自动微分引擎 Value 类:记录每一步计算,自动求导 │
│ ↓ │
│ ④ GPT 模型 Embedding → RMSNorm → Attention → MLP │
│ ↓ │
│ ⑤ 训练循环 前向传播 → 计算损失 → 反向传播 → Adam 更新 │
│ ↓ │
│ ⑥ 推理生成 一个字母一个字母地"编"出新名字 │
│ │
└─────────────────────────────────────────────────────────────────┘
第一章:数据与 Tokenizer —— 给字母编号
类比
你要教一台计算器认字。但计算器只认数字,不认字母。所以第一步:给每个字母一个编号。
代码解读
chars = ['<BOS>'] + sorted(set(''.join(docs)))
stoi = { ch:i for i, ch in enumerate(chars) } # 字母 → 数字
itos = { i:ch for i, ch in enumerate(chars) } # 数字 → 字母
图解
字符表 chars:
┌───────┬───┬───┬───┬───┬───┬─────┐
│ <BOS> │ a │ b │ c │ d │ e │ ... │
│ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ ... │
└───────┴───┴───┴───┴───┴───┴─────┘
编码 "cat": c→3, a→1, t→20 → [3, 1, 20]
解码 [3,1,20]: 3→c, 1→a, 20→t → "cat"
<BOS>(Beginning Of Sequence)是一个特殊标记,表示"开始/结束"。每个名字前后都加上它,就像给句子加上句号。
这就是所谓的"字符级 Tokenizer"——最朴素的分词方式,一个字母就是一个 token。GPT-4 用的 BPE tokenizer 会把常见词组合在一起(比如 "the" 是一个 token),但原理相通。
第二章:自动微分引擎 —— 整个项目的灵魂
类比:连锁反应追踪器
想象一个多米诺骨牌链。你推倒最后一块(loss),想知道:如果我稍微移动第一块骨牌的位置,最后一块会怎么变? 这就是"求梯度"。
Value 类就是一个"智能骨牌"——它不仅记录自己的值,还记录是谁创造了自己以及怎么创造的,这样就能沿着链条反向追踪。
核心结构
class Value:
def __init__(self, data, _children=(), _op=''):
self.data = data # 当前值(前向)
self.grad = 0 # 梯度(反向时填充)
self._backward = lambda: None # 反向传播函数
self._prev = set(_children) # 父节点们
图解:一次加法的计算图
前向传播 (计算值) 反向传播 (计算梯度)
─────────────→ ←─────────────────
a = Value(2.0) a.grad += out.grad = 1.0
↘ ↗
+ ──→ out = Value(5.0) out.grad = 1.0
↗ ↖
b = Value(3.0) b.grad += out.grad = 1.0
规则: 加法的梯度直接"复制"给两个输入
图解:一次乘法的计算图
a = Value(2.0) a.grad += b.data * out.grad
↘ = 3.0 * 1.0 = 3.0
× ──→ out = Value(6.0) out.grad = 1.0
↗ ↖
b = Value(3.0) b.grad += a.data * out.grad
= 2.0 * 1.0 = 2.0
规则: 乘法的梯度是"交叉相乘"——对 a 求导得到 b,对 b 求导得到 a
backward() 方法:拓扑排序 + 反向链式法则
def backward(self):
topo = []
visited = set()
def build_topo(v): # 递归构建拓扑排序
if v not in visited:
visited.add(v)
for child in v._prev:
build_topo(child)
topo.append(v)
build_topo(self)
self.grad = 1 # 起点: d(loss)/d(loss) = 1
for v in reversed(topo): # 反向遍历,逐个调用 _backward
v._backward()
为什么需要拓扑排序? 因为梯度要从后往前传,必须保证一个节点的梯度算完了,才去算它的输入的梯度。拓扑排序恰好保证这个顺序。
拓扑排序的结果 (前向顺序):
a → b → c → d → loss
反转后 (反向传播顺序):
loss → d → c → b → a
先知道 loss 的梯度,才能算 d 的梯度,依此类推
支持的运算一览
| 运算 | 前向 | 反向(梯度规则) |
|---|---|---|
a + b |
a.data + b.data |
两边都加 out.grad |
a * b |
a.data * b.data |
交叉:a.grad += b.data * out.grad |
a ** n |
a.data ** n |
幂规则:n * a^(n-1) * out.grad |
a.log() |
log(a.data) |
(1/a.data) * out.grad |
a.exp() |
exp(a.data) |
exp(a.data) * out.grad(自己就是自己的导数!) |
a.relu() |
max(0, a.data) |
正则传,负则断:(a>0) * out.grad |
常见误区
"为什么 grad 是累加 += 而不是赋值 =?"
因为一个 Value 可能被用在多个地方!比如 y = a + a,a 被用了两次,来自两条路径的梯度要加在一起,不能互相覆盖。这就是多元链式法则中的加法规则。
第三章:模型参数 —— GPT 的"大脑"
类比
如果 GPT 是一个人,参数就是他的神经突触连接强度。初始时随机生成(婴儿状态),通过训练慢慢调整。
架构参数
n_embd = 16 # 每个 token 用 16 个数字表示
n_head = 4 # 4 个注意力头(4 个"视角"同时看)
n_layer = 1 # 1 层 Transformer(真正的 GPT-2 有 12~48 层)
block_size = 8 # 最多看 8 个字符的上下文
head_dim = 4 # 每个头看 4 维 (16 / 4)
参数矩阵一览
state_dict (所有参数):
┌──────────────────────────────────────────────────────────┐
│ wte: [28 × 16] 词嵌入矩阵(28个字符,每个16维向量) │
│ wpe: [8 × 16] 位置嵌入矩阵(8个位置,每个16维向量) │
│ lm_head: [28 × 16] 输出投影(把16维映射回28个字符的概率) │
├──────────────────────────────────────────────────────────┤
│ layer0.attn_wq: [16 × 16] Query 投影 │
│ layer0.attn_wk: [16 × 16] Key 投影 │
│ layer0.attn_wv: [16 × 16] Value 投影 │
│ layer0.attn_wo: [16 × 16] 输出投影 (初始化为0!) │
│ layer0.mlp_fc1: [64 × 16] MLP 第一层 (升维 4x) │
│ layer0.mlp_fc2: [16 × 64] MLP 第二层 (降维回来, 初始化0!) │
└──────────────────────────────────────────────────────────┘
总参数量: 约 5000+ 个 Value 对象
注意 std=0 的初始化:attn_wo 和 mlp_fc2 初始化为全零。这意味着训练开始时,残差连接的"新增部分"为零,模型从"什么都不做"开始,逐渐学会做有用的变换。这是一个重要的训练稳定性技巧。
第四章:GPT 模型 —— 核心前向传播
总体类比
GPT 就像一个读心术选手在玩文字接龙:
- 看到已经写出的字母(embedding)
- 回忆之前所有字母之间的关系(attention)
- 深入思考(MLP)
- 猜下一个字母(logits → softmax → 概率)
完整前向传播流程图
输入: token_id=5 ('e'), pos_id=2
┌─────────────┐ ┌─────────────┐
│ wte[5] │ │ wpe[2] │
│ 词嵌入 16维 │ │ 位置嵌入 16维 │
└──────┬──────┘ └──────┬──────┘
│ │
└────────┬──────────┘
↓
┌─────────┐
│ 相加 │ x = tok_emb + pos_emb
└────┬────┘
↓
┌──────────┐
│ RMSNorm │ 归一化,稳定训练
└────┬─────┘
↓
╔══════════════════════════════╗
║ Transformer Layer 0 ║
║ ┌────────────────────────┐ ║
║ │ Multi-Head Attention │ ║
║ │ ┌──┐┌──┐┌──┐┌──┐ │ ║
║ │ │H0││H1││H2││H3│ │ ║
║ │ └──┘└──┘└──┘└──┘ │ ║
║ └───────────┬────────────┘ ║
║ + residual ║
║ ┌────────────────────────┐ ║
║ │ MLP (FFN) │ ║
║ │ 16→64→ReLU²→64→16 │ ║
║ └───────────┬────────────┘ ║
║ + residual ║
╚══════════════╪═══════════════╝
↓
┌────────────┐
│ lm_head │ 线性投影: 16维 → 28维
└─────┬──────┘
↓
logits [28个原始分数]
↓
┌──────────┐
│ softmax │ 转为概率分布
└─────┬────┘
↓
probs [28个概率值] → 选概率最高的作为下一个字母
4.1 Embedding:给 token 一个"身份证"
tok_emb = state_dict['wte'][token_id] # 查表:字母 → 16维向量
pos_emb = state_dict['wpe'][pos_id] # 查表:位置 → 16维向量
x = [t + p for t, p in zip(tok_emb, pos_emb)] # 加在一起
类比:wte 是"这个字母是什么"(身份),wpe 是"这个字母在第几个位置"(座位号)。加在一起 = 完整的身份信息。
"e" 在位置 2:
wte["e"] = [0.3, -0.1, 0.5, ...] ← "e"的本质特征
wpe[2] = [0.0, 0.2, 0.1, ...] ← 位置2的特征
──────────────────────────────────
x = [0.3, 0.1, 0.6, ...] ← "在位置2的e"
4.2 RMSNorm:稳定信号幅度
def rmsnorm(x):
ms = sum(xi * xi for xi in x) / len(x) # 均方值
scale = (ms + 1e-5) ** -0.5 # 1/sqrt(ms)
return [xi * scale for xi in x] # 缩放
类比:把每个向量的"音量"统一调到标准大小。不管输入是大声还是小声,输出的音量都差不多。这样就不会出现某些值过大或过小导致训练"翻车"的情况。
归一化前: [10.0, 20.0, 30.0, 40.0] ← 数值偏大
RMS = sqrt((100+400+900+1600)/4) = sqrt(750) ≈ 27.4
归一化后: [0.37, 0.73, 1.10, 1.46] ← 数值稳定
与 GPT-2 原版的 LayerNorm 相比,RMSNorm 不减均值,只做缩放,更简单更快。
4.3 Multi-Head Attention:最核心的部分
类比:图书馆查阅系统
想象你在图书馆找资料:
- Query (Q):你的问题——"我想找关于X的信息"
- Key (K):每本书封面上的关键词
- Value (V):每本书的实际内容
你拿着问题 Q 去和每本书的关键词 K 比对,匹配度高的书你就多读一些内容 V。
Multi-Head = 你同时派了 4 个助手去找,每个助手关注不同方面。
# 计算 Q, K, V
q = linear(x, state_dict[f'layer{li}.attn_wq']) # [16维]
k = linear(x, state_dict[f'layer{li}.attn_wk']) # [16维]
v = linear(x, state_dict[f'layer{li}.attn_wv']) # [16维]
keys[li].append(k) # 存入 KV 缓存
values[li].append(v) # 后续 token 还要回头看
图解:4 头注意力
完整 Q/K/V 向量 (16维) 被切成 4 段,每段 4 维:
Q = [q0 q1 q2 q3 | q4 q5 q6 q7 | q8 q9 q10 q11 | q12 q13 q14 q15]
╰─ Head 0 ─╯ ╰─ Head 1 ─╯ ╰─ Head 2 ──╯ ╰── Head 3 ──╯
每个 Head 独立做注意力计算:
Head 0 的注意力计算 (当前是位置 3,回看位置 0,1,2,3):
┌─────────────────────────────────────────────────────────┐
│ │
│ score[t] = Q_h · K_h[t] / √4 │
│ │
│ 位置 0: score=1.2 ──┐ │
│ 位置 1: score=0.3 ──┤── softmax ──→ [0.5, 0.1, 0.1, 0.3]│
│ 位置 2: score=0.1 ──┤ │
│ 位置 3: score=0.8 ──┘ │
│ │
│ output = 0.5·V[0] + 0.1·V[1] + 0.1·V[2] + 0.3·V[3] │
│ 加权平均! │
└─────────────────────────────────────────────────────────┘
4 个 Head 的输出拼接: [head0_out | head1_out | head2_out | head3_out]
╰──────────── 16 维 ────────────────╯
然后通过 attn_wo 线性投影
这里是 Causal (因果) Attention:位置 3 的 token 只能看到位置 0、1、2、3,不能偷看位置 4。这是通过 KV 缓存自然实现的——处理位置 3 时,keys[li] 里只存了位置 0~3 的 key,后面的还没算过,自然看不到。
4.4 MLP:深入思考
x = linear(x, state_dict[f'layer{li}.mlp_fc1']) # 16 → 64(升维)
x = [xi.relu() ** 2 for xi in x] # ReLU² 激活
x = linear(x, state_dict[f'layer{li}.mlp_fc2']) # 64 → 16(降维)
类比:Attention 是"收集信息",MLP 是"消化思考"。先把信息展开到更高维度(16→64,打开思路),经过非线性变换处理,再压缩回来(64→16,得出结论)。
MLP 的 "膨胀-压缩" 结构:
输入 ──→ [16维] ──→ fc1 ──→ [64维] ──→ ReLU² ──→ fc2 ──→ [16维] ──→ 输出
│ │ │
│ 展开思路 │ 非线性处理 │ 压缩结论
╰──────────────────╯ │
│
ReLU²(x) = max(0, x)² │
比普通 ReLU 更平滑,梯度连续 │
┌──────────────────────────────────┐
│ ReLU vs ReLU² │
│ ╱ ╱ │
│ ╱ ╱ ← 更平滑的曲线 │
│ ╱ ───╱ │
│ ────── ──── │
└──────────────────────────────────┘
4.5 残差连接:学习的保险绳
x = [a + b for a, b in zip(x, x_residual)] # 跳跃连接
类比:考试时,"新答案 = 原来的答案 + 修正量"。即使修正量学偏了,至少还有原来的答案保底。这就是为什么 attn_wo 和 mlp_fc2 初始化为 0——训练开始时修正量为零,模型先"什么都不做",然后逐渐学会有用的修正。
╭─────────────────────────╮
│ 直通(保底) │
x_in ─────┤ ├──→ x_out = x_in + f(x_in)
│ → Attention/MLP → f(x) │
╰─────────────────────────╯
即使 f(x) 全是 0,x_out = x_in,信号不会丢失
第五章:Softmax 与 Loss —— 如何衡量"猜得准不准"
Softmax:原始分数 → 概率
def softmax(logits):
max_val = max(val.data for val in logits) # 数值稳定性技巧
exps = [(val - max_val).exp() for val in logits] # 减最大值再 exp
total = sum(exps)
return [e / total for e in exps]
logits: [2.0, 1.0, 0.5, ..., -1.0] ← 28 个原始分数
↓ softmax
probs: [0.35, 0.13, 0.08, ..., 0.01] ← 28 个概率 (总和=1)
↓
选概率最高的 → 预测下一个字母是 'a'
减去 max_val 的技巧:exp(1000) 会溢出,但 exp(1000 - 1000) = exp(0) = 1,结果不变但数值安全。
交叉熵损失:猜错的代价
loss_t = -probs[target_id].log()
如果正确答案是 'm',模型给 'm' 的概率是 0.8:
loss = -log(0.8) = 0.22→ 低损失,猜得好!
如果模型给 'm' 的概率是 0.01:
loss = -log(0.01) = 4.6→ 高损失,猜得差!
损失函数的直觉:
loss ↑
5 │ ╲
4 │ ╲
3 │ ╲
2 │ ╲
1 │ ╲───
0 │ ───────
└──────────────────→ 对正确答案的预测概率
0 0.2 0.4 0.6 0.8 1.0
概率越高 → 损失越低 → 模型猜得越准
第六章:训练循环 —— 让模型变聪明
总流程
┌─────────────────────────────────────────────────────┐
│ 训练循环 (重复 500 次) │
│ │
步骤1 │ 取一个名字: "Emma" │
│ 加 BOS: [BOS, E, m, m, a, BOS] │
│ │
步骤2 │ 对每个位置做预测: │
│ BOS → 预测 E (答案: E) │
│ E → 预测 m (答案: m) │
│ m → 预测 m (答案: m) │
│ m → 预测 a (答案: a) │
│ a → 预测 BOS (答案: BOS, 即"该结束了") │
│ │
步骤3 │ 计算平均损失: loss = (l1+l2+l3+l4+l5) / 5 │
│ │
步骤4 │ 反向传播: loss.backward() │
│ 计算每个参数的梯度 │
│ │
步骤5 │ Adam 优化器更新所有参数 │
│ p.data -= lr * (梯度的聪明加权平均) │
│ │
步骤6 │ 清零梯度: p.grad = 0 │
│ │
└─────────────────────────────────────────────────────┘
Adam 优化器:聪明的梯度下降
类比:普通梯度下降像是在浓雾中下山,只能感觉到脚下的坡度。Adam 像是有一个带惯性的小球在下山:
m(一阶动量):小球的速度,是梯度的指数移动平均。让你保持方向稳定,不会因为一个噪声梯度就突然拐弯v(二阶动量):坡度的剧烈程度。坡度变化大的方向走小步,变化小的方向走大步
m[i] = 0.9 * m[i] + 0.1 * p.grad # 速度: 90% 惯性 + 10% 新梯度
v[i] = 0.95 * v[i] + 0.05 * p.grad ** 2 # 路况: 95% 历史 + 5% 新信息
m_hat = m[i] / (1 - 0.9 ** (step + 1)) # 偏差校正(早期校正偏向0的问题)
v_hat = v[i] / (1 - 0.95 ** (step + 1)) # 偏差校正
p.data -= lr * m_hat / (sqrt(v_hat) + 1e-8) # 更新参数
学习率还有线性衰减:
lr ↑
0.01 │╲
│ ╲
│ ╲
│ ╲
0 │ ╲
└──────────→ step
0 500
开始大步走,结束小步微调
第七章:推理 —— 让模型"说话"
for sample_idx in range(20):
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
token_id = BOS # 从 <BOS> 开始
for pos_id in range(block_size):
logits = gpt(token_id, pos_id, keys, values)
probs = softmax([l / temperature for l in logits]) # 温度调节
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
if token_id == BOS: # 遇到 BOS = 结束
break
print(itos[token_id], end="")
Temperature(温度)的作用
温度 = 0.1 (低温, 保守): 温度 = 1.0 (高温, 疯狂):
概率 概率
0.9 │ █ 0.3 │ █
│ █ │ █ █
0.1 │ █ ░ ░ ░ 0.1 │ █ █ █ █ █
└───────── └─────────
a b c d a b c d
几乎只选最可能的 选择更随机、更有"创意"
原理:logits / temperature。温度低 → logits 被放大 → softmax 后分布更尖锐(集中在最大值)→ 输出更确定。温度高 → 反之 → 更随机。
自回归生成过程
步骤: BOS → ? → ? → ? → ? → BOS(结束)
第1步: 输入 BOS, 模型输出概率分布, 采样得到 'J'
第2步: 输入 'J', 模型回看 [BOS, J], 采样得到 'a'
第3步: 输入 'a', 模型回看 [BOS, J, a], 采样得到 'n'
第4步: 输入 'n', 模型回看 [BOS, J, a, n], 采样得到 'e'
第5步: 输入 'e', 模型回看 [BOS, J, a, n, e], 采样得到 BOS → 结束!
生成结果: "Jane" ← 一个看起来像真名字的名字!
KV Cache 在这里体现了它的价值:每一步生成新 token 时,前面 token 的 K、V 已经算过并存在 keys 和 values 列表里了,不需要重复计算。
第八章:回顾全局——为什么这 243 行如此重要
这个代码和"真正的 GPT"有什么区别?
| 这个迷你 GPT | GPT-2 Small | GPT-4 | |
|---|---|---|---|
| 参数量 | ~5000 | 1.24 亿 | 传闻万亿级 |
| 层数 | 1 | 12 | ? |
| 嵌入维度 | 16 | 768 | ? |
| 训练数据 | 3万个名字 | 整个互联网子集 | 更大 |
| 训练时间 | 几分钟 | 数天 | 数月 |
| 用 GPU? | 不用 | 是 | 大量 |
但核心算法完全一样! 区别只是规模。这就是 Karpathy 说的:
"This file is the complete algorithm. Everything else is just efficiency."
这份代码里没有的"效率"优化
- GPU 并行计算(PyTorch/CUDA)
- 矩阵批量运算(不用一个个 Value 对象)
- Flash Attention(节省显存的注意力实现)
- 混合精度训练(fp16/bf16)
- 分布式训练(多卡/多机)
- KV Cache 优化(PagedAttention 等)
- BPE Tokenizer(更高效的分词)
这些都是工程优化。算法本身,就在这 243 行里。
最关键的一个 Gotcha
这份代码极慢——因为每个标量都是一个 Value Python 对象。一次矩阵乘法 [16×16] × [16] 就要创建几百个 Value 对象、几百次 Python 函数调用。真正的深度学习框架用 C++/CUDA 做同样的事,快上万倍。
但这恰恰是它的教育价值:把被框架隐藏的每一步都暴露出来。当你用 PyTorch 写 loss.backward() 时,底层做的事和这里的 Value.backward() 完全一样——只是用 GPU 上的张量而不是 Python 标量。
总结:一张图看完 243 行
names.txt
│
↓
┌──────────────┐
│ Tokenizer │ "Emma" → [0,5,13,13,1,0]
└──────┬───────┘
│
┌──────↓───────┐ ┌───────────────────┐
│ GPT Model │────→│ Value Autograd │
│ (前向传播) │ │ (记录计算图) │
└──────┬───────┘ └───────────────────┘
│
┌──────↓───────┐
│ Cross-Entropy│ "猜对了吗?" → loss 值
│ Loss │
└──────┬───────┘
│
┌──────↓───────┐ ┌───────────────────┐
│ Backward │────→│ 链式法则反向传播 │
│ (反向传播) │ │ (自动求梯度) │
└──────┬───────┘ └───────────────────┘
│
┌──────↓───────┐
│ Adam │ 用梯度聪明地更新参数
│ Optimizer │
└──────┬───────┘
│
重复 500 次
│
┌──────↓───────┐
│ Inference │ BOS → 'J' → 'a' → 'n' → 'e' → BOS
│ (自回归) │
└──────────────┘
这 243 行代码证明了一件事:GPT 的核心思想,简单到一个文件就能装下。 剩下的一切——GPU、分布式、Flash Attention、BPE——都只是为了让它跑得更快、吃得更多。
If you read this far — thank you.
Come tell me what you thought on X.