nanoGPT

2025年11月更新：nanoGPT 有一个新的、改进的“兄弟”项目叫做 nanochat。你很可能想找/使用的是 nanochat。nanoGPT（本仓库）现在已经非常老旧且不再维护，但我会保留它作为历史记录。

这是用于训练/微调中等规模 GPT 模型最简单、最快速的仓库。它是 minGPT 的重写版本，更注重实用性而非教学。项目仍在积极开发中，但目前 train.py 文件可以在 OpenWebText 数据集上复现 GPT-2 (124M) 模型，在单个 8XA100 40GB 节点上大约需要 4 天的训练时间。代码本身简洁易读：train.py 是一个约 300 行的标准训练循环，model.py 是一个约 300 行的 GPT 模型定义，可以选择性地从 OpenAI 加载 GPT-2 的权重。仅此而已。

由于代码非常简单，很容易根据你的需求进行修改、从头开始训练新模型，或者微调预训练的检查点（例如，目前最大的起点模型是 OpenAI 的 GPT-2 1.3B 模型）。

安装

pip install torch numpy transformers datasets tiktoken wandb tqdm

依赖项：

• pytorch <3
• numpy <3
• transformers 用于 huggingface transformers <3（加载 GPT-2 检查点）
• datasets 用于 huggingface datasets <3（如果你想下载并预处理 OpenWebText）
• tiktoken 用于 OpenAI 的快速 BPE 编码 <3
• wandb 用于可选的日志记录 <3
• tqdm 用于进度条 <3

快速开始

如果你不是深度学习专业人士，只是想感受一下魔力并初步尝试，最快的方式是在莎士比亚作品上训练一个字符级 GPT。首先，我们将其下载为一个单独的 (1MB) 文件，并将原始文本转换为一个大的整数流：

python data/shakespeare_char/prepare.py

这会在数据目录中创建 train.bin 和 val.bin 文件。现在是时候训练你的 GPT 了。模型的大小很大程度上取决于你系统的计算资源：

我有一个 GPU。很好，我们可以使用 config/train_shakespeare_char.py 配置文件中的设置快速训练一个小型 GPT：

python train.py config/train_shakespeare_char.py

如果你查看这个配置文件，会看到我们正在训练一个上下文大小最大为 256 个字符、384 个特征通道的 GPT，它是一个 6 层 Transformer，每层有 6 个头。在一块 A100 GPU 上，这次训练大约需要 3 分钟，最佳验证损失为 1.4697。根据配置，模型检查点会被写入 --out_dir 指定的目录 out-shakespeare-char 中。因此，训练完成后，我们可以通过将采样脚本指向这个目录来从最佳模型中进行采样：

python sample.py --out_dir=out-shakespeare-char

这会生成一些样本，例如：

ANGELO:
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.

DUKE VINCENTIO:
I thank your eyes against it.

DUKE VINCENTIO:
Then will answer him to save the malm:
And what have you tyrannous shall do this?

DUKE VINCENTIO:
If you have done evils of all disposition
To end his power, the day of thrust for a common men
That I leave, to fight with over-liking
Hasting in a roseman.

哈哈 ¯\_(ツ)_/¯。对于在 GPU 上训练了 3 分钟的字符级模型来说，这还不错。通过在这个数据集上微调预训练的 GPT-2 模型，很可能获得更好的结果（参见后面的微调部分）。

我只有一台 Macbook（或其他廉价电脑）。别担心，我们仍然可以训练 GPT，但需要调低一些配置。我建议安装最新的 PyTorch nightly 版本（安装时在此选择），因为它目前很可能使你的代码更高效。但即使没有它，一个简单的训练运行也可以如下所示：

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

