InfoQ 推荐 ( ) • 2024-04-09 18:03

徒手用 1000 行 C 语言实现,不依赖庞大的外部库,Mac 即可运行。

如今这年头,徒手写神经网络代码已经不算事儿了,现在流行手搓大模型训练代码了!这不,今天,特斯拉前AI总监、OpenAI 创始团队成员Andrej Karpathy仅用1000行简洁的C代码,就完成了 GPT-2 大模型训练过程。

 

 

 

几个小时前,Andrej Karpathy推出了一个名为llm.c的项目,旨在用纯C语言训练LLM,这种方法的主要优势在于它显著减少了依赖库的体积——不再需要245MB的PyTorch和107MB的cPython,这样可以使得模型训练过程更为轻量化和高效。该项目还可以立即编译和运行,并且与PyTorch的参考实现完全匹配。

 

Karpathy表示他之所以选择GPT-2作为首个工作示例,是因为它大语言模型鼻祖的定位,亦属现代AI堆栈的首次组合。因此,选择 GPT-2 作为起点,可以让我们更容易地理解和实践大型语言模型训练。

 

徒手实现GPT-2后,Karpathy将这份代码放到了GitHub上,以MIT协议开源。短短几个小时,就超过了2500颗星,并且数据还在不断持续上涨......

 

 

 

将大模型原理拆解得无比简单

 

Andrej Karpathy 是全球人工智能领域的知名科学家,也是 OpenAI 的创始成员和研究科学家。

 

他于2009年本科毕业于多伦多大学,获得计算机科学和物理学学士学位。2011年硕士毕业于不列颠哥伦比亚大学,随后前往斯坦福大学AI Lab(SAIL)读博,师从著名学者李飞飞,是全球最早将深度学习应用于计算机视觉研究的学者之一。

 

在求学期间,Andrej Karpathy曾在谷歌和DeepMind实习,后来在OpenAI刚刚成立时加入并担任研究科学家。直到2017年6月,他被马斯克挖去,担任特斯拉人工智能部门主管,直接向马斯克汇报。在特斯拉工作的五年里,他主导了特斯拉自动辅助驾驶系统Autopilot的开发。这项技术对于特斯拉的完全自动驾驶系统 FSD 至关重要,也是马斯克针对 Model S、Cybertruck 等车型推销的主要卖点。在各大新闻中,他被誉为“特斯拉的秘密武器”。

 

去年Karpathy曾短暂回到OpenAI,然后又在OpenAI众人忙于内斗时抽空录制了一个长达一小时的教学视频《大型语言模型入门》。

 

 

Karpathy 在视频中首先介绍了一些 LLM 入门知识,然后以 Meta 推出的开源大模型 Llama 2-70b 为例进行了讲解。该模型有 700 亿参数,主要包含两个文件,分别是参数文件,文件大小为 140GB,以及运行这些参数的代码,以 C 语言为例需要约 500 行代码。

 

Karpathy 表示只要有这两个文件再加上一台 MacBook,我们就可以构建一个独立的系统,无需联网或其他设施。

 

 

 

大模型训练,可以理解为是对互联网数据进行有损压缩,一般需要一个巨大的GPU集群来完成。以Llama 2-70b为例的话,就是使用了类似网络爬取的约 10TB 的文本,用6000 个 GPU ,耗资 200 万美元,训练约 12 天,最后获得基础模型。

 

 

 

基础模型即上图里140GB的“压缩文件”(压缩率约100倍),就等于靠这些数据对世界形成了理解,那它就可以进行“预测”工作了。

 

Karpathy之前还分享过他的学习经验,就是开始时要尝试从0开始,写一些原生代码,帮助理解消化知识点。也就是说,徒手实现代码才是最有效的学习方式。

 

两年前,Karpathy就曾基于 PyTorch,仅用 300 行左右的代码就写出了一个小型 GPT 训练库,并将其命名为 minGPT,用这份代码揭开了GPT神秘的面纱。

 

截图来源:https://github.com/karpathy/minGPT"

 

因为大多数 GPT 模型的实现都过于庞大,而minGPT 做到了小、干净、可解释和具有教育意义,所以Karpathy的这300行代码是学习 GPT 的最佳资源之一,可以用来深入理解GPT 是如何工作的。

 

