nanoGPT 源码解析(一):项目概览与目录架构

nanoGPT 是Andrej Karpathy 的一个已经"过气"的项目——2025 年底作者在 README 里挂出了醒目的更新提示：新项目 nanochat 已经接棒，nanoGPT 本身"very old and deprecated"，只是保留下来供后人参考。但对于想真正搞懂 Transformer/GPT 训练全链路的系统工程师来说，nanoGPT 恰恰是最值得精读的版本——它比 nanochat 更小、更纯粹，没有被后续工程化诉求（分词器定制、RL 阶段、多阶段流水线）稀释掉核心逻辑。

这篇文章从"这是一个什么项目、代码是怎么组织的"这个角度切入，先把地图画清楚，后面的文章再深入某几个具体的技术点（比如 DDP 梯度同步、Flash Attention fallback、检查点恢复机制）。

https://github.com/karpathy/nanoGPT

nanoGPT是什么

nanoGPT 的自我定位很直白：the simplest, fastest repository for training/finetuning medium-sized GPTs。它是作者早年 minGPT 项目的重写版，设计取舍上"prioritizes teeth over education"——比 minGPT 更看重能跑起来的实际效果，而不是纯粹的教学可读性。

核心事实：

train.py
复现 GPT-2 (124M)，单节点 8×A100 40GB 大约 4 天训练完
train.py
和 model.py 各约 300 行
支持从零训练，也支持加载 OpenAI 发布的 GPT-2 权重做微调

一个常被忽略但很重要的信息是它的血统：GPT-2 官方 TensorFlow 实现和 HuggingFace Transformers 的 PyTorch 实现是 model.py 文件头部明确列出的两个参照对象。也就是说 nanoGPT 不是一个"发明新架构"的项目，而是一次"用最少代码复刻标准架构"的工程练习，这也是为什么它特别适合用来学习标准 GPT 结构，而不必担心某些非主流设计干扰理解。

目录全貌

nanoGPT/├── model.py              # 模型定义（~300行，GPTConfig + GPT + 各子模块）├── train.py              # 训练主循环（~300行）├── sample.py             # 推理/采样脚本├── bench.py              # 纯性能基准测试，剥离了训练循环里的评估/日志/checkpoint逻辑├── configurator.py       # ~20行的配置覆写工具├── config/               # 各训练/评估场景的参数集合│   ├── train_gpt2.py│   ├── train_shakespeare_char.py│   ├── finetune_shakespeare.py│   └── eval_gpt2*.py├── data/                            # 数据集准备脚本，每个子目录一个数据集│   ├── shakespeare/prepare.py       # 词级，走 GPT-2 BPE tokenizer│   ├── shakespeare_char/prepare.py  # 字符级，自建 vocab│   └── openwebtext/prepare.py       # 生产级预训练语料（GPT-2 复现用）├── scaling_laws.ipynb               # 分析笔记本，不参与训练主流程├── transformer_sizing.ipynb         # 分析笔记本，不参与训练主流程└── assets/                          # README 用图

这个结构没有 src/、没有 lib/、没有多层嵌套的包结构——所有核心逻辑都摊平在根目录下，这本身就是一种设计声明：代码总量小到不需要用目录层级去管理复杂度。对比一下你如果去看 HuggingFace Transformers 那种数万文件的仓库，会立刻感受到这种"克制"是刻意为之，不是偷懒。

按职责拆分各模块

