llm.c：实现了大语言模型(LLM)训练的简单、纯 C/CUDA 版本，无需 PyTorch 或 cPython

项目简介

LLM简单、纯 C/CUDA 的训练。不需要 245MB 的 PyTorch 或 107MB 的 cPython。例如，训练 GPT-2 （CPU， fp32） 是单个文件中的 ~1,000 行干净代码。它可以立即编译和运行，并且与 PyTorch 参考实现完全匹配。我选择 GPT-2 作为第一个工作示例，因为它是 LLMs的祖父，第一次将现代堆栈放在一起。

目前，我正在研究：

• 直接实现 CUDA，这将明显更快，并且可能接近 PyTorch。

• 使用 SIMD 指令加速 CPU 版本，x86 上的 AVX2 / ARM 上的 NEON（例如 Apple Silicon）。

• 更现代的建筑，例如 Llama2、Gemma 等。

对于存储库，我想维护干净、简单的参考实现以及更多优化的版本，这些版本可以接近 PyTorch，但代码和依赖项的一小部分。

快速上手

下载数据集并对其进行标记化。tinyshakespeare 数据集是下载和标记化速度最快的数据集：

pythonprepro_tinyshakespeare.py

这打印：

Saved32768tokenstodata/tiny_shakespeare_val.binSaved305260tokenstodata/tiny_shakespeare_train.bin

.bin 文件是 int32 数字的原始字节流，指示带有 GPT-2 分词器的令牌 ID。或者，您也可以使用 prepro_tinystories.py 对 TinyStories 数据集进行标记化。

原则上，我们已经准备好在这里训练模型了。但是，基线 CPU/fp32 参考代码效率低下，因此从头开始训练这些模型还不切实际。取而代之的是，我们使用 OpenAI 发布的 GPT-2 权重进行初始化，然后进行微调。为此，我们必须下载 GPT-2 权重并将它们保存为我们可以在 C 中加载的检查点：

pythontrain_gpt2.py

你会将 nanoGPT 中的这段代码识别为 PyTorch 中的简单 GPT-2 参考实现。该脚本将下载 GPT-2 （124M） 模型，对单批数据进行 10 次迭代的超拟合，运行几个生成步骤，最重要的是它将保存两个文件：1） 包含用于在 C 中加载的原始模型权重 gpt2_124M.bin 的文件，以及 gpt2_124M_debug_state.bin ，它还包含更多调试状态：输入、目标、logits 和 loss。这对于调试 C 代码、单元测试以及确保我们完全匹配 PyTorch 参考实现非常有用。现在，我们只关心 gpt2_124M.bin .我们现在可以使用它们进行初始化并在原始 C 中进行训练。首先编译代码：

maketrain_gpt2

您可以查看其内部 Makefile 及其评论。它将尝试自动检测您的系统上是否有 OpenMP 可用，这对于以非常低的代码复杂性成本加快代码速度非常有帮助。编译完成后 train_gpt2 ，您可以运行它：

OMP_NUM_THREADS=8./train_gpt2

您应该根据 CPU 的内核数来调整线程数。该程序将加载模型权重、标记，它将使用 Adam lr 1e-4 运行微调循环进行几次迭代，然后从模型生成样本。该文件（我认为）非常可读，您应该看一下。简单地说，所有层都有向前和向后传递的实现，它们被串在一起形成一个大的、手动的、向前/向后/更新的循环。在我的MacBook Pro（Apple Silicon M3 Max）上，输出如下所示：

[GPT-2]max_seq_len:1024vocab_size:50257num_layers:12num_heads:12channels:768num_parameters:124439808traindatasetnum_batches:1192valdatasetnum_batches:128num_activations:73323776valloss5.252026step0:trainloss5.356189(took1452.121000ms)step1:trainloss4.301069(took1288.673000ms)step2:trainloss4.623322(took1369.394000ms)step3:trainloss4.600470(took1290.761000ms)...(trunctated)...step39:trainloss3.970751(took1323.779000ms)valloss4.107781generated:50256167731816221986111981368126323875198315226211773291019811696002638625832862621185819820424428313575969953675131984048140773617903113297036029706408219842826102811286332631119810594407198270445468010282621027288602861983237323step40:trainloss4.377757(took1366.368000ms)

该代现在只为您提供令牌 ID，我们必须将其解码回文本。我们也可以很容易地在 C 中实现这一点，因为解码非常简单，它只是字符串块查找和打印。现在我们可以使用 tiktoken：

importtiktokenenc=tiktoken.get_encoding("gpt2")print(enc.decode(list(map(int,"50256167731816221986111981368126323875198315226211773291019811696002638625832862621185819820424428313575969953675131984048140773617903113297036029706408219842826102811286332631119810594407198270445468010282621027288602861983237323".split()))))

其中打印：

<|endoftext|>ComeRunningAway,GreaterconquerWiththeImperialbloodtheheaviesthostofthegodsintothiswondrousworldbeyond.Iwillnotbackthee,forhowsweetafterbirthNetflixagainstrepounder,willnotflourishagainsttheearlocksofAllay

我喜欢Netflix的出现方式，很明显，过去训练的阴影仍然潜伏在模型中。我没有尝试调整微调超参数，所以这很可能需要改进很多，尤其是如果一个人要训练更长的时间。

测试

我还附加了一个简单的单元测试，以确保我们的 C 代码与 PyTorch 代码一致。编译和运行时使用：

maketest_gpt2./test_gpt2

现在，它会加载 gpt2_124M_debug_state.bin 文件，运行前向传递，将日志和损失与 PyTorch 参考实现进行比较，然后使用 Adam 进行 10 次训练迭代，并确保损失与 PyTorch 匹配。