microgpt: 一个最小化的 GPT 实现

这是我新艺术项目 microgpt 的简要指南，一个包含 200 行纯 Python 代码的单个文件，没有依赖项，可以训练和推断 GPT。这个文件包含了所需的全部算法内容：文档数据集、分词器、自动求导引擎、GPT-2 类神经网络架构、Adam 优化器、训练循环和推断循环。其他一切都是效率问题。 我无法再简化这个脚本了。这脚本是多个项目（micrograd、makemore、nanogpt 等）的结晶，经过十年来简化 LLM 的执着努力，我认为它很美丽 🥹。 它甚至可以完美地分成 3 列：

找到它

• 这个 GitHub gist 包含了完整的源代码：microgpt.py
• 它也可以在这个网页上找到：https://karpathy.ai/microgpt.html
• 还可以作为一个 Google Colab 笔记本

数据集

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

# 让我们有一个输入数据集 `docs`：文档列表（例如一个名字数据集）
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()] # 文档列表
random.shuffle(docs)
print(f"num docs: {len(docs)}")

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

emma olivia ava isabella sophia charlotte mia amelia harper ... (~32, 000 个名字后面跟着)

模型的目标是学习数据中的模式，然后生成类似的新文档，分享统计模式。作为预览，脚本结束时我们的模型将生成（“幻想”！）新的、听起来合理的名字。跳过一些内容，我们将得到：

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

看起来不太像什么，但从 ChatGPT 等模型的角度来看，和它的对话是同样“文档”。当您初始化文档时，模型的响应就是从其角度看的统计文档完成。

分词器

神经网络在背后工作的是数字，而不是字符，所以我们需要一种方法将文本转换为整数令牌 ID 的序列和反过来。生产级分词器，如 tiktoken（GPT-4 使用的分词器），在效率上操作字符块，但最简单的分词器只是为数据集中的每个唯一字符分配一个整数：

# 让我们有一个分词器，将字符串转换为离散符号和反过来
uchars = sorted(set(''.join(docs))) # 数据集中的唯一字符变成令牌 ID 0..n-1
BOS = len(uchars) # 令牌 ID 为特殊的开始序列（BOS）令牌
vocab_size = len(uchars) + 1 # 总令牌数，+1 是 BOS 令牌
print(f"vocab size: {vocab_size}")

在上面的代码中，我们收集数据集中的所有唯一字符（就是所有小写字母 a-z），排序它们，并为每个字母分配一个 ID。请注意，整数本身没有任何意义；每个令牌都是一个独立的离散符号。取而代之的是 0、1、2 等，可能更好地使用不同的 emoji。另外，我们创建一个特殊令牌 BOS（开始序列），它作为分隔符：它告诉模型“一个新文档开始/结束”。在训练期间，后面每个文档都被包围在 BOS 两侧： [BOS, e, m, m, a, BOS]。模型学习 BOS 初始化一个新名字，另一个 BOS 结束它。因此，我们有一个最终词汇表 27 个元素（26 个可能的小写字母 a-z 和 +1 个 BOS 令牌）。

自动求导

训练神经网络需要梯度：对于模型中的每个参数，我们需要知道“如果我轻微地推动这个数字，损失会上升还是下降，幅度有多大？””. 计算图有很多输入（模型参数和输入令牌），但都汇集到一个单一的标量输出：损失（我们稍后将定义损失）。反向传播从这个单一输出开始，沿着图向后工作，计算损失对每个输入的梯度。它依赖于微积分中的链式法则。在生产中，库如 PyTorch 会自动处理这个问题。在这里，我们从头开始实现它在一个名为 Value 的单个类中：

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

    def __init__(self, data, children=(), local_grads=()):
        self.data = data # 计算过程中的此节点的标量值
        self.grad = 0 # 损失对此节点的导数，计算在反向传播过程中
        self._children = children # 此节点在计算图中的子节点
        self._local_grads = local_grads # 此节点对其子节点的局部导数

    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{∂(a⋅b)}{∂a}=b$ 和 $\frac{∂(a⋅b)}{∂b}=a$ . 全部积木块：

操作	前向传播	局部导数
a + b	$a+b$	$\frac{∂}{∂a}=1,\frac{∂}{∂b}=1$
a * b	$a⋅b$	$\frac{∂}{∂a}=b,\frac{∂}{∂b}=a$
a ** n	$a^{n}$	$\frac{∂}{∂a}=n⋅a^{n-1}$
log(a)	$ln(a)$	$\frac{∂}{∂a}=\frac{1}{a}$
exp(a)	$e^{a}$	$\frac{∂}{∂a}=e^{a}$
relu(a)	$max(0,a)$	$\frac{∂}{∂a}=1_{a>0}$