数据层：data/*/prepare.py

三个数据集脚本各自独立，互不依赖，但产出格式完全统一：每个目录跑完之后都会在本地生成 train.bin、val.bin（uint16 裸字节流，存的是 token id 序列），字符级数据集额外产出一个 meta.pkl（记录 vocab_size 和字符/id 双向映射表）。

这里有个很值得学习的接口设计：train.py 完全不知道、也不关心数据是莎士比亚还是 OpenWebText，它只认目录下有没有 train.bin/val.bin 这两个文件。加一个新数据集，只需要照着现有 prepare.py 写一份新的分词/切片逻辑，训练代码一行都不用改。这是典型的"用文件格式做接口契约"而不是"用代码接口做契约"——契约足够简单（裸 token id 数组），复杂度可以完全下沉到各自的 prepare.py 里。

模型层：model.py

严格意义上的"纯函数式"模块：GPTConfig（dataclass 定义的超参集合）+ GPT（nn.Module），不引用任何训练相关的全局状态（没有 optimizer、没有 DDP、没有数据加载）。对外暴露的能力边界很清晰：

forward(idx, targets)
—— 训练/评估用，给 targets 就顺带算 loss
generate(idx, max_new_tokens, ...)
—— 推理用，自回归采样
from_pretrained(model_type)
—— 加载 HuggingFace 上的 GPT-2 官方权重
configure_optimizers(...)
—— 把参数分组（矩阵参数走 weight decay，bias/LayerNorm 不走）
estimate_mfu(...)
—— 算力利用率估算，纯离线分析用

这种"计算图定义"与"训练调度"彻底解耦的写法，换来的直接好处是：train.py、sample.py、bench.py 三个完全不同用途的脚本能共享同一份 model.py，不需要为不同场景维护不同的模型定义版本。

训练层：train.py

这是整个项目里逻辑最重的文件，但依然控制在 300 行左右。它承担几件事：

环境识别
：靠检查 RANK 环境变量是否存在来判断是否处于 torchrun 拉起的 DDP 环境，单卡调试和多卡集群训练用的是完全同一份代码
三种初始化路径
：scratch（从零训练）/ resume（从 checkpoint 续训）/ gpt2*（加载官方权重做微调），由 init_from 一个字符串参数分发
主循环
：学习率调度（带 warmup 的 cosine decay）→ 梯度累积的 micro-step 循环 → 梯度裁剪 → optimizer step → 定期验证 + 保存 checkpoint

它和 model.py 之间只有一处硬编码耦合：from model import GPTConfig, GPT。这条边界切得很干净，也是为什么这个项目虽然极简却依然工程感十足的原因之一。

配置层：config/ + configurator.py

这是整个仓库里最"取巧"的一块设计。configurator.py 只有二十来行，核心逻辑是：先收集 train.py 顶部所有基础类型（int/float/bool/str）的全局变量名作为 config_keys，然后用 exec(open(config_file).read()) 直接把传入的配置文件当 Python 代码执行，从而原地覆盖这些全局变量；命令行传的 --key=value 走同样逻辑处理。

没有 argparse，没有 Config 类，没有 schema 校验。代价是没有类型安全、没有 IDE 补全提示；换来的是新增一个配置项不需要改任何"框架代码"——直接在 train.py 顶部声明一个新变量，它自动就能被 config/*.py 或命令行覆盖。这是典型的"教学优先"取舍，如果是要长期维护的生产系统，这套机制通常是要被替换掉的。

推理与基准：sample.py / bench.py

两者都是"训练主流程的旁支"，都只依赖 model.py 和一个 checkpoint，彼此独立、互不调用。bench.py 的存在很值得注意——它把 train.py 主循环里的核心 forward/backward 逻辑单独抽出来，去掉了评估、日志、checkpoint 这些"噪音"，专门用来做性能剖析。这种"从主流程里再抽一份纯净的性能测试版本"的做法，在系统软件里也很常见（想想你平时怎么从完整业务代码里抽 microbenchmark）。

模块间的数据流总览

data/*/prepare.py     │  (生成 train.bin / val.bin / meta.pkl)     ▼train.py ──uses──> model.py (GPTConfig, GPT)     │     ├──> configurator.py (读取 config/*.py 做参数覆盖)     │     └──> out/ckpt.pt (model + optimizer + iter_num + config 全量状态)                │                ▼          sample.py (独立加载 ckpt.pt 做推理)ta/*/prepare.py      │  (生成 train.bin / val.bin / meta.pkl)      ▼ train.py ──uses──> model.py (GPTConfig, GPT)      │      ├──> configurator.py (读取 config/*.py 做参数覆盖)      │      └──> out/ckpt.pt (model + optimizer + iter_num + config 全量状态)                 │                 ▼           sample.py (独立加载 ckpt.pt 做推理)

整个链路里，model.py 是唯一的"中心节点"，其余模块都是围绕它的调度/数据准备/结果消费。这种星型依赖结构，是这个项目虽小但读起来不觉得混乱的关键原因——任何时候你只需要问"这段代码是在准备数据、调度训练，还是消费模型"，就能立刻定位到该看哪个文件。

小结

从目录结构就能读出 nanoGPT 的设计哲学：用文件系统层面的极简（摊平、少嵌套）配合职责边界的清晰切分，把复杂度死死摁在 300 行的量级上。它刻意放弃了灵活的模型架构变体、复杂的数据管线、类型安全的配置系统——这些放弃换来的是任何人都能在一两个小时内读完整个仓库，建立起对 GPT 训练全链路的完整心智模型。

下一篇会往下钻一层，看几个具体的技术实现细节：DDP 训练里梯度同步是怎么被"手动"优化掉冗余 all-reduce 的、Flash Attention 的 fallback 逻辑、以及 checkpoint 恢复机制里那些容易被忽略的边界条件处理。