支持AMD GPU的llm.c

anthonix/llm.c: LLM training in simple, raw C/HIP for AMD GPUs (github.com)

llm.c for AMD devices

This is a fork of Andrej Karpathy's llm.c with support for AMD devices.

性能

在单个7900 XTX显卡上使用默认设置，目前的训练步骤耗时约为79毫秒，相比PyTorch的夜间版本（2.4.0.dev20240513）的约97毫秒，以及tinygrad的约440毫秒来说，表现更优。

对于多GPU训练，在装有四个7900 XTX显卡的机器上，吞吐量达到了每秒约210,000个令牌。

更新（2024年5月28日）：在单个7900 XTX显卡上，快速注意力分支的训练步骤时间已经降低到58.340831毫秒，或者在四个7900 XTX显卡上达到了每秒318,777个令牌的吞吐量。目前正在研究双缓冲技术以进一步推动性能提升。

状态

train_gpt2_fp32（基线，最小改动）：使用32位浮点数（FP32）训练GPT-2模型（基线版本，仅进行最小改动）
train_gpt2 with BF16（基线，最小改动）：使用半精度浮点数（BF16）训练GPT-2模型（基线版本，仅进行最小改动）
train_gpt2 with BF16 and multiple GPUs：使用半精度浮点数（BF16）并在多个GPU上训练GPT-2模型
RDNA3 优化内核（进行中）：针对RDNA3架构优化的内核（仍在开发中）
CDNA3 优化内核：针对CDNA3架构优化的内核（具体状态未提及）

快速入门（AMD目标）

安装ROCm 6.1.1，检出仓库，并执行以下步骤：

pip install -r requirements.txt （安装所需的依赖项）
python prepro_tinyshakespeare.py （预处理tinyshakespeare数据集）
export HF_ENDPOINT=https://hf-mirror.com
python train_gpt2.py （（此步骤可能用于生成某种训练数据或配置，但具体细节未在给定指令中明确））
make train_gpt2amd （编译AMD特定版本的GPT-2训练程序）

ROCm 6.0.2报错：

ld.lld: error: unable to find library -ldevice_gemm_operations
ld.lld: error: unable to find library -ldevice_other_operations
ld.lld: error: unable to find library -lstdc++
clang: error: linker command failed with exit code 1 (use -v to see invocation)
make: *** [Makefile:285：train_gpt2amd] 错误 1

不知道ROCm 6.1.1是否能编译成功。

6. ./train_gpt2amd （运行AMD特定版本的GPT-2训练程序）

[原始README]

karpathy/llm.c

在简单、纯C/CUDA中进行LLM（大型语言模型）训练。无需245MB的PyTorch或107MB的cPython。在单个文件llm.c/train_gpt2.c中，使用CPU和fp32（32位浮点数）训练GPT-2大约需要1,000行干净的代码，而在llm.c/train_gpt2.cu中使用GPU训练大约需要3,000行代码（增加了CUDA内核）。代码立即编译并运行，它与PyTorch的参考实现完全匹配，并且当前速度略快于（编译后的）PyTorch（使用bf16、torch编译和flash attention）。我选择GPT-2作为第一个工作示例，因为它是LLM的鼻祖，也是现代堆栈首次结合在一起的例子。

我们当前的目标是复现GPT-2。要了解当前正在进行的工作的概述，请参阅最新的State of the Union帖子。

2024年5月28日更新：一个有用的近期帖子可能是“在llm.c中用90分钟和20美元复现GPT-2（124M）”，我在其中详细说明了从零开始复现124M/350M模型的GPT-2微系列的步骤。展示启动命令的文件本身是run124M.sh和run350M.sh。

我希望这个仓库只维护C和CUDA代码。这个仓库到其他语言的移植非常受欢迎，但应该在单独的仓库中完成，然后我很乐意在下面的“显著分支”部分链接到它们，就像我在llama2.c分支中所做的那样。

快速入门（GPU，速度慢但稳定且适合学习）

对于“我不在乎其他任何事情，我只想训练，并且我有GPU”的这部分用户。请运行以下命令：

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

