Andrej Karpathy 重写minGPT库，NanoGPT上线速揽2500星

前特斯拉 AI 负责人 Andrej Karpathy 基于 PyTorch，仅用 300 行左右的代码就写出了一个小型 GPT 训练库，并将其命名为 minGPT。这个 minGPT 能够进行加法运算和字符级的语言建模，而且准确率还不错。

时隔两年，minGPT 迎来更新，Karpathy 又上线新版本，并命名为 NanoGPT，该库用于训练和微调中型大小的 GPT。上线短短几天，狂揽 2.5K 星。

在项目介绍中，Karpathy 这样写道：「NanoGPT 是用于训练和微调中型尺度 GPT 最简单、最快的库。是对 minGPT 的重写，因为 minGPT 太复杂了，以至于我都不愿意在使用它。NanoGPT 还在开发当中，当前致力于在 OpenWebText 数据集上重现 GPT-2。

NanoGPT 代码设计目标是简单易读，其中 train.py 是一个约 300 行的代码；model.py 是一个约 300 行的 GPT 模型定义，它可以选择从 OpenAI 加载 GPT-2 权重。」

为了呈现数据集，用户首先需要将一些文档 tokenize 为一个简单的 1D 索引数组。

$ cd data/openwebtext
$ python prepare.py


这将生成两个文件：train.bin 和 val.bin，每个文件都包含一个代表 GPT-2 BPE token id 的 uint16 字节原始序列。该训练脚本试图复制 OpenAI 提供的最小的 GPT-2 版本，即 124M 版本。

$ python train.py


假如你想使用 PyTorch 分布式数据并行（DDP）进行训练，请使用 torchrun 运行脚本。

$ torchrun --standalone --nproc_per_node=4 train.py


为了让代码更有效，用户也可以从模型中进行取样：

$ python sample.py

Karpathy 表示，该项目目前在 1 个 A100 40GB GPU 上一晚上的训练损失约为 3.74，在 4 个 GPU 上训练损失约为 3.60。在 8 x A100 40GB node 上进行 400,000 次迭代（约 1 天）atm 的训练降至 3.1。

至于如何在新文本上微调 GPT，用户可以访问 data/shakespeare 并查看 prepare.py。与 OpenWebText 不同，这将在几秒钟内运行。微调只需要很少的时间，例如在单个 GPU 上只需要几分钟。下面是运行微调的一个例子

$ python train.py config/finetune_shakespeare.py