C 语言从零复现 GPT-2 完整实现

通知 必须登录才能更改通知设置

加载时发生错误。请重新加载此页面。

加载时发生错误。请重新加载此页面。

karpathy 2024 年 5 月 28 日 维护者

让我们用 llm.c（约 4,000 行 C/CUDA 代码），花 90 分钟和 20 美元复现 GPT-2（124M）。124M 模型是 OpenAI 于 2019 年发布的 GPT-2 系列中最小的模型。如今，即使对 GPU 资源有限的人来说，它实际上也相当容易上手。llm.c 的效率相当高，模型 FLOPS 利用率最高可达约 60%；使用一个配备 8 张 A100 80GB SXM 的节点复现该模型，大约需要 90 分钟。例如，在 Lambda 上，这种节点的价格约为 14 美元/小时，因此如今复现该模型的总成本约为 20 美元。你也可以只用一张 GPU 训练模型，只是耗时会按比例增加（例如，根据 GPU 的不同，大约需要 4～24 小时）。此外，llm.c 仍有许多待完成的优化，而且人们还没有尝试以 cramming 的方式调优训练，因此我认为这个数字很可能还会得到显著改善。下面就是这次训练：使用 100 亿个 FineWeb token，训练一个 12 层、12 个注意力头、768 维、124M 参数的 Transformer：

左图显示，在 FineWeb 留出的验证数据集上，我们的表现超过了 OpenAI 发布的 checkpoint。这并不是理想的指标，因为 GPT-2 的数据分布并不相同（它是在从未公开发布的「WebText」数据集上训练的），而且 5 年前互联网的统计特征可能也有所不同，因此这并不是一次特别公平的比较。因此，我们还在右图绘制了 HellaSwag 准确率。HellaSwag 是一个常用于评估 LLM 能力的基准，其曲线良好、平滑而且表现稳定。我主要会关注 HellaSwag，不过 FineWeb 验证集也是一个不错的佐证。话虽如此，HellaSwag 不包含数学或代码内容，因此它会略微偏向我们的训练设置（类似 Common Crawl 的数据）。另一个参考点是，GPT-3 论文附录 H 中给出的 GPT-3 Small（124M）模型 HellaSwag 准确率为 33.7。我们在这里达到了 29.9，超过了 GPT-2（124M）的 29.4。请记住，我们这里只训练了 100 亿个 token，而所有 GPT-3 模型都训练了 3,000 亿个 token。

下面是你亲自复现这个结果的最短路径。你需要一张 GPU。我喜欢在 Lambda Labs 上运行自己的工作（他们慷慨赞助了 llm.c 的开发），不过其可用实例有时会比较有限。市面上还有许多其他提供商，你可以在下方的讨论区交流相关技巧。以下流程以 Linux x86 64 位 Ubuntu 22.04 和 CUDA 12 为例（这大致是目前默认的「现代」配置）。如果你使用的是其他系统，主 README 文件中的注释和讨论可能会有所帮助。

# install miniconda
mkdir -p ~/miniconda3
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
rm -rf ~/miniconda3/miniconda.sh
~/miniconda3/bin/conda init bash
source ~/.bashrc

# pytorch nightly (optional) https://pytorch.org/get-started/locally/
# conda install --yes pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia

# pip installs so we can tokenize the FineWeb dataset
yes | pip install tqdm tiktoken requests datasets

# install cudnn so we can use FlashAttention and run fast (optional)
# https://developer.nvidia.com/cudnn-downloads
# for me, CUDA 12 (run `nvcc --version`) running on Linux x86_64 Ubuntu 22.04
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install libcudnn9-dev-cuda-12

# "install" cudnn-frontend to ~/
git clone https://github.com/NVIDIA/cudnn-frontend.git

# install MPI (optional, if you intend to use multiple GPUs)
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev

# tokenize the FineWeb dataset 10B tokens sample (takes ~1 hour, get lunch?)
# writes ~19GB of raw GPT-2 tokens to dev/data/fineweb10B
# and ~46GB in ~/.cache/huggingface/datasets/HuggingFaceFW___fineweb
git clone https://github.com/karpathy/llm.c.git
cd llm.c
python dev/data/fineweb.py --version 10B