以上命令（1）会下载tinyshakespeare数据集，并使用GPT-2的Tokenizer进行分词处理，（2）下载并保存GPT-2（124M）的权重，（3）在C/CUDA中从这些权重初始化，并在tinyshakespeare数据集上使用AdamW优化器进行一轮（epoch）的训练（使用批量大小4，上下文长度1024，总共74步），评估验证损失，并生成一些文本样本。请注意，在这个快速入门中，我们使用的是CUDA代码的fp32版本train_gpt2_fp32.cu。在下一节中，我们将介绍当前“主流”的train_gpt2.cu，它使用混合精度，运行速度大约快2倍。

快速入门（GPU，最前沿的优化）

我想看到它运行得更快。在这种情况下，切换到我们最主要的、优化度最高的 train_gpt2.cu。运行如下命令：

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

如果你额外安装了cuDNN（请参见下面的CUDA部分），你可以通过flash attention实现更快的速度。调整make命令，如下编译带有cudnn/flash attention的版本：

make train_gpt2cu USE_CUDNN=1
./train_gtp2cu

这段话的意思是，如果你已经安装了cuDNN，那么在编译和运行一个名为train_gpt2cu的程序时，可以通过设置USE_CUDNN=1来启用cuDNN的支持，从而利用cuDNN提供的加速功能来提高程序的运行速度。这里的make是一个构建工具，用于自动化编译过程，而./train_gpt2cu则是运行编译后的程序。
请注意，默认的批量大小非常小（4）。如果你的GPU有足够的内存，我建议你将其增加到例如32：

./train_gtp2cu -b 32

我的标准单GPU "生产" 运行（例如使用A100 40GB）不是训练TinyShakespeare，而是训练TinyStories，示例如下：

python dev/data/tinystories.py
make train_gtp2cu USE_CUDNN=1
./train_gtp2cu -i dev/data/tinystories/TinyStories_train.bin \
               -j dev/data/tinystories/TinyStories_val.bin \
               -v 250 -s 250 -g 144 -o stories.log -b 32

-i 标志是输入数据的通配符模式，`-j` 是验证数据。另外，我将验证损失和采样的频率减少到每250步，并且在采样阶段采样144个tokens（大约能装下一篇故事），批量大小为32。
如果你想要训练实际的、真实的预训练数据，查看最近添加的对 fineweb数据集的支持。与上面的数据集不同，这里的训练/验证tokens 不是放在一个.bin文件中，而是现在有多个数据分片。这里有一个示例：

# 将FineWeb数据以1亿个标记的分片写入到dev/data/fineweb10B
python dev/data/fineweb.py -s 100000000
# 编译并运行
./train_gtp2cu -i "dev/data/fineweb10B/fineweb_train_*.bin" \
               -j "dev/data/fineweb10B/fineweb_val_*.bin" \
               -v 250 -s 250 -g 144 -o fineweb.log -b 32

其中，你会注意到使用了通配符 * 来匹配所有的训练分片。

快速入门（多GPU）

很好，让我们更进一步。我们将使用MPI和NCCL来进行多GPU训练。上面部分的内容依然适用，但需要做以下改变：

# 安装MPI的示例：
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev

# 运行命令现在需要以 mpirun 开头：
mpirun -np <你机器上的GPU数量> ./train_gpt2cu

在最后的命令中，替换为你想要运行的GPU数量。上面部分讨论的所有标志在这里也适用。

快速开始（CPU）

“我是如此贫穷以至于连GPU都没有”的部分。你仍然可以进行训练！但你不会走得太远。你仍然可以微调一个GPT-2小模型（1.24亿参数模型）以输出像莎士比亚式的文本，作为示例：

pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py
make train_gpt2
OMP_NUM_THREADS=8 ./train_gpt2

上面的行（1）下载了tinyshakespeare 数据集，使用GPT-2的分词器对其进行分识，（2）下载并保存GPT-2（124M）的权重，（3）在C中从它们初始化并在tineshakespeare上训练40步骤使用AdamW（使用批量大小4，上下文长度只有64），评估验证损失，并抽取一些文本。诚实地讲，除非你有一个强大的CPU（并且可以在启动命令中增加OMP线程的数量），你在CPU上训练大型语言模型（LLMs）不会走得太远，但它可能是一个不错的演示/参考。

训练：更多细节