用C语言实现LLM

 

这次,Andrej Karpathy单纯通过C/CUDA实现大语言模型训练,且无需245 MB PyTorch或107 MB cPython。例如,训练GPT-2(CPU,fp32单精度)需要在单个文件中使用约1000行简洁代码,可立即编译并运行、且与PyTorch参考实现完全匹配。

 

 

从某种意义上说,Karpathy确实在尝试重新设计LLM的架构。他通过llm.c项目探索一种更简单、更高效的训练LLM方法。与现有LLM架构相比,这种新架构的主要亮点包括:

 

1. 代码简洁性:仅使用约1000行代码就能完成GPT-2模型的训练,相比之下显著降低了复杂度。

2. 独立性:不依赖庞大的外部库如PyTorch或cPython,使得部署和运行更加轻便快捷。

3. 高效性:直接使用C/CUDA进行编程有望提高计算效率和训练速度。

 

有网友问Karpathy为何不用Rust,Karpathy回复说,“我完全理解Rust的吸引力。然而,我仍然觉得 C 语言非常棒。它简单、干净、可移植,在审美上也十分优美。使用 C 语言就像直接与机器交流一样。”

 

 

这种语言选择也让网友们纷纷感叹:

 

“我们正在掀起一场 C 语言复兴!”

 

“真男人就应该用 C 语言编程。”

 

 

Karpathy以更简单、更原始的C/CUDA架构来做LLM的训练,其中还涉及算法优化、计算资源管理等多个方面。

 

 

你会看到,项目在开始时一次性分配所有所需的内存,这些内存是一大块 1D 内存。然后在训练过程中,不会创建或销毁任何内存,因此内存占用量保持不变,并且只是动态的,将数据批次流过。这里的关键在于手动实现所有单个层的前向和后向传递,然后将它们串联在一起。例如,这里是 layernorm 前向和后向传递。除了 layernorm 之外,我们还需要编码器、matmul、自注意力、gelu、残差、softmax 和交叉熵损失。

 

“一旦你拥有了所有的层,接下来的工作只是将它们串在一起。讲道理,写起来相当乏味和自虐,因为你必须确保所有指针和张量偏移都正确排列, ”Karpathy 表示。

 

 

另外Karpathy还在doc/layernorm/layernorm.md中附上了短小的使用教程。里面是一份简单的分步指南,用于实现GPT-2模型中的单一层,即layernorm层,希望能成为大家理解在C中实现大语言模型各层的理想起点。

 

更重要的是,他还用自己的MacBook Pro(苹果M3 Max芯片)演示了整个训练过程,对照他之前的大模型入门教程,就可以轻松了解如今炙手可热的LLM是怎么一回事儿了。

 

训练过程:

 

首先下载数据集并token化。

 

python 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化器的token id。或者也可以使用prepro_tinystories.py对TinyStories数据集进行标注。

 

原则上,到这里就已经可以开始训练模型。为提高效率,可以使用OpenAI发布的GPT-2权重进行初始化,而后加以微调。为此需要下载GPT-2权重并将其保存为可在C中加载的检查点:

 

python train_gpt2.py

 

该脚本会下载GPT-2(124M)模型,对单批数据进行10次过拟合迭代,运行多个生成步骤,最重要的是保存两个文件:1)gpt2_124M.bin文件,包含用于在C中加载的模型权重;2)以及gpt2_124M_debug_state.bin,包含包括input、target、logits及loss等更多调试状态,对于调试C代码、单元测试及确保能够与PyTorch参考实现完全匹配非常重要。现在我们可以使用这些权重进行初始化并在原始C代码中进行训练。首先编译代码:

 

make train_gpt2

 

在train_gpt2编译完成后即可运行:

 

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

generated: 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

step 40: train loss 4.377757 (took 1366.368000 ms)

 

现在的生成结果仅给出token ids,需要将其解码回文本形式:

 

<|endoftext|>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

 

 

参考链接:

https://twitter.com/karpathy/status/1777427947126936026"

https://github.com/karpathy/llm.c"

https://www.youtube.com/watch?v=zjkBMFhNj_g"