# compile llm.c (mixed precision, with cuDNN flash-attention)
# first compilation is ~1 minute, mostly due to cuDNN
make train_gpt2cu USE_CUDNN=1

# train on a single GPU
./train_gpt2cu \
    -i "dev/data/fineweb10B/fineweb_train_*.bin" \
    -j "dev/data/fineweb10B/fineweb_val_*.bin" \
    -o log124M \
    -e "d12" \
    -b 64 -t 1024 \
    -d 524288 \
    -r 1 \
    -z 1 \
    -c 0.1 \
    -l 0.0006 \
    -q 0.0 \
    -u 700 \
    -n 5000 \
    -v 250 -s 20000 \
    -h 1

# if you have multiple GPUs (e.g. 8), simply prepend the mpi command, e.g.:
# mpirun -np 8 ./train_gpt2cu \ ... (the rest of the args are same)

参数指南。这里的许多超参数遵循 GPT-3 论文，而不是 GPT-2 论文，因为 GPT-3 论文中的细节要丰富得多。参数说明：

-i 和 -j 分别是训练集与验证集的 token 文件，由 fineweb.py 写入。

-o 是用于写入日志和 checkpoint 的输出目录。

-e "d12" 表示从头初始化一个深度为 12 层的 GPT-2 模型。

-b 64 将微批次大小设置为 64。如果内存不足，请减小该值，例如依次尝试 32、16、8，必要时一直降到 1。

-t 1024 将最大序列长度设置为 1024，与 GPT-2 相同。

-d 524288 要求每次更新的总批次大小约为 50 万个 token。代码会获取这个期望的批次大小，并计算优化过程中所需的梯度累积「内循环」步数。例如，在 8 张 GPU 上，当 -b 64 且 -t 1024 时，每个微批次处理的 token 数恰好为 8 X 64 X 1024 = 524288，因此无需进行梯度累积。但如果只有 1 张 GPU，代码就会将梯度累积步数设置为 8，并执行 8 次内循环迭代，从而使每一步累积到这个「总批次大小」。虽然 GPT-2 训练时使用的批次大小未知，但对于该模型规模，约 50 万这个数字来自 GPT-3 论文中的表格。

-r 1 将重计算设置设为 1，因此我们会重新计算 GeLU 激活值。这会略微增加运行时间，但能够节省大量内存，使我们可以增大批次大小，并最终提升 token 吞吐量。

-z 1 会在多张 GPU 之间启用 ZeRO-1（即优化器状态分片）。如果你使用超过 1 张 GPU 进行训练，这个设置无需多想，基本上应该始终开启。在 1 张 GPU 上，该设置不起作用。

-c 0.1 将权重衰减设置为 0.1。与 GPT-2 完全相同，只有（二维）权重会进行衰减，而这个数值来自 GPT-3 论文。

-l 0.0006 设置最大学习率，该值来自 GPT-3 论文。

-q 0.0 表示在整个训练过程中，我们会将学习率衰减到 0。

-u 700 表示在前 700 次迭代中，将学习率从 0 逐渐提升到最大学习率。按照约 50 万的总批次大小计算，这相当于 3.5 亿个 token，遵循 GPT-3 论文的设置。

-n 5000 表示每隔 5,000 步保存一次模型 checkpoint。

-v 250 表示每隔 250 步评估并记录一次验证损失。

-s 20000 表示每隔 20,000 步采样一些 token。由于总步数会小于该值（见下文），这实际上相当于关闭生成，我们基本上只会在最后采样一次。

-h 1 表示评估 HellaSwag 准确率，以便与不同论文中的结果进行比较。

由于我们没有使用 -x 标志设置最大步数，因此默认会完整遍历训练数据一次，即训练一个 epoch，也就是 100 亿个 token。因为总批次大小约为 50 万，而 token 总数为 100 亿，所以总步数约为 10B/0.5M = 20K 步。

