Reproducing GPT-2 (124M) in llm.c in 90 minutes for $20（粗略翻译）

https://github.com/karpathy/llm.c/discussions/481

让我们用 llm.c（约 4000 行 C/CUDA 代码）在 90 分钟内以 20 美元的成本复现 GPT-2（124M）。124M 模型是 OpenAI 在 2019 年发布的 GPT-2 系列中最小的模型，实际上对于现在的硬件来说是相当可行的，即使是对 GPU 资源有限的人来说也是如此。使用 llm.c，可以实现高达约 60% 模型浮点运算效率，使用一个 8X A100 80GB SXM 节点大约需要 90 分钟来复现这个模型。例如，在 Lambda 平台上，这个节点的费用约为 14 美元/小时，因此今天复现这个模型的总成本大约是 20 美元。你也可以用单个 GPU 来训练这个模型，只是需要相应更长的时间（例如，根据 GPU 的不同，大约需要 4 到 24 小时）。此外，llm.c 仍有许多待优化的部分，而且人们还没有尝试按照紧凑训练的方式来调整训练，所以我认为在这个时间上我们有望看到显著的改进。因此，这是一次运行，在 10 亿个 FineWeb 语料的训练数据上训练一个具有 12 层、12 个头、768 维度、124M 参数的 Transformer 模型：

左侧的图表显示，我们在 FineWeb 保留的验证数据集上超越了 OpenAI 发布的 GPT-2 检查点。这并不是一个理想的衡量标准，因为 GPT-2 的数据分布不同（它是在从未公开的 "WebText" 数据集上训练的），而且互联网上的数据统计在 5 年前可能也有所不同，所以这并不是一个非常公平的比较。因此，在右侧的图表中，我们还绘制了 HellaSwag 准确率，这是一个常用于评估大型语言模型（LLM）能力的基准，非常稳定且表现良好。我更倾向于关注 HellaSwag，但 FineWeb 验证集也是一个不错的确认。需要注意的是，HellaSwag 不包含数学或代码相关内容，因此它稍微偏向于我们的设置（类似 common crawl 的数据）。另外一个参考点是 GPT-3 的附录 H，其中引用的 HellaSwag 准确率为 GPT-3 Small（124M）模型的 33.7。我们在这里达到了 29.9，超过了 GPT-2（124M）的 29.4。请注意，我们在这里训练了 100 亿个 tokens，而 GPT-3 模型都是在 3000 亿个 tokens 上训练的。

现在，这里是你自己复现这个结果的最短路径。你需要一块 GPU。我个人喜欢并在 Lambda 实验室上进行工作（他们慷慨地赞助了 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 是输出目录，用于写入日志和检查点。
• -e "d12" 表示从头开始初始化一个深度为 12 的 GPT-2 模型。
• -b 64 将微批次大小设置为 64。如果你的内存不够，可以降低这个值，比如尝试 32、16、8，一直到 1。
• -t 1024 将最大序列长度设置为 1024，与 GPT-2 一样。
• -d 524288 请求单次更新的总批次大小为约 0.5M tokens。代码将根据所需的批次大小计算优化器所需的梯度累积“内循环”步骤。例如，在 8 个 GPU 上，设置 -b 64 和 -t 1024，每个微批次恰好处理 8 X 64 X 1024 = 524288 个 tokens，因此不需要进行梯度累积。但如果只有 1 个 GPU，代码会将其设置为 8，并执行 8 次内循环，以达到每步的这个“总批次大小”。虽然用于训练 GPT-2 的批次大小未知，但这个约 0.5M 的数值来自 GPT-3 论文表格中针对该模型大小的建议。
• -r 1 将重计算设置为 1，因此我们将重新计算 GeLU 激活函数。这样做稍微增加了运行时间，但节省了不少内存，从而允许我们增加批次大小，进而提高 token 吞吐量。
• -z 1 启用 ZeRO-1（即优化器状态分片）在多个 GPU 间分配。如果你用超过 1 个 GPU 进行训练，这个设置是必须的，基本上应该始终启用。在单个 GPU 上，这个设置无效。
• -c 0.1 将权重衰减设置为 0.1。只有（二维）权重会被衰减，这与 GPT-2 的方法完全一致，这个值来源于 GPT-3 论文。
• -l 0.0006 将最大学习率设置为 0.0006，这也来源于 GPT-3 论文。
• -q 0.0 表示在整个训练过程中，我们将学习率衰减到 0。
• -u 700 表示在前 700 次迭代中，我们将学习率从 0 提升到最大学习率。对于总批次大小 0.5M，这相当于前 3.5 亿个 tokens，遵循了 GPT-3 论文的设置。
• -n 5000 每 5000 步保存一次模型检查点。
• -v 250 每 250 步评估并记录验证损失。
• -s 20000 每 20000 步生成一些 tokens。由于总步数会少于这个值（见下文），这基本上关闭了生成功能，最终我们只会在最后生成一次。
• -h 1 要求评估 HellaSwag 准确率，这是我们可以在论文间进行比较的指标。

