从零实现 GPT：Karpathy nanoGPT 学习笔记

最近学习了 Karpathy 的《Let's build GPT: from scratch, in code, spelled out》。跟着教程实现了一个简易的字符级、基于 decoder-only transformer 的语言模型，nanoGPT。

起点

以前知道 Transformer 和 LLM，但是对具体概念还是不清楚，于是自己找教程学习。

这篇笔记不打算再次完整介绍 Transformer 架构，而是沿着数据流，看看一段 token 序列如何经过 GPT，最终产生下一 token 的概率分布。

从字符到 GPT

一条输入序列首先会被 tokenizer 转成一列数字，称为 token ID 序列，然后经过 token_embedding_table 查表后得到 token embedding，加上位置编码，进入 Transformer block 进行处理，经过 lm_head，最终变成 next token。

从文本到 token

以 nanoGPT 的教程为例。在一个莎士比亚数据集上面训练了一个约 5M 参数的语言模型，数据集以 input.txt 存储，原始文本类似于这样：

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

由于这里我们训练的是字符级语言模型，所以：

统计所有出现过的字符，词表大小设置为当前字符种类的数量（65）。

创建 stoi 和 itos 映射和编解码，把字符映射到 token ID，同时准备把模型的最终结果映射回字符。

最终得到了一个一维 tensor，比如：

"hello"
↓
[3,1,5,5,8]

然后一维长序列里随机截取很多窗口，堆成 [B,T] 的二维 tensor。

例如：

"hello"
 h e l l o
 0 1 2 3 4

T = 5，由于是并行处理，这里教程设置的 B 是64，可以想象同时有64个序列并行进入网络。

由于预测目标是“预测下一个字符”，例如原序列：

hello!

则可以形成：

x = h e l l o
y = e l l o !

这样一个 [B,T] batch 会同时提供 B*T 个 next-token prediction 目标。

token 在网络内的流转

接上，我们得到了输入 tensor [B,T]，[B,T] 的 token ID tensor 经过 token_embedding_table 查表后得到 [B,T,C] 的 token embedding。这里模型设置的 C = 384，即隐藏状态向量的维度为384。

然后还需要加上位置编码，为每个 token 注入其所在序列位置的信息。这里 nanoGPT 采用的是可学习的 position embedding，位置 ID 0...T-1 经过 nn.Embedding(block_size, n_embd) 得到 [T,C]，再与 token embedding 相加。

这样我们得到了输入向量 x，[B,T,C]，这个向量会被送入 Transformer blocks 进行注意力计算。

处理后的结果经过 lm_head 得到 logits，即 [B,T,C] → [B,T,V]，其中 V = vocab_size = 65。这里的 logits 是“每个 token 位置对整个词表的原始分数”。

在推理中，会使用 softmax 得到下一个 token 的概率分布，然后使用 torch.multinomial 从分布中采样一个 token ID，然后将生成的这个新 token 拼到已生成的 token 后面，重复直到生成指定数量的新 token。

Attention Block 内部

输入向量 x [B,T,C] 被送入 Transformer blocks 进行注意力计算。对于每个 Block 内，有多头注意力 MHA、FFN 前馈网络、LayerNorm 归一化层和残差连接等组件。

单头注意力

具体注意力计算的实现：

代码定义了3个线性层（nn.Linear）K、Q、V，即可学习的权重矩阵（大小 n_embd × head_size），然后注册了一个固定的下三角矩阵 causal mask。

对于每次前向传播：

输入 Tensor x，形状为 [B,T,C]，这里 B 是批次并行化处理，实际上计算就是 T × C @ C × head_size（本教程 head_size 为 96），单个 head，得到 k、q、v。然后计算 q @ k 的转置，这里是 q @ k.transpose(-2, -1)，形状 [B,T,T]。

按照标准 Scaled Dot-Product Attention 公式，这里的缩放维度应为单个 Head 的 head_size（即 dk）。在本模型中为 96。本实现仍保留了 Karpathy 教程代码中的 C**-0.5。