上面的细节很多，但简而言之，我们正在使用 100 亿个 FineWeb token，从头训练一个 12 层的 GPT-2（124M），最大序列长度为 1,024 个 token。如果内存不足，我会先确认已经开启 -r 1，然后开始将批次大小 -b 不断除以 2，直到能够正常运行。成功运行后，我会再看看是否可以重新开启 -r 0，以恢复一点速度。

训练。随着时间推移，代码会输出类似下面的内容（这里以一张 A100 40GB PCIe GPU 为例，价格为 1.29 美元/小时）：

step   80/18865 | train loss 7.577051 | norm 1.1461 | lr 6.86e-05 | 2950.68 ms | 49.0% A100 fp16 MFU | 177968 tok/s
step   81/18865 | train loss 7.540626 | norm 1.4001 | lr 6.94e-05 | 2952.59 ms | 49.0% A100 fp16 MFU | 177948 tok/s
step   82/18865 | train loss 7.465753 | norm 1.0613 | lr 7.03e-05 | 2953.98 ms | 48.9% A100 fp16 MFU | 177924 tok/s
step   83/18865 | train loss 7.472681 | norm 1.1553 | lr 7.11e-05 | 2955.67 ms | 48.9% A100 fp16 MFU | 177897 tok/s

发生了什么？我们有 10B 个训练 token，批大小约为 0.5M，因此我们预计总共会有约 10B/0.5M ~= 20K 步。实际结果正好是 18,865 步，因为其中一个数据分片被保留用作验证数据，确切的批大小是一个漂亮的 2 的幂 @ 524,288。所以我们现在处于步骤 80/18865，总耗时 2950.68ms。MFU 是"模型浮点运算利用率"的缩写。A100 声称可提供 312 TFLOPS，但实际上这很难实现，因为训练受内存限制，我们无法充分喂饱执行矩阵乘法的 TensorCore。在这个 A100 40GB PCIe GPU 上，当我们计算执行的 FLOPs 并除以时间时，我们粗略达到理论峰值 FLOPS 的一半，这相当不错。如果使用内存带宽更高、热设计功耗更大的 A100 80GB SXM，这个比例会上升到约 60%。（如果你使用的不是 A100，请忽略这个数字，因为它的单位是 A100 fp16 FLOPS）。我们还看到实现的 token 吞吐量约为 178K tok/s。接下来，当前损失为 7.577。这个值越低，我们的模型在平均预测序列中下一个 token 时的表现越好。步骤 80 是这里训练的很早阶段。因为困惑度是 exp(7.577) ~= 2K，我们的模型对每个下一个 token 的困惑程度，相当于从 2,000 个 token 中随机猜测。完整词汇表大小为 50,257。优化结束时我们会达到约 3.29，就好像在每个时间步长时从 exp(3.29) ~= 27 个 token 中均匀随机猜测。最后我们看到梯度范数是 1.1461。当这个数字激增时，梯度爆炸了，这很糟糕。为了缓解梯度爆炸，如惯例所示，llm.c 使用梯度裁剪，阈值为 1.0，因此如果梯度范数超过 1.0（就像这个时间步长中一样），我们会强制缩放它，使其范数达到 1.0。优化后期，梯度范数通常会"平复"到较低的值。

可视化。最后，你会想要制作漂亮的图表，就像我上面发布的那样。为此，我们的程序将一些非常基础的日志打印到一个即兴的 log124M/main.log 文件中。我附上了一个示例 Jupyter notebook，它解析这些文件并按上述风格可视化它们。

分词器。当你在上面进行训练时，你会看到一个警告，说 llm.c 找不到 GPT-2 分词器 .bin 文件。这对于训练来说完全没问题，但这意味着我们无法解码——即我们无法将我们采样的整数 token 转换成小的字符串片段，来生成我们能读的文本。以下是我们如何生成它的方法：

# 安装 pytorch nightly
conda install --yes pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia

# 安装 huggingface transformers
pip install transformers

# 预处理 TinyShakespeare 数据集（非常快，比 FineWeb 快得多）
python dev/data/tinyshakespeare.py

