llm.c：用纯 C/CUDA 从零训练 LLM

项目简介

用纯 C/CUDA 实现 LLM，无需 245MB 的 PyTorch 或 107MB 的 cPython。当前重点是预训练，特别是复现 GPT-2 和 GPT-3 系列，同时在 train_gpt2.py 中提供了平行的 PyTorch 参考实现。你会认出这个文件是我之前项目 nanoGPT 的稍微改进版本。目前，llm.c 的性能比 PyTorch Nightly 快约 7%。除了主干代码 train_gpt2.cu 之外，我们还有一个简单的参考 CPU fp32 实现，约 1000 行干净代码，全部在一个文件 train_gpt2.c 中。我希望这个仓库仅维护 C 和 CUDA 代码。欢迎移植到其他语言或仓库，但应该在独立的仓库中进行，我很乐意在下方的"notable forks"部分链接它们。开发者协调在 Discussions 和 Discord 上进行，可以是 Zero to Hero 频道的 #llmc，或 GPU MODE Discord 的 #llmdotc。

学习 llm.c 仓库最好的方法是复现 GPT-2（124M）模型。Discussion #481 详细介绍了这一点。我们可以在 llm.c 和平行的 PyTorch 实现中复现 GPT-2 和 GPT-3 系列的其他模型。看看 scripts README。

调试提示：当你运行 make 命令构建二进制文件时，修改它，用 -g 替换 -O3，这样你可以在你喜欢的 IDE（例如 vscode）中逐步执行代码。

快速开始（1 GPU，fp32 仅）

如果你不会在多节点上训练，对混合精度不感兴趣，又对学习 CUDA 感兴趣，那么 fp32（legacy）文件可能对你有兴趣。这些文件是在 llm.c 历史早期"检查点"并冻结在时间中的。它们更简单、更可移植，可能更容易理解。像这样运行 1 GPU、fp32 代码：

chmod u+x ./dev/download_starter_pack.sh
./dev/download_starter_pack.sh
make train_gpt2fp32cu
./train_gpt2fp32cu

download_starter_pack.sh 脚本是一个快速简便的开始方式，它下载一堆 .bin 文件来帮助你上手。这些包含：1) 以 fp32 和 bfloat16 保存的 GPT-2 124M 模型，2) 一个"调试状态"用于单元测试（一小批数据、目标激活和梯度），3) GPT-2 分词器，以及 4) 分词化的 tinyshakespeare 数据集。或者，你可以手动重新创建这些工件，而不是运行 .sh 脚本：

pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py

"我是如此 GPU 贫困，甚至都没有一个 GPU" 部分

你仍然可以看到 llm.c 训练！但你不会走太远。就像上面的 fp32 版本一样，CPU 版本是 llm.c 历史中更早的一个检查点，回到它只是一个简单的 C 参考实现的时代。例如，你可以微调 GPT-2 小版本（124M）来输出莎士比亚风格的文本，作为一个例子：

chmod u+x ./dev/download_starter_pack.sh
./dev/download_starter_pack.sh
make train_gpt2
OMP_NUM_THREADS=8 ./train_gpt2

如果你更想避免运行 starter pack 脚本，那么如前一部分所述，你可以通过运行 python dev/data/tinyshakespeare.py 然后 python train_gpt2.py 来重新生成完全相同的 .bin 文件和工件。

上面的代码行 (1) 下载一个已分词的 tinyshakespeare 数据集并下载 GPT-2（124M）权重，(2) 用 C 初始化它们，用 AdamW 在 tineshakespeare 上训练 40 步（使用批大小 4，上下文长度仅 64），评估验证损失，以及采样一些文本。老实说，除非你有一个很强的 CPU（并且可以在启动命令中增加 OMP 线程数），否则你在 CPU 上训练 LLM 不会走太远，但它可能是一个好的演示/参考。输出在我的 MacBook Pro（Apple Silicon M3 Max）上看起来像这样：

[GPT-2]
max_seq_len: 1024
vocab_size: 50257
num_layers: 12
num_heads: 12
channels: 768
num_parameters: 124439808
train dataset num_batches: 1192
val dataset num_batches: 128
num_activations: 73323776
val loss 5.252026
step 0: train loss 5.356189 (took 1452.121000 ms)
step 1: train loss 4.301069 (took 1288.673000 ms)
step 2: train loss 4.623322 (took 1369.394000 ms)
step 3: train loss 4.600470 (took 1290.761000 ms)
... (trunctated) ...
step 39: train loss 3.970751 (took 1323.779000 ms)
val loss 4.107781
generating:
---
Come Running Away,
Greater conquer
With the Imperial blood
the heaviest host of the gods
into this wondrous world beyond.
I will not back thee, for how sweet after birth
Netflix against repounder,
will not
flourish against the earlocks of
Allay
---

