Karpathy 发布 microgpt：200 行纯 Python 实现 GPT 训练与推理

这是一份关于我新艺术项目 microgpt 的简要指南，它是一个 200 行纯 Python 的单文件，无任何依赖，可以训练和推理一个 GPT。这个文件包含了所需算法的全部内容：文档数据集、分词器、自动求导引擎、类似 GPT-2 的神经网络架构、Adam 优化器、训练循环和推理循环。其余的一切都只是效率问题。我无法再进一步简化它了。这个脚本是多个项目（micrograd、makemore、nanogpt 等）的结晶，也是我十年来痴迷于将 LLM 简化到最本质的成果，我觉得它很美 🥹。它甚至能完美地分成 3 列：

在哪里找到它：

• 这个 GitHub gist 包含完整源代码：microgpt.py
• 它也可以在这个网页上找到：https://karpathy.ai/microgpt.html
• 也可以作为 Google Colab notebook 获取
• 新：在我的艺术商店 karpathy.art 购买 microgpt 三联画 :)

以下是我引导感兴趣的读者逐步理解代码的指南。

数据集

大型语言模型的燃料是文本数据流，可选地分隔成一组文档。在生产级应用中，每个文档会是一个互联网网页，但对于 microgpt，我们使用一个更简单的例子：32,000 个名字，每行一个：

# Let there be an input dataset `docs`: list[str] of documents (e.g. a dataset of names)
if not os.path.exists('input.txt'):
    import urllib.request
    names_url = 'https://raw.githubusercontent.com/karpathy/makemore/refs/heads/master/names.txt'
    urllib.request.urlretrieve(names_url, 'input.txt')
docs = [l.strip() for l in open('input.txt').read().strip().split('\n') if l.strip()] # list[str] of documents
random.shuffle(docs)
print(f"num docs: {len(docs)}")

数据集看起来像这样。每个名字是一个文档：

emma
olivia
ava
isabella
sophia
charlotte
mia
amelia
harper
... (~32,000 names follow)

模型的目标是学习数据中的模式，然后生成具有相同统计模式的新文档。作为预览，到脚本结束时，我们的模型将生成（“幻觉”！）新的、听起来合理的名字。提前看一下，我们会得到：

sample  1: kamon
sample  2: ann
sample  3: karai
sample  4: jaire
sample  5: vialan
sample  6: karia
sample  7: yeran
sample  8: anna
sample  9: areli
sample 10: kaina
sample 11: konna
sample 12: keylen
sample 13: liole
sample 14: alerin
sample 15: earan
sample 16: lenne
sample 17: kana
sample 18: lara
sample 19: alela
sample 20: anton

看起来不多，但从像 ChatGPT 这样的模型的角度来看，你与它的对话只是一个有趣的“文档”。当你用提示初始化文档时，从模型的角度来看，它的响应只是一个统计文档补全。

分词器

在底层，神经网络使用数字而不是字符，所以我们需要一种方法将文本转换为整数 token id 序列，然后再转换回来。像 tiktoken（GPT-4 使用）这样的生产分词器为了提高效率而操作字符块，但最简单的分词器只是为数据集中的每个唯一字符分配一个整数：

# Let there be a Tokenizer to translate strings to discrete symbols and back
uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1
BOS = len(uchars) # token id for the special Beginning of Sequence (BOS) token
vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS
print(f"vocab size: {vocab_size}")

在上面的代码中，我们收集数据集中的所有唯一字符（它们只是所有小写字母 a-z），对它们进行排序，每个字母根据其索引获得一个 id。注意，整数值本身没有任何意义；每个 token 只是一个单独的离散符号。它们可以是不同的表情符号，而不是 0、1、2。此外，我们创建了一个额外的特殊 token 称为 BOS（序列开始），它充当分隔符：它告诉模型“一个新文档在这里开始/结束”。稍后在训练期间，每个文档两边都会用 BOS 包裹：[BOS, e, m, m, a, BOS]。模型学会 BOS 开始一个新名字，另一个 BOS 结束它。因此，我们有一个最终的词汇表，大小为 27（26 个可能的小写字符 a-z 加上 1 个 BOS token）。

自动求导