# 在 Python/PyTorch 中运行一个小的训练循环
# 它会保存许多 .bin 文件，包括分词器
python train_gpt2.py

Python 脚本是 llm.c 的平行实现，用于错误检查和单元测试（但没有完全的功能奇偶性）。特别是，如果我们像上面那样运行它，它将写入文件 gpt2_tokenizer.bin，C 代码可以读取并使用它来在采样期间输出漂亮的文本。

采样。该代码目前并不是真正为推理而设计的，但你可以通过类似下面的方式hack该代码来非常低效地进行推理（没有任何 kv-cache 等）：

make train_gpt2cu USE_CUDNN=1
./train_gpt2cu \
    -i "dev/data/fineweb10B/fineweb_train_*.bin" \
    -j "dev/data/fineweb10B/fineweb_val_*.bin" \
    -e "log124M/gpt2_124M_00018865.bin" \
    -b 1 -t 1024 \
    -x 1 \
    -l 0.0 \
    -s 1 -g 256

-i -j 标志是多余的。-e 标志指向我们 GPT-2 124M 模型的最终检查点，llm.c 将从中初始化模型。-b 1 表示仅使用单个批元素（长度为 1024 个 token 的一行，我们从左到右采样）。-x 1 表示我们仅想运行单个步骤，-l 0.0 将学习率设置为零，因此我们实际上不会在这个单步上训练模型。最后 -s 1 表示"每步采样"，-g 256 表示采样 256 个 token。

现在，上面只是无条件采样。可以hack该代码来进行条件采样，即序列补全。例如，我要求我们的 124M 模型完成文本"The GitHub project llm.c is a"，它继续说："free service to enhance the scholarly infrastructure of the academic community."。然后我用不同的种子重新采样，得到"The GitHub project llm.c is a collaborative effort that rocks GitHub itself"。所以，还不错吧 :) 我必须通过将 gen_tokens[1:10] 设置为提示 token 464、21722、1628、32660、76、13、66、318、257（来自 tiktokenizer 谢了）来直接hack代码，然后hack采样的循环索引以从 token 位置 10 开始，...你明白了吧。总之，条件生成目前并未真正支持，但原则上是可能的，可能即将推出。

代码。95% 的重活在 train_gpt2.cu 文件中。它最初是一份漂亮的 1,000 行 C 代码，但已经增长了不少，现在接近 3,500 行，另外还有 4 个文件用于文件 I/O 工具、分词器、数据加载器和随机数生成。粗略来说，前 500 行只是 MPI、NCCL、cuDNN、cuBLAS 等的基本设置。接下来的 1,500 行是 Transformer 的所有层及其在高效 CUDA 代码中的前向和反向实现。所有这些文件的 CUDA 内核开发都在 dev/cuda 中进行。例如，有一个 gelu_forward()，然后还有一个 gelu_backward()，所有其他层也是如此。接下来的 1,000 行是 gpt2 模型，它只是将各层串联起来，本身有一个大的 gpt2_forward() 和 gpt2_backward()。最后的 1,000 行是 int main()，它包含主训练循环和所有相关的簿记与参数解析，以及大量围绕例如从之前的检查点恢复训练等的繁琐代码。

350M 模型。一夜之间我还复现了 350M 参数模型。查看文件 scripts/run_gpt2_350M.sh 了解确切的启动命令。我发现 10B token 对于 350M 模型来说不够，所以你必须下载并预处理 FineWeb100B（或尝试在上面的 10B 上进行多个 epoch，可能会有效，我还没有检查）。我将其配置为训练 30B token，所以我们有：

使用 6ND 近似计算 FLOPS：

124M on 10B tokens => 6 * 124e6 * 10e9 = 7.44e18 ~= 7e18 能力模型

350M on 30B tokens => 6 * 350e6 * 31.5e9 = 6.615e19 ~= 7e19 能力模型 (~10X)