/dev/data/(dataset).py 内的数据文件负责下载、分词和将令牌保存到 .bin 文件，从 C 中易于读取。所以例如当你运行：

python dev/data/tinyshakespeare.py

我们下载并分词 tinyshakespeare 数据集。这个的输出看起来像这样：

writing 32,768 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_val.bin
writing 305,260 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_train.bin

.bin 文件包含一个短标头（1024 字节），然后是一流的 uint16 令牌，表示使用 GPT-2 分词器的令牌 ID。更多数据集可在 /dev/data 中获得。

我还附加了一个简单的单元测试，以确保我们的 C 代码与 PyTorch 代码一致。例如在 CPU 上，用以下方式编译和运行：

make test_gpt2
./test_gpt2

这现在加载由 train_gpt2.py 写入的 gpt2_124M_debug_state.bin 文件，运行前向通过，将逻辑和损失与 PyTorch 参考实现进行比较，然后进行 10 次训练迭代与 Adam，并确保损失与 PyTorch 匹配。要测试 GPU 版本，我们运行：

# fp32 test (cudnn not supported)
make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
# mixed precision cudnn test
make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu

这测试 fp32 路径和混合精度路径。测试应该通过并打印整体好：1。

我在 doc/layernorm/layernorm.md 中附加了一个非常小的教程。这是实现 GPT-2 模型的单个层 layernorm 层的一个简单、分步指南。这是一个很好的起点，可以理解 C 中层是如何实现的。

Flash Attention

自 2024 年 5 月 1 日以来，我们使用来自 cuDNN 的 Flash Attention。因为 cuDNN 将编译时间从几秒钟增加到约一分钟，而这个代码路径现在非常新，所以它默认被禁用。你可以这样编译来启用它：

make train_gpt2cu USE_CUDNN=1

这将尝试使用 cudnn 编译并运行它。你必须在你的系统上安装了 cuDNN。cuDNN 安装说明与 apt-get 将抓取默认的 cuDNN 包集。对于最小设置，cuDNN 开发包就足够了，例如在 Ubuntu 22.04 上用于 CUDA 12.x：

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

在此基础上，你需要 cuDNN 前端，但这只是头文件。只需将仓库克隆到你的磁盘。Makefile 目前在你的主目录或当前目录中查找它。如果你把它放在其他地方，添加 CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include 到 make 命令行。

多 GPU 训练

确保安装了 MPI 和 NCCL，例如在 Linux 上：

sudo apt install openmpi-bin openmpi-doc libopenmpi-dev

对于 NCCL，按照官方网站上的说明进行（例如网络安装程序）

make train_gpt2cu
mpirun -np <number of GPUs> ./train_gpt2cu

或者简单地运行 ./scripts/ 下的我们的脚本之一。

多节点训练

确保你已按照多 GPU 部分的说明安装了 NCCL。

我们目前支持 3 种允许你运行多节点训练的方式：

使用 OpenMPI 交换 nccl id 并初始化 NCCL。有关详细信息，请参阅例如 ./scripts/multi_node/run_gpt2_124M_mpi.sh 脚本。

使用共享文件系统初始化 NCCL。有关详细信息，请参阅 ./scripts/multi_node/run_gpt2_124M_fs.sbatch 脚本。

使用 TCP 套接字初始化 NCCL。有关详细信息，请参阅 ./scripts/multi_node/run_gpt2_124M_tcp.sbatch 脚本。

如果你在 slurm 环境中运行，而你的 slurm 不支持 PMIx（我们假设这将是一个常见情况，因为 slurm-wlm 放弃了 PMIx 支持），你将不得不使用 FS（2）或 TCP（3）方法。要测试你的 slurm 是否支持 PMIx，运行：srun --mpi=list 并查看你是否在输出中得到 pmix。

如果你没有设置 slurm，你可以使用 mpirun - MPI（1）启动多节点运行。

这 3 种方法中没有一种是优越的，我们只是为你提供选项，这样你就可以在你的特定环境中运行。

学习率扫描示例

仅作为一个例子，在具有 4 个 GPU 的机器上在 TinyStories 上扫描学习率。运行一个 shell 脚本 sweep.sh（当然在你 chmod u+x sweep.sh 之后）：

#!/bin/bash