这里，由于我们在 CPU 上运行而不是 GPU，我们必须同时设置 --device=cpu 并关闭 PyTorch 2.0 编译 --compile=False。然后，当我们进行评估时，会得到一个更嘈杂但更快的估计（--eval_iters=20，从 200 下调），我们的上下文大小只有 64 个字符而不是 256，每次迭代的批量大小只有 12 个样本，而不是 64。我们还将使用一个更小的 Transformer（4 层，4 个头，128 维嵌入大小），并将迭代次数减少到 2000（并相应地通过 --lr_decay_iters 将学习率衰减到大约 max_iters）。因为我们的网络非常小，我们也降低了正则化强度（--dropout=0.0）。这仍然需要大约 3 分钟的运行时间，但得到的损失仅为 1.88，因此样本质量也较差，但仍然很有趣：

python sample.py --out_dir=out-shakespeare-char --device=cpu

生成如下样本：

GLEORKEN VINGHARD III:
Whell's the couse, the came light gacks,
And the for mought you in Aut fries the not high shee
bot thou the sought bechive in that to doth groan you,
No relving thee post mose the wear

对于在 CPU 上运行约 3 分钟来说，能捕捉到一点正确的角色特征，这还不错。如果你愿意等待更长时间，可以随意调整超参数，增加网络大小、上下文长度（--block_size）、训练时长等。

最后，在 Apple Silicon Macbook 上并使用较新的 PyTorch 版本时，请确保添加 --device=mps（“Metal Performance Shaders”的缩写）；PyTorch 随后会使用芯片上的 GPU，这可以显著加速训练（2-3 倍）并允许你使用更大的网络。更多信息请参见 Issue 28。

复现 GPT-2

更专业的深度学习从业者可能对复现 GPT-2 的结果更感兴趣。那么，我们开始吧——首先我们对数据集进行分词，这里使用的是 OpenWebText，它是 OpenAI（私有的）WebText 的一个开源复现：

python data/openwebtext/prepare.py

这会下载并分词 OpenWebText 数据集。它将创建 train.bin 和 val.bin 文件，其中存储了 GPT2 BPE 分词 ID 的序列，以原始 uint16 字节形式存储。然后我们就可以开始训练了。要复现 GPT-2 (124M)，你至少需要一个 8XA100 40GB 节点并运行：

torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

这将使用 PyTorch 分布式数据并行（DDP）运行大约 4 天，损失降至约 2.85。现在，一个仅在 OWT 上评估的 GPT-2 模型的验证损失约为 3.11，但如果你对其进行微调，它会下降到约 2.85 的范围（由于明显的领域差距），使两个模型大致匹配。

如果你在集群环境中，并且有幸拥有多个 GPU 节点，你可以让 GPU 飞速运行，例如跨 2 个节点：

# 在第一个（主）节点上运行，示例 IP 为 123.456.123.456：
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=123.456.123.456 --master_port=1234 train.py
# 在工作节点上运行：
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=123.456.123.456 --master_port=1234 train.py

最好对你的互连性能进行基准测试（例如使用 iperf3）。特别是，如果你没有 Infiniband，还需要在上述启动命令前加上 NCCL_IB_DISABLE=1。你的多节点训练将可以工作，但很可能非常缓慢。默认情况下，检查点会定期写入 --out_dir 目录。我们可以通过简单的 python sample.py 从模型中采样。

最后，要在单个 GPU 上训练，只需运行 python train.py 脚本。查看它的所有参数，该脚本力求非常易读、可修改和透明。你很可能需要根据你的需求调整其中许多变量。

基线

OpenAI GPT-2 检查点允许我们为 openwebtext 建立一些基线。我们可以通过以下方式获取数值：

$ python train.py config/eval_gpt2.py
$ python train.py config/eval_gpt2_medium.py
$ python train.py config/eval_gpt2_large.py
$ python train.py config/eval_gpt2_xl.py

并观察训练和验证的以下损失：

然而，我们必须注意，GPT-2 是在（封闭的、从未发布的）WebText 上训练的，而 OpenWebText 只是对该数据集的最佳开源复现。这意味着存在数据集领域差距。实际上，取 GPT-2 (124M) 检查点并在 OWT 上直接微调一段时间后，损失可以下降到约 2.85。这便成为更合适的复现基线。

微调