训练神经网络需要梯度：对于模型中的每个参数，我们需要知道“如果我把这个数值稍微调大一点，损失会上升还是下降，变化多少？”。计算图有许多输入（模型参数和输入 token），但最终汇聚成一个标量输出：损失（我们将在下面精确定义损失是什么）。反向传播从这单个输出开始，沿着计算图反向进行，计算损失相对于每个输入的梯度。它依赖于微积分中的链式法则。在生产环境中，PyTorch 等库会自动处理这一点。在这里，我们用一个名为 Value 的类从头实现它：

class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads')

    def __init__(self, data, children=(), local_grads=()):
        self.data = data                # scalar value of this node calculated during forward pass
        self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass
        self._children = children       # children of this node in the computation graph
        self._local_grads = local_grads # local derivative of this node w.r.t. its children

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))
    def log(self): return Value(math.log(self.data), (self,), (1/self.data,))
    def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))
    def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))
    def __neg__(self): return self * -1
    def __radd__(self, other): return self + other
    def __sub__(self, other): return self + (-other)
    def __rsub__(self, other): return other + (-self)
    def __rmul__(self, other): return self * other
    def __truediv__(self, other): return self * other**-1
    def __rtruediv__(self, other): return other * self**-1

    def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._children:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1
        for v in reversed(topo):
            for child, local_grad in zip(v._children, v._local_grads):
                child.grad += local_grad * v.grad

我意识到这是数学和算法上最密集的部分，我为此制作了一个 2.5 小时的视频：micrograd 视频。简而言之，一个 Value 包装单个标量数值（.data）并记录它是如何计算出来的。把每个运算看作一块小乐高积木：它接收一些输入，产生一个输出（前向传播），并且知道它的输出会如何随每个输入变化（局部梯度）。这就是 autograd 从每个积木块中需要的全部信息。其余的一切都只是链式法则，把这些积木块串联起来。

每当你用 Value 对象进行数学运算（加法、乘法等）时，结果是一个新的 Value，它记住自己的输入（_children）以及该运算的局部导数（_local_grads）。例如，__mul__ 记录了 \(\frac{\partial(a \cdot b)}{\partial a} = b\) 和 \(\frac{\partial(a \cdot b)}{\partial b} = a\)。完整的乐高积木集合：

backward() 方法按逆拓扑顺序遍历此图（从损失开始，到参数结束），在每一步应用链式法则。如果损失为 \(L\)，节点 \(v\) 有一个子节点 \(c\)，其局部梯度为 \(\frac{\partial v}{\partial c}\)，那么：

\[\frac{\partial L}{\partial c} \mathrel{+}= \frac{\partial v}{\partial c} \cdot \frac{\partial L}{\partial v}\]

如果你对微积分不太熟悉，这看起来有点吓人，但这实际上只是以直观的方式将两个数相乘。一种理解方式如下：“如果汽车的速度是自行车的两倍，而自行车的速度是步行者的四倍，那么汽车的速度就是步行者的 2 x 4 = 8 倍。”链式法则也是同样的道理：你沿着路径将变化率相乘。

我们首先在损失节点处设置 self.grad = 1，因为 \(\frac{\partial L}{\partial L} = 1\)：损失相对于自身的变化率显然是 1。从那里开始，链式法则只是沿着回到参数的每条路径将局部梯度相乘。

注意 +=（累加，而非赋值）。当一个值在图中被多处使用（即图出现分支）时，梯度会沿每个分支独立回流，并且必须求和。这是多变量链式法则的结果：如果 \(c\) 通过多条路径对 \(L\) 有贡献，那么总导数就是各路径贡献之和。

在 backward() 完成后，图中的每个 Value 都有一个 .grad，其中包含 \(\frac{\partial L}{\partial v}\)，它告诉我们如果微调该值，最终损失会如何变化。

这里有一个具体例子。注意 a 被使用了两次（图发生了分支），因此它的梯度是两条路径之和：

a = Value(2.0)
b = Value(3.0)
c = a * b       # c = 6.0
L = c + a       # L = 8.0
L.backward()
print(a.grad)   # 4.0 (dL/da = b + 1 = 3 + 1, via both paths)
print(b.grad)   # 2.0 (dL/db = a = 2)