backward() 方法沿着图的反向拓扑顺序（从损失开始，结束于参数）向后工作，应用链式法则。在每个步骤中，如果损失为 $L$ 且一个节点 $v$ 有一个子节点 $c$ 有局部导数 $\frac{∂v}{∂c}$ , 则：

$$\frac{∂L}{∂c}+=\frac{∂v}{∂c}⋅\frac{∂L}{∂v}$$

这看起来很吓人，如果您不熟悉微积分，但这只是在直觉上进行乘法。一个方法是这样看：如果一辆汽车比自行车快两倍，而自行车比行人快四倍，那么汽车比行人快 2 x 4 = 8 倍。链式法则是相同的想法：您将沿路径的导数相乘。

我们从损失节点开始，将 self.grad = 1 设置为 1，因为 $\frac{∂L}{∂L}=1$ : 损失对自身的导数是显而易见的 1。从那里，链式法则就只是沿每条路径将局部导数相乘。

我们从损失节点开始，设置 self.grad = 1，因为 $\frac{∂L}{∂L}=1$：损失的变化率与自身的变化率是微不足道的。从那里，链式规则只是沿着每条路径反向传播局部梯度。

注意 +=（累加，而不是赋值）。当一个值在图中被使用多次（即图分叉）时，梯度沿着每个分叉独立流动并必须相加。这是多变量链式规则的结果：如果 $c$ 通过多条路径贡献到 $L$，则总导数是每条路径贡献的总和。

当 backward() 完成时，每个 Value 在图中都有一个 .grad，包含 $\frac{∂L}{∂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$，损失会如何变化？这里，损失对 $a$ 的导数是 4.0，这意味着如果我们增加 $a$ 的一个小量（例如 0.001），损失将增加约 4 倍的量（0.004）。同样，$b.grad = 2.0$ 表示同样的轻微推动 $b$ 将增加损失约 2 倍的量（0.002）。换句话说，这些梯度告诉我们每个个体输入对最终输出（损失）的方向（正或负，取决于符号）和陡度（幅度）的影响。这样我们就可以逐步推动神经网络的参数来降低损失，从而改进其预测。

参数

参数是模型的知识。它们是一大堆浮点数（包装在 Value 中用于 Autograd），从随机开始并在训练过程中逐步优化。每个参数的具体作用将在下面定义模型架构时变得更加清晰，但现在我们只需要初始化它们：

n_embd = 16 # embedding 维度
n_head = 4 # 注意力头数
n_layer = 1 # 层数
block_size = 16 # 序列最大长度
head_dim = n_embd // n_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 有 1.6 亿个参数，现代 LLM 有数百亿个参数。

架构

模型架构是一个无状态函数：它接受一个令牌、一个位置、参数和缓存的键/值对（来自前面的位置），并返回对模型认为下一个令牌应该出现的序列中的令牌的分数。我们遵循 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），可以范围从 $-∞$ 到 $+∞$，转换为概率分布：所有值都在 $[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 (根均方归一化) 将一个向量的值缩放到其根均方为 1。这保持激活值在网络中不增长或缩小，从而稳定训练。它是原始 GPT-2 中使用的 LayerNorm 的一个更简单的变体。

现在是模型本身：

def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id] # 令牌嵌入
    pos_emb = state_dict['wpe'][pos_id] # 位置嵌入
    x = [t + p for t, p in zip(tok_emb, pos_emb)] # 联合令牌和位置嵌入
    x = rmsnorm(x)

    for li in range(n_layer):
        # 1) 注意力块
        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 块
        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

函数处理一个令牌（id token_id)在特定时间位置（pos_id)，一些上下文信息从前面的迭代中概括为 keys 和 values，称为 KV 缓存。以下是每一步的详细信息：

嵌入。神经网络无法直接处理一个原始令牌 id，如 5。它只能处理向量（数字列表）。因此，我们为每个可能的令牌学习一个向量，并将其作为其神经签名。令牌 id 和位置 id 分别从它们的嵌入表中查找一行（wte 和 wpe)。这两个向量相加，给模型提供一个表示令牌的表示，既包含令牌的内容，又包含其在序列中的位置。现代 LLM 通常跳过位置嵌入并引入其他相对定位方案，例如 RoPE。