微调与训练没有区别，我们只需确保从预训练模型初始化并以较小的学习率进行训练。关于如何在新的文本上微调 GPT 的示例，请转到 data/shakespeare 并运行 prepare.py 来下载小型的莎士比亚数据集，并使用 GPT-2 的 OpenAI BPE 分词器将其转换为 train.bin 和 val.bin。与 OpenWebText 不同，这将在几秒钟内完成。微调可能只需要很少的时间，例如在单个 GPU 上只需几分钟。运行一个示例微调如下：

python train.py config/finetune_shakespeare.py

这将加载 config/finetune_shakespeare.py 中的配置参数覆盖（不过我没有过多调整它们）。基本上，我们使用 init_from 从 GPT2 检查点初始化，并像正常一样训练，只是时间更短且学习率较小。如果内存不足，尝试减小模型大小（它们是 {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}）或者可能减小 block_size（上下文长度）。最佳检查点（最低验证损失）将位于 out_dir 目录中，例如根据配置文件默认为 out-shakespeare。然后你可以运行 sample.py --out_dir=out-shakespeare 中的代码：

THEODORE:
Thou shalt sell me to the highest bidder: if I die,
I sell thee to the first; if I go mad,
I sell thee to the second; if I
lie, I sell thee to the third; if I slay,
I sell thee to the fourth: so buy or sell,
I tell thee again, thou shalt not sell my
possession.

JULIET:
And if thou steal, thou shalt not sell thyself.

THEODORE:
I do not steal; I sell the stolen goods.

THEODORE:
Thou know'st not what thou sell'st; thou, a woman,
Thou art ever a victim, a thing of no worth:
Thou hast no right, no right, but to be sold.

哇哦，GPT，你进入了一些黑暗的地方。我并没有在配置中过多调整超参数，请随意尝试！

采样 / 推理

使用脚本 sample.py 从 OpenAI 发布的预训练 GPT-2 模型或你自己训练的模型中进行采样。例如，以下是从最大的可用 gpt2-xl 模型采样的方法：

python sample.py \
    --init_from=gpt2-xl \
    --start="What is the answer to life, the universe, and everything?" \
    --num_samples=5 --max_new_tokens=100

如果你想从自己训练的模型中采样，使用 --out_dir 来指向正确的目录。你也可以用文件中的一些文本来提示模型，例如 python sample.py --start=FILE:prompt.txt。

效率说明

对于简单的模型基准测试和分析，bench.py 可能很有用。它与 train.py 训练循环核心部分发生的事情相同，但省略了许多其他复杂性。

请注意，代码默认使用 PyTorch 2.0。在撰写本文时（2022年12月29日），这使得 torch.compile() 在 nightly 版本中可用。这一行代码带来的改进是显著的，例如将每次迭代时间从约 250 毫秒/迭代减少到 135 毫秒/迭代。PyTorch 团队干得漂亮！

待办事项

• 研究并添加 FSDP 以替代 DDP
• 在标准评估集上评估零样本困惑度（例如 LAMBADA? HELM? 等）
• 优化微调脚本，我认为超参数设置得不是很好
• 在训练期间实现线性批量大小增长计划
• 整合其他嵌入方法（旋转嵌入、alibi 嵌入）
• 在检查点中将优化器缓冲区与模型参数分离
• 增加关于网络健康状况的额外日志记录（例如梯度裁剪事件、幅度）
• 围绕更好的初始化等进行更多研究

故障排除

请注意，默认情况下本仓库使用 PyTorch 2.0（即 torch.compile）。这是相当新且实验性的功能，并非在所有平台上都可用（例如 Windows）。如果你遇到相关的错误信息，尝试通过添加 --compile=False 标志来禁用它。这会减慢代码速度，但至少能运行。

关于这个仓库、GPT 和语言建模的背景，观看我的 Zero To Hero 系列可能会有所帮助。特别是，如果你有一些先前的语言建模背景，GPT 视频很受欢迎。

如果有更多问题/讨论，欢迎来 Discord 的 #nanoGPT 频道：

致谢

所有 nanoGPT 实验都由 Lambda labs 上的 GPU 提供支持，这是我最喜欢的云 GPU 提供商。感谢 Lambda labs 对 nanoGPT 的赞助！