这正是 PyTorch 的 .backward() 所提供的结果：

import torch
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)
c = a * b
L = c + a
L.backward()
print(a.grad)   # tensor(4.)
print(b.grad)   # tensor(2.)

这与 PyTorch 的 loss.backward() 所运行的算法相同，只是作用于标量而非张量（标量数组）——算法上完全一致，规模显著更小、更简单，但当然效率低得多。

让我们解释一下上面 .backward() 给出的结果。Autograd 计算出，如果 L = a*b + a，且 a=2 和 b=3，那么 a.grad = 4.0 告诉我们的是 a 对 L 的局部影响。如果你微调输入 a，L 会朝哪个方向变化？这里，L 关于 a 的导数是 4.0，意味着如果我们将 a 增加一个很小的量（比如 0.001），L 将增加约 4 倍（0.004）。类似地，b.grad = 2.0 意味着对 b 的同样微调会使 L 增加约 2 倍（0.002）。换句话说，这些梯度告诉我们每个单独输入对最终输出（损失）影响的方向（正或负取决于符号）和陡峭程度（幅度）。这进而使我们能够迭代地微调神经网络参数以降低损失，从而改进其预测。

参数

参数是模型的知识。它们是一大组浮点数（为自动求导而封装在 Value 中），初始时是随机的，并在训练过程中被迭代优化。每个参数的确切作用在下面定义模型架构后会更清楚，但现在我们只需要初始化它们：

n_embd = 16     # embedding dimension
n_head = 4      # number of attention heads
n_layer = 1     # number of layers
block_size = 16 # maximum sequence length
head_dim = n_embd // n_head # dimension of each head
matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}
for i in range(n_layer):
    state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)
params = [p for mat in state_dict.values() for row in mat for p in row]
print(f"num params: {len(params)}")

每个参数都初始化为从高斯分布中抽取的小随机数。state_dict 将它们组织成命名矩阵（借用 PyTorch 的术语）：嵌入表、注意力权重、MLP 权重和最终输出投影。我们还将所有参数展平为一个列表 params，以便优化器稍后可以遍历它们。在我们的微型模型中，这总共是 4,192 个参数。GPT-2 有 16 亿个，而现代 LLM 有数千亿个。

架构

模型架构是一个无状态函数：它接收一个 token、一个位置、参数以及来自先前位置的缓存键/值，并返回关于模型认为序列中下一个应该是什么 token 的 logits（分数）。我们遵循 GPT-2 并做了少量简化：用 RMSNorm 代替 LayerNorm，没有偏置，用 ReLU 代替 GeLU。首先，三个小型辅助函数：

def linear(x, w):
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

linear 是矩阵-向量乘法。它接收一个向量 x 和一个权重矩阵 w，并为 w 的每一行计算一个点积。这是神经网络的基本构建块：一种学习到的线性变换。

def softmax(logits):
    max_val = max(val.data for val in logits)
    exps = [(val - max_val).exp() for val in logits]
    total = sum(exps)
    return [e / total for e in exps]

softmax 将原始分数（logits）向量（范围可从 \(-\infty\) 到 \(+\infty\)）转换为概率分布：所有值最终都在 \([0, 1]\) 内且总和为 1。为了数值稳定性，我们先减去最大值（这在数学上不改变结果，但可防止 exp 中的溢出）。

def rmsnorm(x):
    ms = sum(xi * xi for xi in x) / len(x)
    scale = (ms + 1e-5) ** -0.5
    return [xi * scale for xi in x]

rmsnorm（均方根归一化）重新缩放向量，使其值具有单位均方根。这可以防止激活值在网络中流动时增大或缩小，从而稳定训练。它是原始 GPT-2 中使用的 LayerNorm 的更简单变体。

现在来看模型本身：

