nanoGPT

这是一个用于训练/微调中等规模GPT模型的最简单、最快速的代码库。它是minGPT的重写版本，优先考虑实用性而非教育性。目前仍在积极开发中，但现在的train.py文件可以在单个8XA100 40GB节点上重现GPT-2（124M）模型，在OpenWebText数据集上训练约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数据集 <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

如果您查看其内容，您会看到我们正在训练一个GPT模型，它的上下文大小最多为256个字符，384个特征通道，是一个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每日构建版（在这里选择），因为它很可能会让你的代码更高效。但即使没有它，一个简单的训练运行可能如下所示：

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附近）。因为我们的网络非常小，我们还降低了正则化（--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），你至少需要一个8X A100 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，下载小型莎士比亚数据集并将其渲染为train.bin和val.bin，使用GPT-2的OpenAI BPE分词器。与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中的代码：

西奥多：
你要把我卖给出价最高的人：如果我死了，
我就把你卖给第一个人；如果我疯了，
我就把你卖给第二个人；如果我
撒谎，我就把你卖给第三个人；如果我杀人，
我就把你卖给第四个人：所以买或卖，
我再告诉你一遍，你不能卖掉我的
财产。

朱丽叶：
如果你偷窃，你就不能出卖自己。

西奥多：
我不偷窃；我卖的是赃物。

西奥多：
你不知道你在卖什么；你，一个女人，
你永远是受害者，一个毫无价值的东西：
你没有权利，没有权利，只能被出售。

哇，GPT，你在那里进入了一些黑暗的地方。我没有真正调整配置中的超参数，欢迎尝试！

采样/推理

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

python sample.py \
    --init_from=gpt2-xl \
    --start="生命、宇宙以及一切的答案是什么？" \
    --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()在夜间版本中可用。这一行代码带来的改进显著，例如将迭代时间从每次迭代约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！