由于我们没有使用 -x 参数设置最大步数，因此它默认为对训练数据进行整整一个周期，即 100 亿个 tokens。由于总批次大小约为 0.5M 和总 token 数为 100 亿，总步数将为约 10B / 0.5M = 20000 步。

上面的内容信息量很大，但总结来说，我们从零开始训练一个 12 层的 GPT-2 模型（124M 参数），使用 100 亿个 FineWeb tokens，最大序列长度为 1024 个 tokens。如果你遇到内存不足的问题，首先确保 -r 1 已开启，然后开始将批次大小 -b 减半，直到可以正常运行。一旦可以运行，你可以尝试将 -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

这段代码在做什么呢？我们有 100 亿个训练 tokens，批次大小约为 0.5M，所以预计总共会有大约 10B / 0.5M ≈ 2 万步（steps）。实际上，总步数恰好是 18865，因为其中一个数据分片是保留用于验证数据的，并且精确的批次大小是一个 2 的幂（524288）。因此，这里显示的是我们处于第 80 步/18865 步，总耗时 2950.68 毫秒。

MFU 是 "Model Flops Utilization"（模型浮点运算利用率）的缩写。A100 声称提供 312 TFLOPS，但实际上很难实现这个水平，因为训练是内存受限的，我们无法给执行矩阵乘法的 TensorCores 提供足够的数据。在这个 A100 40GB PCIe GPU 上，当我们计算执行的 FLOPs 并除以时间时，发现我们大约达到了理论最大 FLOPS 的一半，这已经相当不错了。如果使用内存带宽更高、最大热设计功率更大的 A100 80GB SXM，这个比例可以提升到大约 60%。（如果你使用的 GPU 不是 A100，请忽略这个数字，因为它是以 A100 fp16 FLOPS 为单位的。）

我们还看到，我们目前实现的 token 吞吐量约为 178K tok/s。接下来，我们的当前损失是 7.577。损失越低，模型在平均上预测序列中下一个 token 的能力就越好。第 80 步还处于训练的早期。因为困惑度（perplexity）是 exp(7.577) ≈ 2000，我们的模型在预测每个下一个 token 时的困惑程度，相当于从 2000 个 token 中随机猜测。完整词汇表的大小是 50257。在优化结束时，我们的困惑度会降到大约 3.29，也就是说，在每个时间步预测下一个 token 时，相当于从 exp(3.29) ≈ 27 个 token 中随机猜测。

最后，我们看到梯度范数（gradient norm）是 1.1461。当这个数字激增时，表示梯度爆炸，这是非常不好的。为了缓解梯度爆炸，llm.c 使用了标准的梯度裁剪（gradient clipping）策略，将其裁剪到 1.0。所以如果梯度范数超过 1.0（比如在这个时间步），我们会强制将其缩小，使其范数最多为 1.0。优化的后期，梯度范数通常会“平静”下来，达到较低的值。

可视化。最后，你可能希望像我上面发布的那样制作漂亮的图表。为此，我们的程序会将一些非常基础的日志输出到一个自制的 log124M/main.log 文件中。我已经附上了一个 Jupyter notebook 示例，它可以解析这些文件，并以上述风格进行可视化。

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

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

# install huggingface transformers
pip install transformers

# preprocess the TinyShakespeare dataset (very fast, much faster than FineWeb)
python dev/data/tinyshakespeare.py

# run a little training loop in Python/PyTorch
# it saved a lot of .bin files, including the Tokenizer
python train_gpt2.py

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

采样（Sampling）。 当前代码主要不是为了推理而设计的，但你可以通过一些小改动让代码执行推理，尽管这样效率很低（没有任何键-值缓存等），实现方法类似如下：

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 个 tokens 的一行，我们从左到右进行采样）。-x 1 表示我们只想运行一步，-l 0.0 将学习率设置为零，这样我们实际上不会在这一步训练模型。最后，-s 1 表示“每一步都采样”，-g 256 表示采样 256 个 tokens。

上面只是无条件采样。我们可以通过修改代码来实现条件采样，即序列补全。例如，我让我们的 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] 设置为提示 tokens 464, 21722, 1628, 32660, 76, 13, 66, 318, 257（使用 tiktokenizer），然后将采样的循环索引修改为从 token 位置 10 开始，等等。简而言之，条件生成目前尚不支持，但理论上是可能的，可能很快就会实现。