注意力块。当前令牌被投影为三个向量：一个查询（Q）、一个键（K）和一个值（V）。直观地说，查询说“我在寻找什么？”，键说“我包含什么？”，值说“如果我被选中，我会提供什么？”例如，在名字“emma”中，当模型在第二个“m”时，试图预测接下来出现的令牌时，它可能会学习一个查询“最近出现的哪些元音？”早先的“e”将有一个匹配这个查询的键，因此它将获得一个高注意力权重，其值（元音信息）将流入当前位置。键和值被追加到 KV 缓存中，以便之前的位置可用。每个注意力头计算其查询与所有缓存键的点积（乘以 $\sqrt{d_{head}}$），应用 softmax 得到注意力权重，并取缓存值的加权和。所有头的输出被连接并通过 attn_wo 投影。值得注意的是，注意力块是令牌在位置 t 时可以“看”之前的令牌 0..t-1 的唯一地方。注意力是一种令牌通信机制。

注意力块。 当前的令牌被投射到三个向量：一个查询（Q）、一个密钥（K）和一个值（V）。直观地说，查询说“我在寻找什么？”，密钥说“我包含什么？”，值说“如果选择我，我会提供什么？”。例如，在名称“emma”中，当模型位于第二个“m”时，试图预测接下来会发生什么，它可能会学习一个查询“最近出现的元音是什么？”早期的“e”将有一个匹配这个查询的密钥，因此它将获得一个高注意力权重，其值（元音的信息）将流入当前位置。密钥和值将被追加到KV缓存中，以便以前的位置可用。每个注意力头计算其查询与所有缓存密钥（乘以$\sqrt{d_{head}}$）的点积，应用softmax以获得注意力权重，并取缓存值的加权和。所有头部的输出将被连接并通过attn_wo投影。值得强调的是，注意力块是唯一一个令牌在位置t时可以“看”过去的位置0..t-1的位置。注意力是一个令牌通信机制。

MLP块。 MLP是“多层感知器”的缩写，它是一个两层的前馈网络：将嵌入维度投影到4倍，应用ReLU，投影回去。这是模型在每个位置做最多的“思考”。与注意力不同，这个计算是完全局部的到时间t。Transformer在通信（注意力）和计算（MLP）之间交替。

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

输出。 最终的隐藏状态将被投影到词汇大小（lm_head），产生每个词汇中的一个对数概率。我们的例子中，只有27个数字。更高的对数概率意味着模型认为对应的令牌更有可能是下一个。

你可能会注意到，我们在训练期间使用了KV缓存，这是非典型的。人们通常将KV缓存与推理相关联。但是，KV缓存在概念上始终存在，即使在训练期间。生产实现中，它通常被隐藏在高度向量化的注意力计算中，该计算处理整个序列中的所有位置。由于微GPT处理一个令牌一个令牌（没有批次维度，没有并行时间步），我们显式地构建了KV缓存。与典型的推理设置不同，在这里缓存的密钥和值是活跃的Value节点，因此我们实际上可以通过它们进行反向传播。

训练循环

现在我们将所有内容连接起来。训练循环反复：（1）选择一个文档，（2）将模型前向传播到其令牌，（3）计算损失，（4）反向传播以获取梯度，（5）更新参数。

# 让我们有一个Adam，圣光的优化器及其缓冲区
学习率，beta1，beta2，eps_adam = 0.01，0.85，0.99，1e-8
m = [0.0] * len(params) # 第一时间缓冲区
v = [0.0] * len(params) # 第二时间缓冲区

# 重复序列
num_steps = 1000 # 训练步骤的数量
for step in range(num_steps):

    # 取一个文档， tokenize 它，包围它的 BOS 特殊令牌
    doc = docs[step % len(docs)]
    tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
    n = min(block_size, len(tokens) - 1)

    # 将令牌序列前向传播到模型，构建 KV 缓存
    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) # 文档序列中的平均损失。愿你的损失低。

    # 反向传播损失，计算每个参数的梯度
    loss.backward()

    # Adam 优化器更新：根据每个参数的梯度更新参数
    lr_t = learning_rate * (1 - step / num_steps) # 线性学习率衰减
    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 + 1:4d} / {num_steps:4d} | 损失 {loss.data:.4f}")

让我们逐步走过每个部分：

令牌化。 每个训练步骤都选择一个文档并将其包围在 BOS 特殊令牌上：名称“emma”变成 [BOS, e, m, m, a, BOS]。模型的任务是预测每个下一个令牌，给定之前的令牌。

