600行代码复现GPT-2:Karpathy nanoGPT源码全解
600行代码复现GPT-2:Karpathy nanoGPT源码全解
当整个生态都在卷万亿参数时,Karpathy 用两个文件、约 600 行代码,把 GPT-2 训练这件事讲明白了。
理解大模型,没有比 nanoGPT 更合适的起点。它不是玩具,而是真能在 8 张 A100 上花 4 天复现 GPT-2 (124M) 的可用实现;它也不是工业框架,而是把训练循环压到 300 行、模型定义压到 300 行的极简范本。读懂这 600 行,你就读懂了一个 GPT 从初始化到生成 token 的全部主干。
下面把 model.py 和 train.py 逐块拆开讲,配上能直接抄走的工程技巧。
本文提纲
- nanoGPT 是什么,以及它现在「过气」了吗
- 仓库结构:两个文件撑起一切
model.py全解:300 行写完一个 GPTtrain.py全解:300 行写完一个训练循环- GPT-2 复现的成绩单
- 30 分钟上手:从莎士比亚到你的 MacBook
- 值得偷走的 8 个工程技巧
- 现状:nanochat 与 build-nanogpt 视频
nanoGPT 是什么,以及它现在「过气」了吗
nanoGPT 是 Andrej Karpathy 对自己早期项目 minGPT 的重写。定位上他自己说得很直白:prioritizes teeth over education——优先要「能咬人的牙」,而不是教学版的软糖。换句话说,minGPT 偏教学,nanoGPT 偏「真能训」。
它的核心承诺只有一句:train.py 是约 300 行的样板训练循环,model.py 是约 300 行的 GPT 模型定义,可选加载 OpenAI 的 GPT-2 权重。仅此而已。正因为代码这么简单,你很容易 hack 成自己想要的样子:从头训、微调预训练 checkpoint,都没区别。
2025 年 11 月 Karpathy 在 README 顶部加了一条更新:nanoGPT 已有更现代化的继任者 nanochat,本仓库标记为 deprecated 但保留以供后人参考。所以别再用它去做生产级训练,但作为学透 GPT 内部结构的教材,它至今没有对手。配套的 build-nanogpt 仓库把「从空文件到复现 GPT-2 (124M)」拆成 44 个 git commit,配合 YouTube 视频逐 commit 讲解,现在复现 124M 只要约 1 小时、约 10 美元云 GPU 成本。
一句话定位:nanoGPT 是「读源码学 GPT」的事实标准,本文要做的就是把这两个文件讲透。
仓库结构:两个文件撑起一切
graph TB
subgraph "nanoGPT Repository"
A["model.py
~300 lines
GPT definition"]
B["train.py
~300 lines
Training loop"]
C["configurator.py
CLI config override"]
D["sample.py
Inference / sampling"]
E["bench.py
Throughput benchmark"]
F["config/*.py
Preset hyperparameters"]
G["data/
Dataset prepare scripts"]
end
A --> B
C --> B
F --> C
G --> B
B --> D
style A fill:#FF6B6B,color:#000000
style B fill:#FF6B6B,color:#000000
style C fill:#4ECDC4,color:#000000
style D fill:#4ECDC4,color:#000000
style E fill:#4ECDC4,color:#000000
style F fill:#96CEB4,color:#000000
style G fill:#96CEB4,color:#000000真正需要读的只有红框两个文件。其余都是辅助:
configurator.py:一个 ~50 行的小魔法,让命令行--n_layer=8或config/train_gpt2.py里的变量覆盖train.py顶部的默认值。实现就是exec(open(...).read())加一点 introspection。sample.py:采样脚本,可以从 OpenAI 的预训练 GPT-2,也可以从你自己--out_dir里的 checkpoint 生成。bench.py:纯吞吐量 benchmark,只保留训练循环里的核心前向反向。config/:预设超参,例如train_shakespeare_char.py、train_gpt2.py、finetune_shakespeare.py、eval_gpt2.py等。data/:每个数据集一个目录,里面有prepare.py,把原始文本 tokenize 成train.bin/val.bin(uint16 的 token id 长流)。
理解了这个布局,下面直接进源码。
model.py 全解:300 行写完一个 GPT
整个模型由五个组件构成:LayerNorm、CausalSelfAttention、MLP、Block、GPT。从下往上看,每个都极短。
GPTConfig:唯一的数据类
@dataclass
class GPTConfig:
block_size: int = 1024
vocab_size: int = 50304 # GPT-2 vocab 50257,向上取整到 64 的倍数提效
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
dropout: float = 0.0
bias: bool = True注意 vocab_size = 50304 这个细节:GPT-2 真实词表是 50257,但 50304 是 64 的倍数。在现代 GPU 上,矩阵维度对齐到 64 能让 tensor core 跑满,这点浪费换来的吞吐值得。bias=False 是个更快更好的选项,加载 OpenAI checkpoint 时会被强制设回 True。
LayerNorm:为什么要自己写
class LayerNorm(nn.Module):
def __init__(self, ndim, bias):
super().__init__()
self.weight = nn.Parameter(torch.ones(ndim))
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
def forward(self, input):
return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)PyTorch 内置 nn.LayerNorm 不支持 bias=False,所以 Karpathy 自己包了一层。直接调 F.layer_norm,传 None 当 bias 即可。这种「为了一行 API 差异自己写一个 module」的取舍贯穿整个 nanoGPT——能少一层抽象就少一层。
CausalSelfAttention:一行 Flash Attention
这是全文件含金量最高的 30 行:
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
# Q/K/V 三个投影合并成一个大 Linear,一次算完
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.n_head = config.n_head
self.n_embd = config.n_embd
self.dropout = config.dropout
# PyTorch 2.0+ 自带 Flash Attention
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
if not self.flash:
# 退化路径:预计算一个下三角 causal mask
self.register_buffer("bias",
torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
def forward(self, x):
B, T, C = x.size()
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
# 把 head 维提到 batch 维:(B, nh, T, hs)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
if self.flash:
y = F.scaled_dot_product_attention(
q, k, v, attn_mask=None,
dropout_p=self.dropout if self.training else 0,
is_causal=True)
else:
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v
y = y.transpose(1, 2).contiguous().view(B, T, C)
y = self.resid_dropout(self.c_proj(y))
return y三个值得注意的点:
第一,Q/K/V 合并投影。 GPT-2 原版用三个独立的 Conv1D,nanoGPT 合并成 nn.Linear(n_embd, 3 * n_embd),一次 matmul 出 Q/K/V 再 split。少两次 kernel launch,在长序列上能省下可观的时间。
第二,Flash Attention 的优雅退化。 不需要装 flash-attn 这个麻烦的包,直接用 PyTorch 2.0 内置的 scaled_dot_product_attention,传 is_causal=True 让 kernel 自己处理 causal mask——既不用物化 T×T 的注意力矩阵,也不用显式传 mask buffer。老版本 PyTorch 走手写 softmax 路径,靠预注册的下三角 buffer 做掩码。两套实现并存,是这个项目「能跑就行 + 跑得快」哲学的缩影。
第三,head 维提前。 transpose(1, 2) 把 (B, T, nh, hs) 变成 (B, nh, T, hs),让 head 成为 batch 维。这是后续所有 attention 计算能 batch 化的前提。
MLP:4 倍扩展 + GELU
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
self.gelu = nn.GELU()
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
self.dropout = nn.Dropout(config.dropout)标准 GPT-2 配方:隐层扩 4 倍、GELU 激活、再投回去。没什么花活,但注意 4 * n_embd 这个比例是 Transformer 原论文定下的,nanoGPT 没改。
Block:Pre-LN 残差结构
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention(config)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return xPre-LN:LayerNorm 在子层之前,残差走在外层。两行 forward 就是这个结构的全部。这跟原始 Transformer 的 Post-LN 不同——Post-LN 把残差加完再归一化,深层训练不稳;Pre-LN 让梯度能沿残差主干无衰减地回流,是 GPT-2/3 及之后所有大模型的标配。nanoGPT 用最直白的方式把这个事实摆在代码里。
GPT 主类:权重共享与特殊初始化
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd), # token embedding
wpe = nn.Embedding(config.block_size, config.n_embd), # position embedding
drop = nn.Dropout(config.dropout),
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f = LayerNorm(config.n_embd, bias=config.bias),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# 权重共享:embedding 和 lm_head 用同一份权重
self.transformer.wte.weight = self.lm_head.weight
self.apply(self._init_weights)
# 残差投影做缩放初始化
for pn, p in self.named_parameters():
if pn.endswith('c_proj.weight'):
torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))三个关键决策:
权重共享(weight tying)。 最后一层 lm_head 和 token embedding wte 共用同一矩阵。理由是直觉性的:embedding 把 token 映射到语义空间,lm_head 把语义空间映回词表,二者互为逆映射,没必要学两份。这一招能省下 vocab_size × n_embd 个参数——在 GPT-2 124M 里大约是 38M,占总参数近 30%。
残差投影的缩放初始化。 所有 c_proj.weight(每个 Block 里 attention 和 MLP 的输出投影)的初始化标准差从 0.02 缩放到 0.02 / sqrt(2 * n_layer)。这是 GPT-2 论文的做法:深层网络里残差路径会累加,输出投影缩放后能让残差贡献随层数衰减,避免训练初期方差爆炸。2 * n_layer 是因为每层有两个残差(attn + mlp)。
参数计数扣掉 position embedding。 get_num_params() 默认把 wpe 减掉,因为 position embedding 严格说不算「模型容量」。token embedding 因为和 lm_head 共享,仍计入。
forward 与 generate 的不对称
def forward(self, idx, targets=None):
b, t = idx.size()
pos = torch.arange(0, t, dtype=torch.long, device=device)
tok_emb = self.transformer.wte(idx)
pos_emb = self.transformer.wpe(pos)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
x = block(x)
x = self.transformer.ln_f(x)
if targets is not None:
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
targets.view(-1), ignore_index=-1)
else:
# 推理时只算最后一个位置的 logits,省一大笔计算
logits = self.lm_head(x[:, [-1], :])
loss = None
return logits, loss训练时算全部位置的 logits 做 cross-entropy;推理时只算最后一位。x[:, [-1], :] 用列表索引保留时间维,这是个容易踩的小坑——写成 x[:, -1, :] 会少一维,后面 softmax 就炸。这个不对称是自回归生成的关键优化:每生成一个 token,前 T-1 个位置的 logits 根本用不上。
generate 是标准的「喂回采样」循环,支持 temperature 和 top-k:
@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
for _ in range(max_new_tokens):
idx_cond = idx if idx.size(1) <= self.config.block_size \
else idx[:, -self.config.block_size:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx注意超出 block_size 时直接裁窗口(idx[:, -block_size:]),不做 KV-cache——这是教学取舍,工业实现会做 cache。
从预训练权重加载:from_pretrained
GPT.from_pretrained 把 HuggingFace 的 GPT2LMHeadModel 权重搬到 nanoGPT 结构。唯一需要小心的点是 OpenAI 原版用 Conv1D(权重形状是 (in, out)),nanoGPT 用 nn.Linear(形状 (out, in)),所以四个权重矩阵需要转置:
transposed = ['attn.c_attn.weight', 'attn.c_proj.weight',
'mlp.c_fc.weight', 'mlp.c_proj.weight']
for k in sd_keys_hf:
if any(k.endswith(w) for w in transposed:
sd[k].copy_(sd_hf[k].t()) # 转置
else:
sd[k].copy_(sd_hf[k])这种「显式列出需要特殊处理的键」的做法,比写一堆 if-else 还清楚。
configure_optimizers:分组的 AdamW
def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
# 2D 及以上参数做 weight decay,1D(bias、layernorm)不做
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{'params': decay_params, 'weight_decay': weight_decay},
{'params': nodecay_params, 'weight_decay': 0.0}
]
# 用 fused AdamW(CUDA 上更快)
fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and device_type == 'cuda'
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate,
betas=betas, **(dict(fused=True) if use_fused else {}))
return optimizer两个工程细节:按维度分组 decay——所有 2D+(权重矩阵、embedding)才施加 weight decay,1D(bias、LayerNorm 的 weight/bias)不施加。这是 GPT-2 以来的标准做法,因为对归一化参数做 decay 会破坏其标定作用。fused AdamW——CUDA 上用融合实现,少几次 kernel launch。inspect.signature 检测可用性的写法很干净,比 try/except 优雅。
train.py 全解:300 行写完一个训练循环
train.py 顶部是一大段默认配置,然后是 DDP 初始化、数据加载、学习率调度、训练循环。逐块看。
默认配置:GPT-2 124M 的 Chinchilla 调参
# data
dataset = 'openwebtext'
gradient_accumulation_steps = 5 * 8 # 梯度累积模拟大 batch
batch_size = 12 # micro-batch
block_size = 1024
# model
n_layer = 12
n_head = 12
n_embd = 768
# adamw optimizer
learning_rate = 6e-4
max_iters = 600000
weight_decay = 1e-1
beta1, beta2 = 0.9, 0.95
grad_clip = 1.0
# learning rate decay
decay_lr = True
warmup_iters = 2000
lr_decay_iters = 600000
min_lr = 6e-5 # ~= learning_rate / 10几个数字背后的道理:
gradient_accumulation_steps = 5 * 8 = 40,单卡 micro-batch 12、上下文 1024,8 卡累积 40 步,等效 batch 约 0.5M tokens。这是为了在显存有限的情况下逼近 GPT-2 原版的大 batch。beta2 = 0.95(默认是 0.999),GPT-3 论文里专门调过,语言模型梯度方差大,更小的 beta2 让二阶矩估计跟得更紧。min_lr = learning_rate / 10,lr_decay_iters ≈ max_iters——这是 Chinchilla 论文里推荐的 cosine decay 终值和衰减长度。max_iters = 600000,按上面 tokens/iter 算约 295B tokens。GPT-2 124M 的 Chinchilla 最优是 ~2.5B tokens,所以 nanoGPT 是严重过训练的——这正是 GPT-2 当年的做法,用算力换下游任务表现。
DDP 与梯度累积的配合
ddp = int(os.environ.get('RANK', -1)) != -1
if ddp:
init_process_group(backend=backend)
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)
master_process = ddp_rank == 0
# 梯度累积步数按 world_size 等比缩放
assert gradient_accumulation_steps % ddp_world_size == 0
gradient_accumulation_steps //= ddp_world_sizeRANK 环境变量存在即说明是 torchrun 启动的 DDP。关键优化在训练循环里:
for micro_step in range(gradient_accumulation_steps):
if ddp:
# 只在最后一个 micro step 同步梯度,省掉前面 N-1 次的 all-reduce
model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
with ctx:
logits, loss = model(X, Y)
loss = loss / gradient_accumulation_steps
X, Y = get_batch('train') # 异步预取下一批
scaler.scale(loss).backward()require_backward_grad_sync 这个标志位是 DDP 内部 no_sync() 上下文管理器toggle 的同一个变量。Karpathy 在注释里说得很坦白:官方 no_sync() 写法会让代码变臃肿、还要重复代码,干脆直接拨这个 flag,效果一样。这是 nanoGPT 风格的典型体现——不为了「正确性」引入不必要的抽象。
「穷人的数据加载器」:memmap
def get_batch(split):
if split == 'train':
data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
else:
data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
if device_type == 'cuda':
x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
else:
x, y = x.to(device), y.to(device)
return x, y整个数据加载:把全量 token 存成 uint16 的 train.bin,每次 np.memmap 重新打开(注释说为了避免 memmap 的内存泄漏),随机抽 batch_size 个起始位置,切出 x 和错位一位的 y。pin_memory().to(non_blocking=True) 让 H2D 拷贝异步——配合上面训练循环里「先 forward 再预取下一批」的顺序,数据搬运和计算重叠。
没有 DataLoader、没有 sampler、没有 collate。朴素到极致,但跑起来就是快。
Cosine 学习率调度:warmup + cosine + floor
def get_lr(it):
if it < warmup_iters:
return learning_rate * (it + 1) / (warmup_iters + 1) # 线性 warmup
if it > lr_decay_iters:
return min_lr # 衰减到底后保持
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (learning_rate - min_lr) # cosine 衰减三段式:前 2000 步线性 warmup(防止初期 Adam 二阶矩未稳时发散),中间从 learning_rate cosine 衰减到 min_lr,之后保持 min_lr。没有用 torch.optim.lr_scheduler,手写一个函数在循环里手动设 param_group['lr']——又是「少一层抽象」的取舍。
训练循环主体
去掉日志和 checkpoint 后,核心循环就是:
while True:
lr = get_lr(iter_num) if decay_lr else learning_rate
for param_group in optimizer.param_groups:
param_group['lr'] = lr
for micro_step in range(gradient_accumulation_steps):
if ddp:
model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
with ctx: # autocast 上下文,自动混合精度
logits, loss = model(X, Y)
loss = loss / gradient_accumulation_steps
X, Y = get_batch('train')
scaler.scale(loss).backward()
if grad_clip != 0.0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True) # set_to_none=True 比 zero 更省内存
iter_num += 1
if iter_num > max_iters:
break完整训练循环就这么多。混合精度用 torch.amp.autocast + GradScaler:bfloat16 时不缩放(数值范围够),float16 时用 scaler 防下溢。zero_grad(set_to_none=True) 把梯度直接置 None 而不是填 0,省下一次 memset,长训练能省下可观显存。
torch.compile 一行提速
if compile:
print("compiling the model... (takes a ~minute)")
model = torch.compile(model) # PyTorch 2.0README 里给了实测:单卡迭代时间从 ~250ms 降到 135ms,约 1.85×。代价是第一次 forward 会花一分钟编译。注意 compile 默认 True,在某些平台(Windows、老 PyTorch)上会报错,这时加 --compile=False 退化。
GPT-2 复现的成绩单
跑 torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py,8×A100 40GB 训练约 4 天,val loss 落到 ~2.85。对照表(用 OpenAI 公开的 GPT-2 checkpoint 在 OpenWebText 上评测):
| 模型 | 参数量 | 训练 loss | 验证 loss |
|---|---|---|---|
| gpt2 | 124M | 3.11 | 3.12 |
| gpt2-medium | 350M | 2.85 | 2.84 |
| gpt2-large | 774M | 2.66 | 2.67 |
| gpt2-xl | 1558M | 2.56 | 2.54 |
注意 GPT-2 原版是在闭源的 WebText 上训的,OpenWebText 只是社区复刻,存在域差距。把 OpenAI 的 124M checkpoint 拿到 OWT 上微调一阵,loss 也能降到 ~2.85——所以 nanoGPT 从零训出的 124M 实际上「打平」了 OpenAI 的官方模型(在 OWT 域上)。这个数字是判断「我真的复现了 GPT-2」的基准线。
30 分钟上手:从莎士比亚到你的 MacBook
不想烧钱跑 4 天?nanoGPT 准备了 3 分钟入门套餐。先 tokenize 莎士比亚全集(1MB 文本)成字符级 token:
python data/shakespeare_char/prepare.py然后三种机器三套配置:
有 A100——直接用预设,3 分钟训完,val loss 1.4697:
python train.py config/train_shakespeare_char.py预设是 block_size=256, n_layer=6, n_head=6, n_embd=384 的小 GPT。采样看看:
python sample.py --out_dir=out-shakespeare-charANGELO:
And cowards it be strawn to my bed,
And thrust the gates of my threats,
Because he that ale away, and hang'd
An one with him.3 分钟、字符级、能看出莎士比亚的「腔调」——值这个票价。
只有 CPU——缩小模型、关 compile、关 GPU 路径:
python train.py config/train_shakespeare_char.py \
--device=cpu --compile=False --eval_iters=20 --log_interval=1 \
--block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 \
--max_iters=2000 --lr_decay_iters=2000 --dropout=0.0约 3 分钟,val loss 1.88,差一些但能跑。关键点:CPU 必须 --device=cpu 且 --compile=False。
Apple Silicon Mac——加 --device=mps,用 M 系列的 GPU 加速 2-3×,可以上更大网络。
微调也很简单:data/shakespeare 里有 BPE 版的莎士比亚,python train.py config/finetune_shakespeare.py 从 GPT-2 checkpoint 初始化、用更小学习率继续训,单卡几分钟就能出像样的结果。
值得偷走的 8 个工程技巧
读完源码,下面这些是任何 PyTorch 训练项目都能直接抄走的:
- vocab_size 对齐到 64 的倍数——GPU tensor core 友好,几行代码换稳定提速。
- Q/K/V 合并投影——
nn.Linear(n_embd, 3*n_embd)一次算完,少 kernel launch。 - 直接用
scaled_dot_product_attention+is_causal=True——不用装flash-attn,免维护 causal mask buffer。 - 权重共享 embedding 和 lm_head——省一大块参数,还常常带来轻微性能提升。
- AdamW 按维度分组 decay——2D+ decay、1D 不 decay,是 LLM 训练的隐形标准。
- 梯度累积时手动拨
require_backward_grad_sync——比no_sync()上下文管理器更省代码,N-1 次 all-reduce 免了。 memmap+ 随机切窗当数据加载器——Token 流够小(GPT-2 词表 uint16)时,这比任何 DataLoader 都轻。optimizer.zero_grad(set_to_none=True)——一行配置省一次 memset,长训练里能省下可观显存。
这八条加起来,往往就是一个普通训练脚本和「能跑满 A100」的训练脚本的差距。
现状:nanochat 与 build-nanogpt 视频
最后说清楚 nanoGPT 现在的位置。
它已 deprecated,但不是没用。 Karpathy 在 2025 年 11 月明确推荐继任者 nanochat——更现代的架构、更完整的 chat finetuning 流程。如果你要做实际训练,直接上 nanochat。
但 nanoGPT 作为「读源码学 GPT」的教材无可替代。 配套的 build-nanogpt 仓库把从空文件到 GPT-2 (124M) 的过程拆成 44 个 git commit,每步都能 git checkout 看差异,配合 YouTube 视频《Let's reproduce GPT-2 (124M)》逐 commit 讲解。现在云端租 GPU 复现 124M 大约 1 小时、10 美元——这是普通人能负担得起的「亲手训一个 GPT」体验。
需要补一句认知边界:nanoGPT / build-nanogpt 只覆盖预训练(next-token prediction),训出来的是「会做梦互联网文档」的语言模型,不是能对话的 ChatGPT。SFT、RLHF、对话格式这些 chat finetuning 是之后的事,nanochat 才把这条链路补齐。
所以正确的姿势是:用 nanoGPT 读懂主干、用 build-nanogpt 视频走一遍手感、用 nanochat 走完整的 chat 训练链路。三个项目各司其职,构成了 Karpathy 给 LLM 学习者铺的完整台阶。
作者: itech001 来源: 公众号:AI人工智能时代(the-ai-era) 网站: https://www.theaiera.top/ 关注每日最新AI新闻和技术博客,主页有更多的文章的AI 技术参考:https://www.theaiera.top
本文首发于 AI人工智能时代,转载请注明出处。