def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id] # token embedding
    pos_emb = state_dict['wpe'][pos_id] # position embedding
    x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding
    x = rmsnorm(x)

    for li in range(n_layer):
        # 1) Multi-head attention block
        x_residual = x
        x = rmsnorm(x)
        q = linear(x, state_dict[f'layer{li}.attn_wq'])
        k = linear(x, state_dict[f'layer{li}.attn_wk'])
        v = linear(x, state_dict[f'layer{li}.attn_wv'])
        keys[li].append(k)
        values[li].append(v)
        x_attn = []
        for h in range(n_head):
            hs = h * head_dim
            q_h = q[hs:hs+head_dim]
            k_h = [ki[hs:hs+head_dim] for ki in keys[li]]
            v_h = [vi[hs:hs+head_dim] for vi in values[li]]
            attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]
            attn_weights = softmax(attn_logits)
            head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]
            x_attn.extend(head_out)
        x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])
        x = [a + b for a, b in zip(x, x_residual)]
        # 2) MLP block
        x_residual = x
        x = rmsnorm(x)
        x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
        x = [xi.relu() for xi in x]
        x = linear(x, state_dict[f'layer{li}.mlp_fc2'])
        x = [a + b for a, b in zip(x, x_residual)]

    logits = linear(x, state_dict['lm_head'])
    return logits

该函数在特定时间位置（pos_id）处理一个 token（id 为 token_id），并利用由 keys 和 values 中的激活值所概括的、来自先前迭代的上下文，即 KV Cache。以下是逐步发生的过程：

嵌入。 神经网络无法直接处理像 5 这样的原始 token id，它只能处理向量（数字列表）。因此，我们为每个可能的 token 关联一个学习到的向量，并将其作为该 token 的神经特征输入。token id 和位置 id 分别从各自的嵌入表（wte 和 wpe）中查找一行。这两个向量相加，为模型提供了一种既编码 token 是什么、又编码它在哪里的表示。现代 LLM 通常跳过位置嵌入，转而引入其他基于相对位置的方案，例如 RoPE。

注意力块。 当前 token 被投影为三个向量：查询（Q）、键（K）和值（V）。直观地说，查询表示“我在寻找什么？”，键表示“我包含什么？”，值表示“如果被选中，我能提供什么？”。例如，在名字“emma”中，当模型位于第二个“m”并试图预测下一个字符时，它可能学到一个类似“最近出现了哪些元音？”的查询。较早的“e”会有一个与该查询匹配良好的键，因此获得较高的注意力权重，其值（关于它是元音的信息）流入当前位置。键和值被追加到 KV Cache 中，以便先前位置可用。每个注意力头计算其查询与所有缓存键之间的点积（按 \(\sqrt{d_{head}}\) 缩放），应用 softmax 得到注意力权重，并对缓存值取加权和。所有头的输出被拼接起来，并通过 attn_wo 进行投影。值得强调的是，注意力块是位置 t 处的 token 能够“查看”过去 0..t-1 处 token 的唯一确切位置。注意力是一种 token 通信机制。

MLP 块。 MLP 是“多层感知机”的缩写，它是一个两层前馈网络：先投影到嵌入维度的 4 倍，应用 ReLU，再投影回原维度。这是模型在每个位置上进行大部分“思考”的地方。与注意力不同，这一计算完全局限于时间 t。Transformer 将通信（注意力）与计算（MLP）交替进行。

残差连接。 注意力和 MLP 块都会将其输出加回其输入（x = [a + b for ...]）。这使得梯度可以直接流经网络，并使更深的模型可训练。

输出。 最终的隐藏状态通过 lm_head 投影到词表大小，为词表中的每个 token 生成一个 logit。在我们的例子中，这只是 27 个数字。logit 越高 = 模型认为对应 token 接下来出现的可能性越大。

你可能注意到我们在训练时使用了 KV 缓存，这并不常见。人们通常只把 KV 缓存与推理联系在一起。但从概念上讲，KV 缓存始终存在，即使在训练期间也是如此。在生产实现中，它只是隐藏在高度向量化的注意力计算内部，该计算会同时处理序列中的所有位置。由于 microgpt 一次只处理一个 token（没有批次维度，也没有并行时间步），我们显式地构建了 KV 缓存。而且与典型的推理场景中 KV 缓存保存的是分离的张量不同，这里的缓存键和值都是计算图中的活跃 Value 节点，因此我们实际上会通过它们进行反向传播。

训练循环

现在我们把所有部分串联起来。训练循环反复执行以下步骤：(1) 选取一个文档，(2) 对其 token 运行模型前向传播，(3) 计算损失，(4) 反向传播以获得梯度，(5) 更新参数。