得到的注意力分数会经过 causal mask，causal mask 将未来位置置为 -inf，使当前位置只能关注自己和之前的 token，避免训练时通过未来 token 泄漏预测目标，并保持与自回归生成过程一致。

然后在 dim=-1 上做 softmax，使每一行成为“当前 token 对允许看到的各 token 的注意力权重分布”，得到 attention_weights。

这里训练的时候，还会对部分 attention 权重随机 Dropout，作为一种正则化手段，减少模型对少数特征或连接的过度依赖，从而改善泛化能力。

最终的单头输出，attention_weights @ v，形状 [B,T,head_size]。

多头注意力

多头注意力便于模型从不同角度关注序列位置之间的关系，每个单头得到的输出 [B,T,head_size]，通过 concat 从 dim=-1（最后一个特征维）拼接起来，合并成一个大的矩阵 [B,T,n_embd]，再经过一层 Linear(n_embd, n_embd)，重新混合各个 head 的输出。

LayerNorm 层归一化

LayerNorm 的目的是让每一层收到的特征数值尺度更稳定，从而让神经网络更容易优化。它把每个 token 的特征重新整理成大致中心在 0 附近，整体尺度在 1 附近。标准化后还会经过可学习的缩放参数 γ 与平移参数 β，因此 LayerNorm 的最终输出并不被强制永远保持严格的均值 0、方差 1。这样下一层每次收到的输入尺度比较稳定，更容易训练。

具体在每个 block 内，配合残差连接，在 x 进入多头注意力前会 LayerNorm 一次，进入 FFN 前也会 LayerNorm 一次。

残差连接老生常谈，不再赘述。

FFN

FFN 层实现了一个两层 MLP，把 [B,T,C] 映射到 [B,T,4C]，然后 ReLU 一次，再从 [B,T,4C] 返回到 [B,T,C]。

第一层先把特征映射到更宽的中间空间，配合 ReLU 提供更强的非线性特征组合能力；第二层再把中间特征压回 C，以便继续残差相加和送入下一层。

这里 4C 是经典 Transformer 中常见的经验设计，教程也设定为4。

最后返回的 x 会经过多个 Transformer block 后得到最终 logits，本文的模型是3个 block。

训练与实践结果

训练采用了 AdamW 优化器，使用梯度的一阶矩（类似 momentum）和二阶矩，给不同参数自适应调整更新步长，并使用 decoupled weight decay。

超参数按照教程，设定为：

learning_rate = 3e-4
n_embd = 384
n_head = 4
n_layer = 3

模型参数：约5,468,225 ≈ 5.47M

上下文窗口：

模型 block_size=256，即单次前向传播最多接收 256 个 token 的上下文。由于这里采用字符级 tokenizer，因此相当于最多直接看到此前 256 个字符；在 [B,T,C] 中，T ≤ block_size。

对于每一个 optimization step：

xb, yb = get_batch('train')
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

自己训练的实现结果：

step 4500:
train loss 1.2175
val loss 1.5103

可以看到，生成的文本虽然细看没有什么逻辑和剧情可言，但是模型能说出来一堆“莎士比亚风味”的语言，有拼写正确的单词，格式大体上看起来和数据集的原始文本很相似。

当然，现在的模型还不是 ChatGPT 那种对话式的，教程只是做了一个简单的 Pretraining（预训练），还需要经过 Instruction / supervised fine-tuning（SFT）、Preference alignment 偏好对齐。这就是后话了，但是核心仍然是现在的预训练基座。

虽然现在的模型效果和常见的大模型比一言难尽，毕竟现在的模型非常小，只有 5M，同时这里使用的还是字符级 tokenizer，模型规模和训练数据都远小于现代 LLM，因此主要用于理解 GPT 的基本工作机制。

小结

经过《Let's build GPT: from scratch, in code, spelled out》视频的观看和实际的代码实现，我学习了 GPT Pretraining 的简单流程，对 Transformer 底层的结构和人工智能有了进一步理解。