llm.c

LLM 训练使用简单、纯粹的 C/CUDA 实现，无需庞大的 PyTorch 或 cPython。训练 GPT-2（CPU, fp32）代码量约 1,000 行，包括在单文件 train_gpt2.c，而 GPU 训练版本代码量约 2,000 行（添加了 CUDA 内核），在 train_gpt2.cu。这些代码既快速编译运行，又完全一致于 PyTorch 的参考实现，且运行速度与其相当（fp32，无闪存注意力）。我选择 GPT-2 作为首个示例，因为它是 LLM 的鼻祖，首次将现代堆栈整合。

目前，我们的目标是重现 GPT-2，用多节点、混合精度的高效实现。有关最新进展，请查看 State of the Union 帖子。

该仓库仅保留 C 和 CUDA 代码。其他语言的移植版本可以在单独的仓库中完成，我们会在下方“著名分叉”部分链接，如同 llama2.c notable forks。

快速开始（GPU）

如果你“只想训练且有 GPU”，运行：

shell
1pip install -r requirements.txt
2python prepro_tinyshakespeare.py
3python train_gpt2.py
4make train_gpt2fp32cu
5./train_gpt2fp32cu

这将（1）下载 tinyshakespeare 数据集，并用 GPT-2 分词器分词，（2）下载并保存 GPT-2（124M）权重，（3）用 C/CUDA 初始化，在 tinyshakespeare 上训练一轮（AdamW，批量 4，背景 1024，共 74 步），评估验证损失并生成文本。此处使用 CUDA 代码的 fp32 版本 train_gpt2_fp32.cu，在 CUDA 部分介绍 train_gpt2.cu，该版本仍在开发中，使用混合精度并速度翻倍。

快速开始（CPU）

如果你“连 GPU 都没有”，运行：

shell
1pip install -r requirements.txt
2python prepro_tinyshakespeare.py
3python train_gpt2.py
4make train_gpt2
5OMP_NUM_THREADS=8 ./train_gpt2

这将（1）下载 tinyshakespeare 数据集并分词，（2）下载 GPT-2（124M）权重，（3）用 C 初始化并在 tinyshakespeare 上训练 40 步（AdamW，批量 4，背景 64），评估验证损失并生成文本。除非你有强劲 CPU，否则在 CPU 上训练 LLM 进展缓慢，但它作为演示或参考。

快速开始（多 GPU）

您将使用（更前沿的）混合精度代码：

shell
1sudo apt install openmpi-bin openmpi-doc libopenmpi-dev
2pip install -r requirements.txt
3python prepro_tinyshakespeare.py
4python train_gpt2.py
5make train_gpt2cu
6mpirun -np <GPU 数量> ./train_gpt2cu

替换 <GPU 数量>。

训练：更多细节

下载并分词数据集。tinyshakespeare 是最快的：

shell
1python prepro_tinyshakespeare.py

这将输出：

Saved 32768 tokens to data/tiny_shakespeare_val.bin
Saved 305260 tokens to data/tiny_shakespeare_train.bin

这些 .bin 文件是 int32 数字的字节流，表示 GPT-2 分词器 token id。可用 prepro_tinystories.py 对 TinyStories 数据集分词。

我们可以训练模型，但基准 CPU/fp32 参考代码效率太低，因此用 OpenAI GPT-2 权重微调。需下载权重并保存为 C 可用的检查点：

shell
1python train_gpt2.py

这是来自 nanoGPT 的 PyTorch 参考实现。会下载 GPT-2（124M），对数据迭代 10 次，生成文本并保存三文件：1）gpt2_124M.bin（C 加载权重），2）gpt2_124M_debug_state.bin（包含调试信息：输入、目标、logits 和损失），3）gpt2_tokenizer.bin（分词器词汇，将 token id 转为字节序列）。编译：

shell
1make train_gpt2

