从 llm.c 到 llm.rs(一)

llm / python / gpt

我一直想看看大模型是怎么玩的,但是根据我的了解,这是一门学科, 不是简单一两天可以玩的, 平时时间本来就不是很多的我自然是没有可能去系统学习这一块东西, 所以一直也就是放着。近来知道 karpathy 做了一个 1000 行的 gpt-2 llm.c, 我觉得这是个不错的起点。

检查头文件

给自己定了一个小目标,把这 1000 行代码吃透,吃透的方法是用 rust 重写, 所以也就有了这个标题。那么第一步我们要把 llm.c 跑通, 这样我们才有东西可以参考。 下载整个目录,根据 readme 文件,最主要的文件是 train_gpt2.c。 那么我们先看看 include 的头文件:

/*
This file trains the GPT-2 model.
This version is the clean, minimal, reference. As such:
- it runs on CPU.
- it does not make the code too complex; it is readable.
- it does not use any processor-specific instructions, intrinsics and such.
- it _does_ use a few OpenMP pragmas because this is a large speedup at very low cost
There will be other versions of this code that specialize it and make it fast.
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <string.h>
#include <unistd.h>
#ifdef OMP
#include <omp.h>
#endif

这里用到了 unistd.h,也就是 unix 的 std, 那么这个程序能不能在 windows 上跑起来是个问题,不过我们的目标是跑起来看看结果, 不是为了解决这个程序在 windows 上运行,所以我打算用手上的一台 mac 试试。

准备 mac 环境

注解和 readme 都有提到,用 OpenMP 可以让程序跑得快一点,那我们先装上 OpenMP:

brew install llvm libomp

不过 karpathy 的 brew 配置可能和我不同,makefile 文件需要修改一下路径, 才能用到 OpenMP,不然的话,编译是没有问题,但是就吃不到 OpenMP 的加成, 会有一个警告,修改后的 makefile 如下:

# Check if the libomp directory exists
# path modified
ifeq ($(shell [ -d /usr/local/opt/libomp/lib ] && echo "exists"), exists)
  # macOS with Homebrew and directory exists
  CFLAGS += -Xclang -fopenmp -DOMP
  # path modified
  LDFLAGS += -L/usr/local/opt/libomp/lib
  LDLIBS += -lomp
  # path modified
  INCLUDES += -I/usr/local/opt/libomp/include
  $(info NICE Compiling with OpenMP support)
else
  $(warning OOPS Compiling without OpenMP support)
endif

三个路径都需要修改,修改完成后编译的结果如下:

% make train_gpt2
NICE Compiling with OpenMP support
cc -O3 -Ofast -Wno-unused-result -Xclang -fopenmp -DOMP  -I/usr/local/opt/libomp/include  -L/usr/local/opt/libomp/lib train_gpt2.c -lm -lomp -o train_gpt2

接下来我们要根据 readme 去下载那个「小小莎士比亚」了,不过我这里似乎缺了一些库, 还要先安装一下

python3 -m venv myenv
source myenv/bin/activate
# tinyshakespeare
pip install requests tqdm tiktoken numpy
# train_gpt2
pip install torch transformers

然后就可以下载了:

% python prepro_tinyshakespeare.py
Downloading https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt to data/tiny_shakespeare.txt...
data/tiny_shakespeare.txt: 1.06MiB [00:00, 1.73MiB/s]                                     
Saved 32768 tokens to data/tiny_shakespeare_val.bin
Saved 305260 tokens to data/tiny_shakespeare_train.bin

接下来,我们还是要用 python,根据描述,执行这个脚本会下载 gpt-2 的模型, 并生成原始的权重文件和调试态文件,供 c 使用

% python train_gpt2.py    
using device: cpu
loading weights from pretrained gpt: gpt2
config.json: 100%|████████████████████████████████████████| 665/665 [00:00<00:00, 756kB/s]
model.safetensors: 100%|████████████████████████████████| 548M/548M [09:18<00:00, 982kB/s]
generation_config.json: 100%|█████████████████████████████| 124/124 [00:00<00:00, 681kB/s]
loading cached tokens in data/tiny_shakespeare_val.bin
wrote gpt2_124M.bin
wrote gpt2_124M_debug_state.bin
iteration 0, loss: 5.269998073577881
...
iteration 9, loss: 0.376643568277359
<|endoftext|>"If we want to die, we have to die at the front"---------------

接下来跑 c 编译出的程序就是了:

% OMP_NUM_THREADS=4 ./train_gpt2
[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.252019
step 0: train loss 5.356185 (took 20010.307000 ms)
...
step 9: train loss 4.199315 (took 14724.552000 ms)
val loss 4.425501
step 10: train loss 4.288396 (took 13275.060000 ms)
...step 19: train loss 3.371637 (took 13535.751000 ms)
val loss 4.248619
generated: 50256 40 373 523 24822 10785 25 198 1639 561 599 539 329 1775 338 83 198 817 630 268 3675 345 11 705 4246 292 523 1290 1165 30 628 50256 46126 1268 40 2937 25 198 50 1740 6309 1201 1315 2670 318 198 14618 284 10598 25 198 1135 1826 262 48689 286 14494 7027 6 954 36368 30 5875 502 
step 20: train loss 3.880941 (took 13205.788000 ms)
...
step 29: train loss 3.812844 (took 13362.990000 ms)
val loss 4.110466
step 30: train loss 4.028022 (took 13669.109000 ms)
...
step 39: train loss 3.970822 (took 16787.196000 ms)
val loss 4.107841
generated: 50256 16773 18135 2238 11 198 12211 25 12226 13 628 50256 12211 25 314 466 407 8812 1958 616 778 4594 333 628 50256 12211 25 314 25511 257 467 88 198 5962 25 28790 284 616 9482 20246 351 2046 26 26852 3723 379 262 2046 11 198 14574 4151 11 511 1182 11 511 627 359 15224 13 198 12211 25 
step 40: train loss 4.377796 (took 18130.622000 ms)

我的 mac 比 karpathy 的慢多少,大家都知道了,生成最后面的一堆数字, 就是跑出来的结果,我们写一个 parse-result.py 解码一下:

import tiktoken
enc = tiktoken.get_encoding("gpt2")
gen = "50256 16773 18135 2238 11 198 12211 25 12226 13 628 50256 12211 25 314 466 407 8812 1958 616 778 4594 333 628 50256 12211 25 314 25511 257 467 88 198 5962 25 28790 284 616 9482 20246 351 2046 26 26852 3723 379 262 2046 11 198 14574 4151 11 511 1182 11 511 627 359 15224 13 198 12211 25"
print(enc.decode(list(map(int, gen.split()))))

跑一下看看是什么:

% python parse-result.py 
<|endoftext|>Come interfereoo,
Second: excuse.

<|endoftext|>Second: I do not terrify my prasteur

<|endoftext|>Second: I foresee a goy
First: pleasing to my burning burns with fire; gleaming at the fire,
Their eye, their head, their quill jacket.
Second:

看上去有些乱七八糟,不过既然没有报错,那应该 llm.c 我们已经跑通了。

总结

现在我们已经有了一个参照物,如果你和我一样也想把 llm.c 跑起来, 那么希望本篇对你有些帮助,总体上看来,需要一些 python 环境的配置知识, 以及对 makefile 和 brew 的一些了解,原文的 readme 忽略了这些细节。