在`/dev/data/(dataset).py`的数据文件负责下载、分词并将分词保存到文件中。例如，当你运行：

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字节），随后是一个流式的tokens以uint16格式，表明了用GPT-2分词器的分词id。更多数据集可以在`/dev/data`找到。
原则上，一旦我们得到了tokens，我们就准备在这里开始训练模型。然而，当前的代码还不能从零开始训练（很快就会加入），所以我们从OpenAI发布的预训练模型初始化训练并进行微调。为此，我们需要下载GPT-2的权重并将其作为我们可以在C语言中加载的检查点来保存。这就是当你运行以下脚本时发生的事：

python train_gpt2.py

你会认出这段代码，它来自nanoGPT，是一个简单的PyTorch中的GPT-2参考实现。这个脚本将下载GPT-2 (124M)模型，对单批数据进行了10次迭代的超过拟合，运行了几步生成，并且更重要的是它将保存三个文件：1）`gpt2_124M.bin`文件，它包含了原始模型权重以便在C语言中加载，2）`gpt2_124M_debug_state.bin`文件，这也包含了更多的调试状态：输入，目标，逻辑和损失（对于调试和单元测试很有用），最后3）`gpt2_tokenizer.bin`文件，它存储了GPT-2分词器的词汇表，将分词id转换为UTF-8编码字符串片段的字节序列。文件还保存了上述的fp32版本，以及它们的bfloat16版本以供混合精度训练。我们现在可以用这些模型权重来初始化并继续在原始C语言中训练。然后我们用`make`命令来编译训练程序。目前有三个并行实现：

# 简单的，CPU，参考代码版本
make train_gpt2
# 单GPU fp32 CUDA版本
make train_gpt2fp32cu
# 多GPU混合精度CUDA版本
make train_gpt2cu

你可以查阅`Makefile`及其注释。它会尝试自动探测很多工具和库（例如：cuDNN, OpenMP, OpenMPI, nvcc），你要尽量获得尽可能多的勾号。比如当我在我配置完善的机器上运行`make train_gpt2cu USE_CUDNN=1`，我们看到：

✓ cuDNN found, will run with flash-attention
✓ OpenMP found
✓ OpenMPI found, OK to train with multiple GPUs
✓ nvcc found, including GPU/CUDA support

有些人在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

一旦编译好了二进制文件，我们就可以运行它。例如最简单的CPU参考版本运行如下：

OMP_NUM_THREADS=8 ./train_gpt2

你应该根据你的CPU拥有多少核心来调整线程数量。该程序将加载模型权重、tokens，它将运行一个微调循环数次与Adam lr 1e-4，然后从模型中生成一个样本。这个文件很可读，你应该看一看。简单来说，就是所有层的前向和后向传递的实现，并且它们被串接在一个大的、手动的、前向/后向/更新循环中。输出看起来像这样在我的MacBook Pro（苹果硅 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
---

我喜欢Netflix的出现，很明显模型的训练过往仍在影响它。我没有尝试调整微调的超参数，所以这个结果很可能还可以大幅度提高。我还注意到，不同的平台（例如MacOS/Linux）将会（遗憾地）给出非常微小的不同结果，所以可能不要期望得到上文提供的确切的数字或生成结果。

最后，代码还在变动中。如果发生任何你没预料到或之前运行正常的奇怪事情，请尝试`git pull`，重新运行所有上面的命令，回到这个README文件参考，等等。

测试

我还附上了一个简单的单元测试，以确保我们的 C 代码与 PyTorch 代码一致。以 CPU 为例，编译并且执行如下：

make test_gpt2
./test_gpt2

这将加载 gpt2_124M_debug_state.bin 文件，执行一个前向传递，与 PyTorch 参考实现比较 logits 和 loss，然后进行 10 次迭代的 Adam 训练确保损失与 PyTorch 匹配。要测试 GPU 版本，我运行：

# fp32 测试（不支持 cudnn）
make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
# 混合精度 cudnn 测试
make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu

教程

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

CUDA