Makefile 会尝试检测 OpenMP，加速代码。遇到 Ubuntu 问题查看 Issue 19，简而言之，可能要改 CFLAGS：

# 先试这个
CFLAGS="-Ofast -fno-finite-math-only -Wno-unused-result -march=native" make train_gpt2
# 再试这个
CFLAGS="-O3 -Wno-unused-result -march=native" make train_gpt2

完成后运行：

shell
1OMP_NUM_THREADS=8
2
3 ./train_gpt2

调整线程数。会加载权重和 token，微调迭代并生成样本。文件易读，组合成大循环。输出如下：

[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)
...（截断）
step 39: train loss 3.970751 (took 1323.779000 ms)
val loss 4.107781
生成：
---
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
---

Netflix 显示了训练影子在模型中。我没调微调参数，有改进空间。不同平台会有些不同。见 token id 而非文本，更新后再运行 python train_gpt2.py。

测试

单元测试确保 C 与 PyTorch 一致。编译并运行：

shell
1make test_gpt2
2./test_gpt2

加载 gpt2_124M_debug_state.bin 文件，前向传播比对 logits 和损失，再 10 次 Adam 迭代确保一致。

教程

附了教程 doc/layernorm/layernorm.md，逐步实现 GPT-2 规范层，是理解 C 实现的起点。

CUDA

完整训练循环在单文件纯 CUDA 实现，内核优化持续。我们和 PyTorch 速度相当。在 dev/cuda 文件夹中有不断增加的内核，见 dev/cuda/README.md，复制到 train_gpt2cu.cu 中。

WIP 提醒，2024 年 4 月 23 日。我们合并了混合精度训练代码的第一个版本。我将 fp32 版本检查点到不同文件中，包括 _fp32 在文件名中，并希望保留在仓库根目录中，因为它 1）不需要最新的 CUDA 更易于编译，2）更简单并作为参考。实际上，我们希望 fp32 版本向纯 CUDA 方向发展（例如默认不调用 cuBLAS），用于教育参考，甚至可能成为 CUDA 课程的核心。此后，速度方面的开发将转向 train_gpt2.cu 文件，包括混合精度训练。

正确性。首先，我们可以进行 10 次训练迭代，验证代码与 PyTorch 数字完全一致：

shell
1make test_gpt2fp32cu
2./test_gpt2fp32cu

这会输出 overall okay: 1。所有前向激活、后向梯度以及 10 次迭代的损失都完全一致。

训练。要用 CUDA 文件训练 GPT-2，运行训练脚本：

shell
1make train_gpt2fp32cu
2./train_gpt2fp32cu

加载 tiny_shakespeare 数据集验证和训练部分。默认 B=4，T=1024，有 8 验证批次和 74 训练批次。微调一次 lr 1e-4，沿途评估验证表现并生成样本，例如：

step 1/74: train loss 4.367631 (80.639749 ms)
step 2/74: train loss 4.031242 (77.378867 ms)
step 3/74: train loss 4.034144 (77.315861 ms)
step 4/74: train loss 3.859865 (77.357575 ms)
...
step 72/74: train loss 3.085081 (78.850895 ms)
step 73/74: train loss 3.668018 (78.197064 ms)
step 74/74: train loss 3.467508 (78.009975 ms)
val loss 3.516490
生成：
---
?Where will you go?
I take you wherefore I can, myself, and must.
I cast off my beak, that I may look him up on the point;
For on his rock shall he be opencast.

My little nephew:
Keep on with me, my

在 A100 上约 10 秒。PyTorch 约 80ms/迭代，所以略好于 PyTorch。但旧版 PyTorch（2.1.0）下测量，未含 FlashAttention 或 fused 操作。

我们用 torch.compile 和 TensorCores 用 tf32 比较：

shell
1python train_gpt2.py --write_tensors 0 --sequence_length 1024 --batch_size 4 --compile 1 --tensorcores 1