learning_rates=(3e-5 1e-4 3e-4 1e-3)

for i in {0..3}; do
    export CUDA_VISIBLE_DEVICES=$i
    screen -dmS "tr$i" bash -c "./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -l ${learning_rates[$i]} -o stories$i.log"
done

# you can bring these down with
# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit

这个例子打开了 4 个 screen 会话，并用不同的学习率运行四个命令。这将日志文件 stories$i.log 写入所有损失，你可以按照你的意愿在 Python 中绘制。一个如何解析和绘制这些日志文件的快速示例在 dev/vislog.ipynb 中。

关于这个仓库想要成为什么的更多话语

首先，我希望 llm.c 是一个教育的地方。例如，我们的 dev/cuda 文件夹是所有手工编写的、文档齐全的层内核库的地方，从非常简单的内核一直到更复杂/更快的内核。如果你有一个新内核具有各种不同的权衡，请随时在这里贡献它。

也就是说，我也希望 llm.c 非常快，甚至实用地训练网络。例如，首先，我们应该能够复现大型 GPT-2（1.6B）训练运行。这要求我们整合最快的内核，包括使用诸如 cuBLAS、cuBLASLt、CUTLASS、cuDNN 等库。我也认为这样做可以为建立专家上界和测量单位服务，例如你可以说你手写的内核达到 cuBLAS 速度的 80%，等等。然后你可以选择做一个超快速的运行，或者你可以选择"拖放"你想要使用的任何手工内核，并用那些运行。

但是，作为一个约束，我想保持根文件夹中的主线 llm.c 简单且可读。如果有一个 PR，例如提高性能 2%，但它"成本"500 行复杂的 C 代码，以及也许一个异国情调的第三方依赖，我可能会拒绝这个 PR，因为复杂性不值得。作为一个具体的例子 - 将 cuBLAS 用于 matmuls 作为根训练循环中的默认值是显而易见的：它使主线代码快得多，它是一行可解释的代码，并且它是一个非常常见的依赖。在这一边，我们可以有可以与 dev/cuda 中的 cuBLAS 竞争的手工实现。

最后，我对项目根文件夹中的复杂性会更加敏感，它包含项目的主要/默认文件。相比之下，dev/ 文件夹对我们来说有点像一个草稿空间，以开发内核或类库并共享有用的或相关的或教育代码，其中一些代码可能可以在（局部）复杂。

Notable Forks

AMD llm.c by @anthonix：支持 AMD 设备，例如 7900 XTX

C# llm.cs by @azret：这个项目的 C# 移植 Llm.cs by @nietras：这个项目的 C# 移植，重点是易于在任何平台上开始。克隆并运行 ✅

CUDA C++ llm.cpp by @gevtushenko：使用 CUDA C++ Core Libraries 的这个项目的移植 这个 fork 在 GPU MODE Discord Server 的演讲中被覆盖

C++/CUDA llm.cpp by @zhangpiu：使用 Eigen 的这个项目的移植，支持 CPU/CUDA。

WebGPU C++ gpu.cpp by @austinvhuang：一个用于便携 GPU 计算的 C++ 库，使用本机 WebGPU。旨在成为一个通用库，但也将 llm.c 内核移植到 WGSL。

C++ llm.cpp by @GaoYusong：这个项目的移植，具有 C++ 单头 tinytorch.hpp 库

Go llm.go by @joshcarp：这个项目的 Go 移植

Java llm.java by @harryjackson：这个项目的 Java 移植

Metal llm.metal by @regrettable-username：用简单、原始 C/Metal Shading Language 的 LLM 训练

Mojo llm.🔥 by @dorjeduck：这个项目的 Mojo 移植

OpenCL llm.c by @krrishnarraj：这个项目的 OpenCL 移植

Rust llm.rs by @Yijun Yu：一个 Rust 重写，目标是相同的性能 llm.rs by @ToJen：这个项目的 Rust 移植

Swift llm.swift by @otabuzzman：这个项目的 Swift 移植

Zig llm.zig by @saimirbaci：这个项目的 Zig 移植

Habana Gaudi2 llm.tpc by @abhilash1910：这个项目的 Habana Gaudi2 移植

Nim llm.nim by @Vindaar：这个项目的 Nim 移植

组织开发的方式

遇到了仓库的具体问题？使用 Issues。

有一些代码要贡献？开启一个 PR

想讨论仓库、提问等？看看 Discussions。

要快一点？我在我的 Zero to Hero Discord 频道上创建了一个新的 #llmc 频道。