在 8X A100 80GB SXM 上，350M 的步长为 820ms/iter。训练了 60K 步（而不是约 20K 步），总共约 30B token（而不是约 10B token）。总训练时间 14 小时。成本 $14/hr => 14 X 14 ~= $200（124M 的 10 倍）。但是看图，似乎我们可能能够少花一点钱：

即将推出。目前就是这样！我们继续进行 740M，然后当然是实际的"GPT-2"1558M。如果我能找到 GPU...根据非常粗略的粗算，在我的单个 8X A100 80GB GPU 机器上，1558M 模型将耗时约 1 周，成本约 $2.5K。这在可接受的范围内，但我们会想花一些时间来改进当前代码，使其更干净、测试更充分，并添加多节点训练支持。还有，我非常想要重新构建整个东西，从头开始逐块，即将推出给你^TM。

我能从中采样吗？有点，但它效率低下且有点奇怪。

我能和它聊天吗？不能，这目前仅是预训练，不是聊天微调。

你能进行多节点分布式训练吗？原则上能，有一个 slurm PR 已经这样做了，最多支持 50 个节点。实际上我个人还没有尝试过。

你是逐位确定的吗？不是，但我们非常接近，还需要修补一个内核。

你能用 fp8 训练吗？不能，我们目前主要是用 bf16 训练，但即将推出。

我有一个非 NVIDIA GPU（AMD、Apple Silicon 等），我能运行 llm.c 吗？不能，llm.c 仅支持 C/CUDA，但我非常乐意在"notable forks"部分链接任何 fork，或接受能让将 llm.c 移植到其他平台变得更容易的 PR。

我只有一个 CPU，我能玩吗？你将无法复现 GPT-2 模型，但你可以通过在其他数据上微调 OpenAI GPT-2 模型来完成有趣的项目，例如 TinyShakespeare 或 TinyStories。llm.c 中 train_gpt2.c 中存在对这些数据集、初始化和 CPU 微调的支持。（不过它粗糙得多，主要是作为 CUDA 代码的参考）。

这与 PyTorch 相比如何？

llm.c 是一个"直接的" C/CUDA 实现。train_gpt2.py 中的 PyTorch 代码没有完整的功能对等（例如，不进行分片数据加载等），其目的更多是作为参考实现，但我认为你可以通过以下步骤得到类似上面 124M 模型的结果：

torchrun --standalone --nproc_per_node=4 python train_gpt2.py --input_bin dev/data/fineweb10B/fineweb_train_000001.bin --write_tensors 0 --model d12 --batch_size 64 --sequence_length 1024 --total_batch_size 524288 --dtype bfloat16 --compile 1 --tensorcores 1 --flash 1 --num_iterations 18865 --weight_decay 0.1 --overfit_single_batch 0

我感兴趣并会接受 PR，将 PyTorch 训练提升到与 llm.c 训练循环的功能对等水平。

为什么你这么关注 GPT-2？

GPT-2 是 LLM 的祖父级模型，是现代 LLM 栈以可识别的现代形式首次聚集的时刻，并且参数由 OpenAI 发布。GPT-3 实际上对模型改动不大（上下文大小 1024 -> 2048，我想大概就是这样？）。GPT-4 的细节从未发布过。许多其他 LLM 也强烈类似于 GPT-2，尽管它来自 2019 年，例如从架构角度看 Llama 3 是 MLP 中非线性的改变以及 RoPE 相对位置编码的添加。

致谢

向 @ngc92 和 @ademeure 致敬，他们在 llm.c 的各个方面都做出了实质性贡献，特别是在 CUDA 内核优化方面，@chinthysl 和 @PeterZhizhin 在分布式优化 PR 上，以及 @rosslwheeler 在 Windows 支持和工具方面。

欢迎使用讨论区来获取任何常见问题和相关内容，或者如果你想要更快的响应，可以去 Discord 上的 #llmc 或 CUDA MODE Discord 上的 #llmdotc。

回复： 48 条评论 · 83 条回复

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

karpathy 2024 年 5 月 28 日 维护者 作者

我在 HN 线程中回答了一些问题

我在 X 线程中回答了一些问题

加载时出错。请刷新此页面。