整个训练循环也在一个文件中使用纯CUDA实现，但是内核的优化还在进行中。目前，我们的速度略微超过了PyTorch Nightly的速度。我们组织代码的方式是，在`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文件，该文件包含混合精度训练。

在下面的描述中，我现在默认使用fp32版本，因为它当前更加便携和稳定，然后在最后我将介绍新的混合精度版本。

正确性。首先，我们可以做10次训练迭代并验证我们的代码是否与PyTorch完全匹配和再现数字：

make test_gpt2fp32cu
./test_gpt2fp32cu

这会打印出`overall okay: 1`。因此，前向激活、后向梯度和10次迭代的各个损失值都完全匹配。

训练。在单GPU上以fp32进行训练：

make train_gpt2fp32cu
./train_gpt2fp32cu

这将加载tiny_shakespeare数据集的验证和训练划分。在默认设置B=4，T=1024下，有8个验证批次和74个训练批次。该脚本目前配置为以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
generating:
---
?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.

<|endoftext|>My little nephew:
Keep on with me, my

这在我的A100上大约运行了~10秒。我们可以这样与naive PyTorch进行比较，我们开启了`torch.compile`和使用TensorCores，它使用的是tf32类型：

python 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）。训练和测试的命令都一样，只需省略fp32部分：

make train_gpt2cu
./train_gpt2cu

make test_gpt2cu
./test_gpt2cu

如果您有最新的CUDA，应该期望它可以编译OK，并且应该看到性能大幅改进。

Flash Attention。截至2024年5月1日，我们现在支持来自cuDNN的Flash Attention。因为cuDNN使编译时间从几秒增加到约一分钟，并且这个代码路径现在非常新颖，目前默认情况下是禁用的。您可以通过以下方式编译来启用它：

make train_gpt2cu USE_CUDNN=1

这将尝试使用cudnn进行编译并运行。您必须在您的系统上安装cuDNN。通过apt-get安装的cuDNN安装说明将获取默认的cuDNN包集。对于最小安装来说，cuDNN dev包是足够的，例如在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前端，但这只是头文件。只需将仓库克隆到您的磁盘即可。如果您将其放在其他地方，将`CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include`添加到`make`命令行中。

多GPU训练。截至2024年4月26日，现在还支持使用MPI和NCCL进行多GPU训练。确保您安装了MPI，例如在Linux上：

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

然后：

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

fp32版本的代码不支持多GPU。这是因为我们希望GPT-2 fp32版本成为CUDA优化课程的一个不错的教育终点。混合精度版本是我们进行前沿开发的版本，因此这是支持多GPU训练的版本。

实验/扫描

现在基本的 argparse 和日志功能已存在于 .cu 脚本中，我们可以开始进行第一轮学习率扫描。目前这还相当手动，但是下面记录了一个示例过程，展示了在一台拥有4个GPU的机器上针对TinyStories数据集进行学习率扫描的过程。在你当然已经使用 chmod u+x sweep.sh 命令赋予了sweep.sh脚本执行权限后，运行一个名为 sweep.sh 的shell脚本：

#!/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

# 你可以使用以下命令关闭这些屏幕会话
# 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`能够成为的几点想法：
首先，我希望`llm.c`能成为一个教学场所。举个例子，我们的`dev/cuda`文件夹是一个包含了各种手写、文档齐全的内核库，从最简单的内核开始，到更复杂／更快速的内核。如果您有带有不同权衡的新内核，请随时贡献至此。
话虽如此，我也希望`llm.c`能够非常快速，甚至在实际中用于训练网络。比如，首先，我们应该能复现大型GPT-2（1.6B）的训练过程。这需要我们整合各种最快的内核，包含使用如cuBLAS、cuBLASLt、CUTLASS、cuDNN等库。我同样认为这样做对于确立一个专家级的上限，并作为一种度量单位，具有教育意义。比如，你可以说你手写的内核速度达到了cuBLAS的80%等。然后你可以选择进行一个超快的运行，或者你可以选择 "拖放 "任何你想使用的手动内核，并用那些运行。
然而，作为一个限制，我希望保持根目录下的主线`llm.c`简单且可读。如果有一个PR能够改进性能2%，但它“花费了”500行复杂的C代码，可能还有一些非主流的第三方依赖，我可能会拒绝这个PR，因为复杂性不值得。具体的一个例子 - 让cuBLAS成为根训练循环的默认矩阵乘法是明智的：它让主线代码快了很多，它是一行易于理解的代码，而且是一个非常常见的依赖。在这一侧，我们可以在`dev/cuda`中有与cuBLAS竞争的手写实现。
最后，对于项目根目录中包含主要／默认文件的部分，我会对复杂性更加敏感。相比之下，`dev/`文件夹更像是一个草稿空间，供我们开发内核或类的库，分享有用或相关的或教育性代码，其中一些代码是可以接受的（局部）复杂性。