# Let there be Adam, the blessed optimizer and its buffers
learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
m = [0.0] * len(params) # first moment buffer
v = [0.0] * len(params) # second moment buffer

# Repeat in sequence
num_steps = 1000 # number of training steps
for step in range(num_steps):

    # Take single document, tokenize it, surround it with BOS special token on both sides
    doc = docs[step % len(docs)]
    tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
    n = min(block_size, len(tokens) - 1)

    # Forward the token sequence through the model, building up the computation graph all the way to the loss.
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    losses = []
    for pos_id in range(n):
        token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
        logits = gpt(token_id, pos_id, keys, values)
        probs = softmax(logits)
        loss_t = -probs[target_id].log()
        losses.append(loss_t)
    loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.

    # Backward the loss, calculating the gradients with respect to all model parameters.
    loss.backward()

    # Adam optimizer update: update the model parameters based on the corresponding gradients.
    lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay
    for i, p in enumerate(params):
        m[i] = beta1 * m[i] + (1 - beta1) * p.grad
        v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
        m_hat = m[i] / (1 - beta1 ** (step + 1))
        v_hat = v[i] / (1 - beta2 ** (step + 1))
        p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
        p.grad = 0

    print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}")

我们逐一来看每个部分：

分词。 每个训练步骤选取一个文档，并在其两侧加上 BOS：名称“emma”变成 [BOS, e, m, m, a, BOS]。模型的任务是在给定前面 token 的情况下预测每个下一个 token。

前向传播与损失。 我们将 token 逐个送入模型，同时逐步构建 KV 缓存。在每个位置，模型输出 27 个 logits，我们通过 softmax 将其转换为概率。每个位置的损失是正确下一个 token 的负对数概率：\(-\log p(\text{target})\)。这被称为交叉熵损失。直观地说，损失衡量的是预测错误的程度：模型对实际接下来出现的内容有多惊讶。如果模型给正确 token 分配的概率为 1.0，它就完全不惊讶，损失为 0。如果分配的概率接近 0，模型就非常惊讶，损失趋于 \(+\infty\)。我们对文档中各个位置的损失取平均，得到单个标量损失。

反向传播。 一次 loss.backward() 调用会在整个计算图上运行反向传播，从损失一直回溯经过 softmax、模型，进入每个参数。之后，每个参数的 .grad 告诉我们如何修改它来降低损失。

Adam 优化器。 我们本可以只做 p.data -= lr * p.grad（梯度下降），但 Adam 更聪明。它为每个参数维护两个运行平均值：m 跟踪近期梯度的均值（动量，就像滚动的球），v 跟踪近期梯度平方的均值（为每个参数自适应调整学习率）。m_hat 和 v_hat 是偏差修正，用于处理 m 和 v 初始化为零且需要预热的事实。学习率在训练过程中线性衰减。更新后，我们重置 .grad = 0 以供下一步使用。

在 1,000 步中，损失从约 3.3（在 27 个 token 中随机猜测：\(-\log(1/27) \approx 3.3\)）下降到约 2.37。越低越好，最低可能值为 0（完美预测），因此仍有改进空间，但模型显然正在学习名字的统计模式。

推理

训练完成后，我们可以从模型中采样新名字。参数被冻结，我们只需在循环中运行前向传播，将每个生成的 token 作为下一个输入反馈回去：

temperature = 0.5 # in (0, 1], control the "creativity" of generated text, low to high
print("\n--- inference (new, hallucinated names) ---")
for sample_idx in range(20):
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    token_id = BOS
    sample = []
    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:
            break
        sample.append(uchars[token_id])
    print(f"sample {sample_idx+1:2d}: {''.join(sample)}")

我们以 BOS token 开始每个样本，它告诉模型“开始一个新名字”。模型产生 27 个 logits，我们将它们转换为概率，然后根据这些概率随机采样一个 token。该 token 会作为下一个输入被反馈回去，我们重复这一过程，直到模型再次产生 BOS（表示“我完成了”）或达到最大序列长度。