前向传播和损失。 我们将令牌序列前向传播到模型，逐步构建 KV 缓存。在每个位置，模型输出 27 个对数概率，我们通过 softmax 转换为概率。每个位置的损失是正确下一个令牌的负对数概率：$-log(p(\text{target}))$ 。这称为交叉熵损失。直观地说，损失衡量预测的误差：模型对实际出现的令牌有多惊讶。如果模型将概率分配为 1.0，则它对实际出现的令牌没有惊讶，损失为 0。如果它将概率分配为接近 0，则模型对实际出现的令牌非常惊讶，损失趋近于 $+∞$ 。我们将每个位置的损失平均起来，得到一个标量损失。

反向传播。 一个 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 个令牌：$-log(1/27)\approx 3.3$）下降到约 2.37。更低的损失越好，理想情况下是 0（完美预测），所以仍有改进的空间，但模型已经学习了名称的统计模式。

推理

一旦训练完成，我们就可以从模型中采样新的名称。参数被冻结，我们只运行前向传播的循环，将每个生成的令牌作为下一个输入：

温度 = 0.5 # 在 (0, 1] 中控制生成文本的创造性，低到高
print("\n --- 推理（新，幻想的名称）---")
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_idx + 1:2d} : {''.join(sample)}")

我们从 BOS 令牌开始，告诉模型“开始一个新名称”。模型输出 27 个对数概率，我们将它们转换为概率，然后随机采样一个令牌，根据这些概率。该令牌将被作为下一个输入，重复直到模型产生 BOS 令牌（意味着“我完成了”）或达到最大序列长度。

温度 参数控制随机性。前 softmax，我们将对数概率除以温度。温度为 1.0 时，采样直接从模型的学习分布中。较低的温度（如 0.5 在这里）会使分布更尖锐，使模型更保守，更有可能选择其顶部选择。温度接近 0 将始终选择最可能的令牌（贪婪解码）。较高的温度会使分布更平坦，产生更具多样性但可能不那么一致的输出。

运行它

你需要的就是 Python（没有 pip 安装，没有依赖项）。

脚本在我的 MacBook 上需要约 1 分钟才能运行。你将看到每个步骤的损失打印：

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

观看它从约 3.3（随机）下降到约 2.37。损失越低，网络对序列中下一个令牌的预测越好。训练完成后，模型参数中包含了训练令牌序列的统计模式。将这些参数固定，我们现在可以生成新的，幻想的名称。你将看到（再次）：

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

作为替代方案，您可以尝试在谷歌协作笔记本上直接运行脚本并向吉米提问。尝试玩这个脚本！您可以尝试使用不同的数据集。或者，您可以训练更长时间（增加 num_steps）或增加模型大小以获得更好的结果。

进展

要看到代码逐步构建起来，像洋葱一样，建议的进展顺序如下：

文件	添加什么
train0.py	双字计数表 — 无神经网络，无梯度
train1.py	MLP + 手动梯度（数字和分析）+ SGD
train2.py	自动梯度（Value类） — 替换手动梯度
train3.py	位置嵌入 + 单头注意力 + rmsnorm + 残差
train4.py	多头注意力 + 层循环 — 完整的 GPT 架构
train5.py	Adam 优化器 — 这是 train.py

我创建了一个名为 build_microgpt.py 的 gist，其中在 Revisions 中您可以看到所有这些版本和每个步骤之间的差异。 我认为这可能是有帮助的方式来浏览代码库，逐步添加一个组件。

真实内容

microgpt包含训练和运行一个 GPT 的完整算法本质。 但是，相对于生产 LLM，如 ChatGPT，存在一个长列表的变化。 他们没有改变核心算法和整体布局，但他们使它在规模上实际工作。 通过同样的部分顺序：

数据。 相反，生产模型在训练时使用数十亿个令牌的互联网文本：网页、书籍、代码等。数据经过去重、过滤和跨域混合。

Tokenizer。 相反，生产模型使用子词令牌器，如 BPE（字节对编码），它学习将频繁共存的字符序列合并为单个令牌。常见词如“the”变为单个令牌，罕见词被分解为片段。 这给出了约 100,000 个令牌的词汇表，并且由于模型在每个位置看到更多内容而变得更加高效。

Autograd。 microgpt在纯 Python 中操作标量 Value 对象。生产系统使用张量（大型多维数组）并在 GPU/TPU 上运行，执行每秒数十亿个浮点运算。库如 PyTorch 处理张量上的 autograd，CUDA 内核如 FlashAttention 将多个操作融合以提高速度。数学与微小的标量处理相同，只是对多个标量进行并行处理。

