karpathy--llm.c/README.md

Note

本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。
English · 原始项目 · 上游 README
原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。

llm.c

用简洁、纯粹的 C/CUDA 实现 LLM，无需 245MB 的 PyTorch 或 107MB 的 cPython。当前重点是预训练（pretraining），尤其是复现 GPT-2 和 GPT-3 迷你系列，同时在 train_gpt2.py 中提供并行的 PyTorch 参考实现。你会认出该文件是我早期项目 nanoGPT, 的稍加改动版本。目前，llm.c 比 PyTorch Nightly 略快（约 7%）。除了 train_gpt2.cu 中的前沿主线代码外，我们还在单文件 train_gpt2.c 中提供了一个约 1,000 行整洁代码的简单 CPU fp32 参考实现。我希望本仓库仅维护 C 和 CUDA 代码。欢迎移植到其他语言或仓库，但应在独立仓库中进行，我很乐意在下方「notable forks」部分链接它们。开发者协作在 Discussions 以及 Discord 上进行，可在 Zero to Hero 频道的 #llmc 频道，或在 GPU MODE Discord 的 #llmdotc 上交流。

quick start

当前入门 llm.c 仓库的最佳方式是复现 GPT-2 (124M) 模型。Discussion #481 详细介绍了这一过程。我们可以在 llm.c 及并行的 PyTorch 实现中复现 GPT-2 和 GPT-3 系列的其他模型。请查看 scripts README。

debugging tip: 运行 make 命令构建二进制文件时，将其中的 -O3 替换为 -g，即可在你喜爱的 IDE（例如 vscode）中单步调试代码。

quick start (1 GPU, fp32 only)

如果你不会在多节点上训练、对混合精度不感兴趣，并且想学习 CUDA，fp32（legacy）文件可能适合你。这些文件是 llm.c 早期历史中的「checkpoint」版本，被定格在那一刻。它们更简单、更可移植，也可能更容易理解。单 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) 单元测试中使用的「debug state」（小批量数据及目标激活值和梯度），3) GPT-2 分词器，以及 3) 已分词化的 tinyshakespeare 数据集。或者，不运行 .sh 脚本，也可按以下步骤手动重新创建这些产物：

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

quick start (CPU)

「穷到连一块 GPU 都没有」专区。你仍然可以欣赏 llm.c 的训练过程！但走不了太远。与上述 fp32 版本类似，CPU 版本是 llm.c 历史上更早的一个 checkpoint，当时它还只是 C 语言的简单参考实现。例如，你可以不从头训练，而是将 GPT-2 small (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) 权重，（3）在 C 中从它们初始化，在 tineshakespeare 上用 AdamW 训练 40 步（batch size 4，context length 仅 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
---

datasets

/dev/data/(dataset).py 内的数据文件负责下载、分词并将 token 保存为可从 C 轻松读取的 .bin 文件。例如运行：

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 存储的 token 流，表示 GPT-2 分词器的 token id。更多数据集可在 /dev/data 中找到。

test

我还附带了一个简单的单元测试，用于确保 C 代码与 PyTorch 代码一致。以 CPU 为例，编译并运行：

make test_gpt2
./test_gpt2

该测试会加载由 train_gpt2.py 写入的 gpt2_124M_debug_state.bin 文件，执行前向传播，将 logits 和 loss 与 PyTorch 参考实现对比，然后进行 10 次 Adam 训练迭代并确保 loss 与 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 路径和混合精度路径。测试应通过并打印 overall okay: 1。

tutorial

我在此附上了一份很短的教程：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 installation instructions 中通过 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 frontend,，但这只是头文件。将仓库克隆到本地即可。Makefile 目前会在你的主目录或当前目录中查找它。若放在其他位置，请在 make 命令行中添加 CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include。

多 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/ 下的某个脚本。

多节点训练

请确保已按照 multi-GPU 章节的说明安装 NCCL。

目前我们支持 3 种方式来进行多节点训练：

1. 使用 OpenMPI 交换 nccl id 并初始化 NCCL。详见例如 ./scripts/multi_node/run_gpt2_124M_mpi.sh 脚本。
2. 使用共享文件系统初始化 NCCL。详见 ./scripts/multi_node/run_gpt2_124M_fs.sbatch 脚本。
3. 使用 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 种方法并无优劣之分，我们只是提供多种选项，以便你在特定环境中运行。

实验 / 参数扫描（sweeps）

下面是一个示例流程：在一台配备 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，其中包含所有 loss，你可以用 Python 按需绘制。解析并绘制这些日志文件的快速示例见 dev/vislog.ipynb。

仓库

关于我希望这个仓库成为什么样子，再多说几句：

首先，我希望 llm.c 成为一个教学场所。例如，我们的 dev/cuda 文件夹是一个内核库，涵盖各层的手工编写、文档详尽的内核，从非常简单的内核一直到更复杂/更快的内核。如果你有具有不同权衡（tradeoffs）的新内核，欢迎在此贡献。

话虽如此，我也希望 llm.c 非常快，甚至能实际用于训练网络。例如，起步时我们应该能够复现大型 GPT-2（1.6B）训练运行。这要求纳入最快的内核，包括使用 cuBLAS、cuBLASLt、CUTLASS、cuDNN 等库。我认为这样做也有教学意义：建立一个专家级上限和度量单位，例如你可以说你的手写内核达到 cuBLAS 速度的 80% 等。然后你可以选择超快运行，也可以选择“拖放”你想使用的任意手工内核来运行。

不过，作为约束，我希望保持根目录中的主线 llm.c 简单且可读。如果某个 PR 例如将性能提升 2%，但“代价”是 500 行复杂 C 代码，也许还有冷门第三方依赖，我可能会拒绝该 PR，因为复杂度不值得。举个具体例子——在根训练循环中将 cuBLAS 用于 matmul 设为默认是显而易见的：它让主线代码快得多，只是一行可理解的代码，且是非常常见的依赖。与此同时，我们可以在 dev/cuda 中提供能与 cuBLAS 竞争的手工实现。

最后，我对项目根文件夹中的复杂度会敏感得多，根目录包含项目的主/默认文件。相比之下，dev/ 文件夹更像是我们开发内核或类库、分享有用/相关/教学代码的草稿空间，其中部分代码可以（在局部）复杂一些。

值得关注的分支（fork）

• 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 对本项目的移植
    • 该分支曾在 GPU MODE Discord Server 的这场讲座 中介绍

• C++/CUDA

  • llm.cpp by @zhangpiu: 使用 Eigen, 对本项目的移植，支持 CPU/CUDA。

• WebGPU C++

  • gpu.cpp by @austinvhuang: 使用原生 WebGPU 在 C++ 中进行可移植 GPU 计算的库。目标是成为通用库，同时也将 llm.c 内核移植到 WGSL。

• C++

  • llm.cpp by @GaoYusong: 本项目的 C++ 移植版，提供 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 频道

许可证

MIT