temperature 参数控制随机性。在 softmax 之前，我们将 logits 除以温度。温度为 1.0 时直接从模型学习到的分布中采样。较低的温度（如这里的 0.5）会使分布更尖锐，使模型更保守，更可能选择其首选。温度趋近于 0 时总是选择单个最可能的 token（贪婪解码）。较高的温度会拉平分布，产生更多样但可能不太连贯的输出。

运行它

你只需要 Python（无需 pip install，无依赖）：

python train.py

该脚本在我的 macbook 上运行约需 1 分钟。你会看到每一步打印的损失：

train.py
num docs: 32033
vocab size: 27
num params: 4192
step    1 / 1000 | loss 3.3660
step    2 / 1000 | loss 3.4243
step    3 / 1000 | loss 3.1778
step    4 / 1000 | loss 3.0664
step    5 / 1000 | loss 3.2209
step    6 / 1000 | loss 2.9452
step    7 / 1000 | loss 3.2894
step    8 / 1000 | loss 3.3245
step    9 / 1000 | loss 2.8990
step   10 / 1000 | loss 3.2229
step   11 / 1000 | loss 2.7964
step   12 / 1000 | loss 2.9345
step   13 / 1000 | loss 3.0544
...

观察它从约 3.3（随机）下降到约 2.37。这个数字越低，说明网络对序列中下一个 token 的预测越好。训练结束时，训练 token 序列的统计模式知识被提炼到模型参数中。固定这些参数后，我们现在可以生成新的、幻觉的名字。你会（再次）看到：

sample  1: kamon
sample  2: ann
sample  3: karai
sample  4: jaire
sample  5: vialan
sample  6: karia
sample  7: yeran
sample  8: anna
sample  9: areli
sample 10: kaina
sample 11: konna
sample 12: keylen
sample 13: liole
sample 14: alerin
sample 15: earan
sample 16: lenne
sample 17: kana
sample 18: lara
sample 19: alela
sample 20: anton

作为在你自己电脑上运行脚本的替代方案，你可以尝试直接在这个 Google Colab notebook 上运行，并向 Gemini 提问。试着玩一下这个脚本！你可以尝试不同的数据集。或者你可以训练更长时间（增加 num_steps）或增加模型规模，以获得越来越好的结果。

进阶

要像剥洋葱一样逐层查看代码的构建过程，建议的进阶顺序大致如下：

我创建了一个名为 build_microgpt.py 的 Gist，在 Revisions 中你可以看到所有这些版本以及每一步之间的差异。我认为这可能是逐步浏览代码库的一种有用方式，每次添加一个组件。

真实情况

microgpt 包含了训练和运行 GPT 的完整算法精髓。但在这与像 ChatGPT 这样的生产级 LLM 之间，有一长串变化。它们都不会改变核心算法和整体布局，但正是它们使其在规模上真正可行。按相同章节顺序逐一说明：

数据。 生产模型不是用 32K 个短名字训练，而是用数万亿 token 的互联网文本训练：网页、书籍、代码等。数据经过去重、质量过滤，并在各领域间仔细混合。

分词器。 生产模型不是使用单个字符，而是使用像 BPE（Byte Pair Encoding）这样的子词分词器，它学习将频繁共现的字符序列合并为单个 token。像“the”这样的常见词变成单个 token，罕见词被拆分成片段。这给出了约 100K token 的词表，并且效率高得多，因为模型在每个位置看到更多内容。

自动微分。 microgpt 在纯 Python 中对标量 Value 对象进行运算。生产系统使用张量（大型多维数字数组），并在每秒执行数十亿次浮点运算的 GPU/TPU 上运行。PyTorch 等库负责处理张量上的自动微分，而 FlashAttention 等 CUDA 内核则融合多个操作以提升速度。数学原理完全相同，只是对应着大量并行处理的标量。

架构。 microgpt 有 4,192 个参数。GPT-4 级别的模型有数千亿个参数。总体而言，它是一个看起来非常相似的 Transformer 神经网络，只是宽得多（嵌入维度为 10,000+）且深得多（100+ 层）。现代 LLM 还加入了更多类型的积木块并调整了它们的顺序：例如用 RoPE（旋转位置嵌入）替代可学习的位置嵌入，用 GQA（分组查询注意力）减小 KV 缓存大小，用门控线性激活替代 ReLU，以及混合专家（MoE）层等。但 Attention（通信）与 MLP（计算）交错在残差流上的核心结构被很好地保留了下来。