Niskarsh12 2024 年 5 月 28 日

我想试试，但遗憾的是我没有 GPU :/

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

如果其他人也不想花那么多钱，可以尝试 vast.ai，那里同样的机器现在只需要 5$/小时（只需记住它可能不如云计算那样可靠，甚至可能不可用，因为这是一个市场）

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

把你本来要花在云上的钱省下来，直到你能买得起一个 GPU 主机。

这是一个巨大的机会成本。你的时间很宝贵。开始学习这些东西的最佳时间是 10 年前。第二好的时间就是今天（而不是几个月后，等你攒了 10 万多美元买一个 8xA100 SXM 系统）。

出租 GPU 对爱好者来说太棒了。

加载时出错。请刷新此页面。

我同意 thunder golf 的观点，大多数人在周五的饮料上花 100 美元。

如果你的想法足够好，开始构建，GPU 会找到你

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

感谢 @karpathy 在你的 GitHub 仓库中提供的宝贵教学课程。

我克隆了 llm.c 来查看你如何进行 dropout。

我找到了一些在 NVIDIA CUDA GPU 设备上运行的随机数生成函数。

Dropout 在哪里进行？

加载时出错。请刷新此页面。

现在还没有 Dropout。添加弱 Dropout（例如 0.05）可能会好一点，但它在具有训练和评估模式以及处理所有这些的复杂性上引入了一层复杂性，此时处理这些太麻烦了。我们在足够大的数据集上训练非常小的模型，所以我认为对正则化的需求较少。它可能仍然会有帮助。

加载时出错。请刷新此页面。

好的，谢谢你 @karpathy ！

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

我在 https://github.com/karpathy/llm.c/blob/master/train_gpt2.cu#L425 中看到了 MPI_Allgather

如果所有 8 个 A100 80GB SXM GPU 都在同一节点上，为什么要使用 MPI_Allgather？

加载时出错。请刷新此页面。

我们希望立即支持多节点训练。

加载时出错。请刷新此页面。

哦，这太令人兴奋了！

加载时出错。请刷新此页面。

👍 已经有一个 PR 提出来了（在上面的文章中链接），但我还没有抽出时间来处理它。

加载时出错。请刷新此页面。

YuchenJin 2024 年 5 月 28 日

有几个地方应该将 train_gpt2cu 改为 train_gpt2.cu。

加载时出错。请刷新此页面。

@ngc92 那太棒了！

加载时出错。请刷新此页面。

@YuchenJin 你能检查一下当我们在开始时打印出设备名称时 H100 会得到什么名称吗？

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

那么理论上可以在 mps 上做吗？

只是略微开玩笑，但你认为在 m4 ultra 192gb 内存上运行需要多长时间？

加载时出错。请刷新此页面。

是的，我得到了同样的结果。

加载时出错。请刷新此页面。

加载时出错。请刷新此页面。

banyan-god 2024 年 5 月 28 日

这是一个我最近几天正在训练的模型（500M），使用 llama2 架构。希望训练到大约 2000 亿 tokens。这使用的是 fineweb 2024 和 4x 4090 https://wandb.ai/banyan-t/llamac/runs/zjaods8q n_heads:36 n_kv_heads:36 n_layers:36 dim: 970

加载时出错。请刷新此页面。

4090 真的在发挥威力。@karpathy 想知道是否有办法让我使用所有内存。GPU 跑得飞快，我希望能运行批大小大于 32 的。64 会导致内存不足 step 22/18865 | train loss 9.542500 | norm 2.1915 | lr 1.89e-05 | 902.14 ms | 40.1% A100 fp16 MFU | 581589 tok/s step 23/18865 | train loss 9.585844 | norm 2.0341 | lr 1.97e-05 | 902.65 ms | 40.0% A100 fp16 MFU | 581534 tok/s step 24/18865 | train loss 9.502515 | norm 2.0734 | lr 2.06e-05 | 901.97 ms | 40.1% A100 fp16 MFU | 581514 tok/s

加载时出错。请刷新此页面。

bprimal22