代码结构。大部分工作都在 train_gpt2.cu 文件中完成。最初，这个文件是一段 1000 行干净的 C 代码，但现在它已经增长到接近 3500 行代码，还有 4 个辅助文件用于文件 I/O 实用程序、分词器、数据加载器和随机数生成器。大致上，前 500 行代码是 MPI、NCCL、cuDNN、cuBLAS 等的基本设置。接下来的 1500 行是 Transformer 的所有层，以及它们的前向和后向实现的高效 CUDA 代码。所有这些文件的 CUDA 内核开发都发生在 dev/cuda 中。例如，有一个 gelu_forward() 函数和一个相应的 gelu_backward() 函数，其他层的实现方式相同。接下来的 1000 行代码是 GPT-2 模型，它只是将这些层串联起来，并且自身有一个大的 gpt2_forward() 和 gpt2_backward() 函数。最后的 1000 行代码是 int main() 函数，其中包含主训练循环以及所有相关的记录和参数解析代码，以及大量关于从以前的检查点恢复训练等的繁琐代码。

350M 模型。在夜间，我还复现了 350M 参数的模型。可以查看文件 scripts/run_gpt2_350M.sh 获取精确的启动命令。我发现 100 亿 tokens 对于 350M 模型来说不够，所以你需要下载和预处理 FineWeb100B（或者尝试对上面的 10B 数据进行多轮训练，这可能有效，但我没有测试过）。我将其配置为训练 300 亿 tokens，因此我们得到：

使用 6ND 近似法计算 FLOPS：

• 124M 模型在 100 亿 tokens 上的计算：6 * 124e6 * 10e9 = 7.44e18 ≈ 7e18 FLOPS
• 350M 模型在 300 亿 tokens 上的计算：6 * 350e6 * 31.5e9 = 6.615e19 ≈ 7e19 FLOPS（约 10 倍）

在 8 个 A100 80GB SXM 上，350M 模型的每步用时 820ms/iter。训练了 6 万步（而不是 ~2 万步），总共训练了约 300 亿 tokens（而不是 ~100 亿 tokens）。总训练时间 14 小时。费用为 14 美元/小时 => 14 x 14 ≈ 200 美元（是 124M 模型的 10 倍）。不过，从图表上看，我们可能可以稍微减少一些训练时间：

接下来计划
暂时就到这里！我们将继续进行 740M 模型，然后当然是实际的“GPT-2” 1558M 模型。如果我能找到足够的 GPU......通过非常粗略的估算，在我单个 8x A100 80GB GPU 设备上，1558M 模型大约需要 1 周时间，花费大约 2500 美元。这在可接受的范围内，但我们还需要花些时间让现有代码变得更好、更干净、测试更完善，并添加多节点训练支持。另外，我还计划从头开始重新构建整个项目，一步步来，很快就会和大家见面了^TM。

常见问题解答：

• 我可以从中采样吗？ 可以，但是效率低下且有点奇怪。
• 我可以和它聊天吗？ 不行，目前只支持预训练，不支持聊天微调。
• 可以进行多节点分布式训练吗？ 理论上可以，有一个 slurm 的 PR 已经支持最多 50 个节点的训练。但实际上我还没有亲自尝试过。
• 你是按位确定性的吗？ 不是，但我们非常接近，只剩下一个内核需要修补。
• 可以用 fp8 训练吗？ 不能，我们目前主要用 bf16 进行训练，但即将支持 fp8。
• 我有一个非 NVIDIA 的 GPU（例如 AMD、Apple Silicon 等），可以运行 llm.c 吗？ 不能，llm.c 仅支持 C/CUDA，但我很乐意在“notable forks”部分链接到任何分支，或者接受让移植 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。我也很感兴趣，并会接受将 PyTorch 训练更接近 llm.c 训练循环功能的 PR。
• 为什么你这么关心 GPT-2？ GPT-2 是 LLMs 的始祖，是现代 LLM 堆栈首次以现代形式出现，参数由 OpenAI 发布。GPT-3 实际上并没有对模型进行太多更改（上下文大小从 1024 -> 2048，我想就是这样）。GPT-4 的详细信息从未公布。尽管 GPT-2 是在 2019 年发布的，许多其他 LLMs 也非常类似于 GPT-2，例如从架构角度看，Llama 3 是在 MLP 中进行了一些非线性更改，并添加了 RoPE 相对位置编码。

致谢
感谢 @ngc92 和 @ademeure 为 llm.c 在 CUDA 内核优化方面及其他领域做出的重大贡献，感谢 @chinthysl 和 @PeterZhizhin 提供的分布式优化 PR，感谢 @rosslwheeler 提供的 Windows 支持和工具。

如有任何常见问题和相关问题，请随时使用 Discussions 区，如果您希望得到更快速的反馈，请加入 Discord 的 #llmc 频道或 CUDA MODE Discord 的 #llmdotc 频道。