训练。 生产训练不是每步处理一个文档，而是使用大批量（每步数百万个 token）、梯度累积、混合精度（float16/bfloat16）以及精细的超参数调优。训练一个前沿模型需要数千块 GPU 运行数月。

优化。 microgpt 使用 Adam 加上简单的线性学习率衰减，仅此而已。在大规模场景下，优化本身成为一门独立的学问。模型以降低的精度（bfloat16 甚至 fp8）在大型 GPU 集群上训练以提升效率，这带来了自身的数值挑战。优化器设置（学习率、权重衰减、beta 参数、预热调度、衰减调度）必须精确调优，而正确的取值取决于模型规模、批量大小和数据集构成。缩放定律（例如 Chinchilla）指导如何在模型规模和训练 token 数量之间分配固定的计算预算。在大规模场景下弄错其中任何细节都可能浪费数百万美元的计算资源，因此团队会先进行大量小规模实验来预测正确的设置，然后再投入完整的训练运行。

后训练。 训练产出的基础模型（称为“预训练”模型）是一个文档补全器，而不是聊天机器人。将其变成 ChatGPT 需要两个阶段。第一阶段是 SFT（监督微调）：你只需把文档换成精心策划的对话并继续训练。从算法上讲，没有任何变化。第二阶段是 RL（强化学习）：模型生成回复，这些回复被打分（由人类、另一个“评判”模型或算法打分），模型从这些反馈中学习。从根本上说，模型仍然在文档上训练，只不过这些文档现在由来自模型自身的 token 组成。

推理。 向数百万用户提供模型服务需要其自身的工程栈：将请求批处理、KV 缓存管理与分页（vLLM 等）、用于提速的投机解码、量化（以 int8/int4 而非 float16 运行）以减少内存占用，以及将模型分布到多块 GPU 上。从根本上说，我们仍然是在预测序列中的下一个 token，只是投入了大量工程来让它更快。

这些都是重要的工程和研究贡献，但如果你理解了 microgpt，你就理解了算法的本质。

常见问题

模型真的“理解”了什么吗？ 这是个哲学问题，但从机制上说：没有任何魔法发生。模型是一个大型数学函数，将输入 token 映射为下一个 token 上的概率分布。训练期间，参数被调整以使正确的下一个 token 概率更高。这是否构成“理解”由你判断，但机制完全包含在上面的 200 行代码中。

它为什么有效？ 模型有数千个可调参数，优化器每一步将它们微调一点点以降低损失。经过许多步后，参数稳定到能捕捉数据统计规律的值。对于名字，这意味着：名字常以辅音开头，“qu”倾向于一起出现，名字很少连续出现三个辅音，等等。模型不学习显式规则，它学习的是一个恰好反映了这些规则的概率分布。

这与 ChatGPT 有什么关系？ ChatGPT 就是同一个核心循环（预测下一个 token、采样、重复）被极大地扩展，并经过后训练使其能对话。当你与它聊天时，系统提示、你的消息和它的回复都只是序列中的 token。模型在逐个 token 地补全文档，就像 microgpt 补全一个名字一样。

“幻觉”是怎么回事？ 模型通过从概率分布中采样来生成 token。它没有真理的概念，只知道在给定训练数据下哪些序列在统计上看似合理。microgpt“幻觉”出一个像“karia”这样的名字，与 ChatGPT 自信地陈述一个虚假事实是同一现象。两者都是听起来合理但恰好不真实的补全。

它为什么这么慢？ microgpt 在纯 Python 中一次处理一个标量。单次训练步骤需要数秒。同样的数学在 GPU 上并行处理数百万个标量，运行速度快几个数量级。

我能让它生成更好的名字吗？ 可以。训练更久（增加 num_steps），让模型更大（n_embd、n_layer、n_head），或使用更大的数据集。这些正是在大规模下同样重要的调节旋钮。

如果我更改数据集会怎样？ 模型会学习数据中的任何模式。换入一个城市名、宝可梦名字、英语单词或短诗的文件，模型就会学会生成那些内容。其余代码无需更改。