架构。 microgpt 有 4,192 个参数。 GPT-4 类模型有数十亿个参数。总体而言，它是一个非常相似的看起来的 Transformer 神经网络，只是更宽（嵌入维度为 10,000+）和更深（100+ 层）。 现代 LLM 还包含一些额外的 Lego 块类型和它们的顺序：例如，RoPE（旋转位置嵌入）取代学习的位置嵌入，GQA（分组查询注意力）减少 KV 缓存大小，门控线性激活取代 ReLU，混合专家（MoE）层等。 但是，注意力（通信）和 MLP（计算）在残差流中交替的核心结构得以保留。

训练。 相反，生产训练使用大批次（每步数百万令牌），梯度累积，混合精度（float16/bfloat16），以及仔细调参。训练前沿模型需要数千个 GPU 运行几个月。

优化。 microgpt 使用 Adam 优化器，简单的线性学习率衰减，基本上就是这样。 在规模上，优化成为自己的学科。 模型在减少精度（bfloat16 或 fp8）和大 GPU 集群上训练，以提高效率，这引入了自己的数值挑战。 优化器设置（学习率、权重衰减、beta 参数、预热调度、衰减调度）必须精确调整，正确的值取决于模型大小、批次大小和数据集组成。 分辨率定律（例如 Chinchilla）指导如何在固定计算预算下分配模型大小和训练令牌数量。 在规模上错误任何细节都可能浪费数百万美元的计算，因此团队在进行完整训练运行之前会进行广泛的小规模实验来预测正确的设置。

后训练。 基于训练的模型（称为“预训练”模型）是一个完成文档的文档，而不是一个聊天机器人。 将其转换为 ChatGPT 需要两个阶段。 第一阶段是 SFT（监督微调）：您只需将文档替换为经过精心策划的对话并继续训练。 算法上没有变化。 第二阶段是 RL（强化学习）：模型生成响应，得到评分（由人类、另一个“评判”模型或算法评分），并从反馈中学习。 基本上，模型仍在训练文档，但这些文档现在由模型本身生成的令牌组成。

推理。 将模型服务于数百万用户需要自己的工程栈：批量请求，KV 缓存管理和分页（vLLM 等），预测性解码以提高速度，量化（在 int8/int4 中运行而不是 float16）以减少内存，并将模型分布在多个 GPU 上。 基本上，我们仍在预测序列中的下一个令牌，但花费了很多工程精力来提高速度。

所有这些都是重要的工程和研究贡献，但如果您了解 microgpt，则了解了算法本质。

常见问题

模型是否“理解”任何东西？ 这是一个哲学问题，但机械上：没有魔法发生。 模型是一个大数学函数，它将输入令牌映射到下一个令牌的概率分布。 在训练期间，参数被调整以使正确的下一个令牌更可能。 是否这构成了“理解”是由您决定的，但机制完全包含在上面的 200 行代码中。

为什么它有效？ 模型有数千个可调节的参数，优化器每步调整它们一个小小的步骤以使损失下降。 在多步之后，参数定居在捕捉数据统计模式的值上。 对于名字，这意味着：名字通常以辅音开始，“qu”倾向于一起出现，名字很少有三个连续辅音等。 模型没有学习明确的规则，它学习了一个概率分布，它恰好反映了它们。

这与 ChatGPT 有什么关系？ ChatGPT 是这个核心循环（预测下一个令牌，采样，重复）在规模上大大扩展，经过后训练以使其成为对话。 当您与它聊天时，系统提示、您的消息和其回复都是序列中的令牌。 模型正在完成文档，同样地，microgpt 正在完成名字。

什么是“幻觉”？ 模型通过从概率分布中采样生成令牌。 它没有概念的真理，它只知道训练数据中序列的统计可信度。 microgpt “幻觉”一个名字“karia”与 ChatGPT 坚定地声称一个错误事实相同。 都是听起来很可信的完成，但并非真实。

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

您可以让它生成更好的名字吗？ 是的。 训练更长时间（增加 num_steps），使模型更大（n_embd、n_layer、n_head），或使用更大的数据集。 这些是规模上相同的调节器。

如果我改变数据集？ 模型将学习数据中的模式。 将一个城市名称文件、Pokemon 名称文件、英语词汇文件或短诗文件代入其中，模型将学习生成这些内容。 代码的其余部分不需要改变。