AI知识库

53AI知识库

学习大模型的前沿技术与行业应用场景


llm.c:实现了大语言模型(LLM)训练的简单、纯 C/CUDA 版本,无需 PyTorch 或 cPython
发布日期:2024-04-29 14:14:26 浏览次数: 1792


项目简介


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 数据集是下载和标记化速度最快的数据集:

python prepro_tinyshakespeare.py

这打印:

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

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

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

python train_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 中进行训练。首先编译代码:

make train_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: 124439808train dataset num_batches: 1192val dataset num_batches: 128num_activations: 73323776val loss 5.252026step 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.107781generated: 50256 16773 18162 21986 11 198 13681 263 23875 198 3152 262 11773 2910 198 1169 6002 6386 2583 286 262 11858 198 20424 428 3135 7596 995 3675 13 198 40 481 407 736 17903 11 329 703 6029 706 4082 198 42826 1028 1128 633 263 11 198 10594 407 198 2704 454 680 1028 262 1027 28860 286 198 3237 323step 40: train loss 4.377757 (took 1366.368000 ms)

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

import tiktokenenc = tiktoken.get_encoding("gpt2")print(enc.decode(list(map(int, "50256 16773 18162 21986 11 198 13681 263 23875 198 3152 262 11773 2910 198 1169 6002 6386 2583 286 262 11858 198 20424 428 3135 7596 995 3675 13 198 40 481 407 736 17903 11 329 703 6029 706 4082 198 42826 1028 1128 633 263 11 198 10594 407 198 2704 454 680 1028 262 1027 28860 286 198 3237 323".split()))))

其中打印:

<|endoftext|>Come Running Away,Greater conquerWith the Imperial bloodthe heaviest host of the godsinto this wondrous world beyond.I will not back thee, for how sweet after birthNetflix against repounder,will notflourish against the earlocks ofAllay

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

测试

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

make test_gpt2./test_gpt2

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




53AI,企业落地应用大模型首选服务商

产品:大模型应用平台+智能体定制开发+落地咨询服务

承诺:先做场景POC验证,看到效果再签署服务协议。零风险落地应用大模型,已交付160+中大型企业

联系我们

售前咨询
186 6662 7370
预约演示
185 8882 0121

微信扫码

与创始人交个朋友

回到顶部

 
扫码咨询