编译 27 秒，之后 A100 上约 80ms/迭代。

混合精度。新的 CUDA 混合精度版本是 train_gpt2.cu，及其测试 test_gpt2.cu。许多计算发生在低精度格式（fp16 或 bf16），使得运行非常快（大约 TF32 性能的 2 倍）。值得说明的是，我描述基准实现为 fp32，但实际上是 tf32（TensorFloat32）。训练和测试同样命令：

shell
1make train_gpt2cu
2./train_gpt2cu
3
4make test_gpt2cu
5./test_gpt2cu

最新 CUDA 编译正常，速度提高约 2 倍（1.86 倍）。

多 GPU 训练。自 2024 年 4 月 26 日起，现在支持多 GPU 训练，用 MPI 和 NCCL。确保安装 MPI，例如在 Linux：

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

然后：

shell
1make train_gpt2cu
2mpirun -np <GPU 数量> ./train_gpt2cu

fp32 版本不支持多 GPU。因为我们希望 GPT-2 fp32 成为 CUDA 优化课程的教育终点。混合精度版本是我们进行前沿开发的版本，所以支持多 GPU。

实验/扫描

现在 argparse 和日志功能都在 .cu 脚本中，可以做学习率扫描。例如在 4 GPU 机器上对 TinyStories 扫描学习率，运行 shell 脚本 sweep.sh：

shell
1#!/bin/bash
2
3learning_rates=(3e-5 1e-4 3e-4 1e-3)
4
5for i in {0..3}; do
6    export CUDA_VISIBLE_DEVICES=$i
7    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"
8done
9
10# 可以用这个关闭
11# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit

4 个会话运行命令，写 stories$i.log 记录损失，可用 Python 绘制。快速脚本：

python
1import matplotlib.pyplot as plt
2%matplotlib inline
3
4def parse_log(logfile):
5  # 寻找类似 "s:100 tel:1.6952" 的行，第 100 步，验证 1.6952
6    val_steps, val_losses = [], []
7    with open(logfile, "r") as f:
8        lines are read.
9        for line in lines:
10            if "tel" in line:
11                parts are split.
12                step is int(parts[0].split(":")[1]).
13                loss is float(parts[1].split(":")[1]).
14                val_steps.append(step).
15                val_losses.append(loss).
16            return val_steps, val_losses.
17
18results = [parse_log(f"stories{i}.log") for i in range 0, 4)]
19for i, (val_steps, val_losses) in enumerate(results):
20    plt.plot(val_steps, val_losses, label="run {}".format(i))
21plt.xlabel("steps")
22plt.ylabel("loss")
23plt.legend()

仓库理念

希望 llm.c 是教育

场所。如 dev/cuda 手写内核库，文档详细。贡献新内核欢迎。

也希望 llm.c 十分快速，实用训练 GPT-2（1.6B）。需最快内核，包括 cuBLAS, cuBLASLt, CUTLASS, cuDNN。我认为这有教育意义，可说手写内核是 cuBLAS 80%。可快速训练或选手写内核运行。

但 llm.c 根目录保持简单可读。2% 提升需 500 行复杂代码或库 PR 会被拒。更愿 90% PyTorch 速度保持 2,000 行代码。如默认 cuBLAS，是简单实用。可在 dev/cuda 手写内核。

根目录更敏感。dev/ 更像临时空间。

著名分叉

• Mojo
  • llm.🔥 by @dorjeduck：Mojo 移植
• C#
  • llm.cs by @azret：C# 移植
• Rust
  • llm.rs by @ToJen：Rust 移植
• Metal
  • llm.metal by @regrettable-username：Metal 移植
• Zig
  • llm.zig by @saimirbaci：Zig 移植
• Go
  • llm.go by @joshcarp：Go 移植

讨论

• Issues：遇到问题
• PR：贡献代码
• Discussions：讨论
• #llmc 频道在 Zero to Hero Discord

许可证

MIT