純C語言手搓GPT-2，前OpenAI、特斯拉高管新項(xiàng)目火了

作者：機(jī)器之心 2024-04-09 12:23:27

今天凌晨，前特斯拉 Autopilot 負(fù)責(zé)人、OpenAI 科學(xué)家 Andrej Karpathy 發(fā)布了一個(gè)僅用 1000 行代碼即可在 CPU/fp32 上實(shí)現(xiàn) GPT-2 訓(xùn)練的項(xiàng)目「llm.c」。

眾所周知，大語言模型還在快速發(fā)展，應(yīng)該有很多可以優(yōu)化的地方。我用純 C 語言來寫，是不是能優(yōu)化一大截？

也許很多人開過這樣的腦洞，現(xiàn)在有大佬實(shí)現(xiàn)了。

GitHub 鏈接：https://github.com/karpathy/llm.c

消息一出，立即引發(fā)了機(jī)器學(xué)習(xí)社區(qū)的熱烈討論，項(xiàng)目的 Star 量不到七個(gè)小時(shí)就沖上了 2000。有網(wǎng)友表示，大佬從零開始用 C 語言寫大模型只為好玩，我等只能膜拜：

llm.c 旨在讓大模型（LM）訓(xùn)練變得簡單 —— 使用純 C 語言 / CUDA，不需要 245MB 的 PyTorch 或 107MB 的 cPython。例如，訓(xùn)練 GPT-2（CPU、fp32）僅需要單個(gè)文件中的大約 1000 行干凈代碼（clean code），可以立即編譯運(yùn)行，并且完全可以媲美 PyTorch 參考實(shí)現(xiàn)。

Karpathy 表示，選擇從 GPT-2 開始，是因?yàn)樗?LLM 的鼻祖，是大語言模型體系首次以現(xiàn)代形式組合在一起，并且有可用的模型權(quán)重。

原始訓(xùn)練的實(shí)現(xiàn)在這里：https://github.com/karpathy/llm.c/blob/master/train_gpt2.c

你會(huì)看到，項(xiàng)目在開始時(shí)一次性分配所有所需的內(nèi)存，這些內(nèi)存是一大塊 1D 內(nèi)存。然后在訓(xùn)練過程中，不會(huì)創(chuàng)建或銷毀任何內(nèi)存，因此內(nèi)存占用量保持不變，并且只是動(dòng)態(tài)的，將數(shù)據(jù)批次流過。這里的關(guān)鍵在于手動(dòng)實(shí)現(xiàn)所有單個(gè)層的前向和后向傳遞，然后將它們串聯(lián)在一起。

例如，這里是 layernorm 前向和后向傳遞。除了 layernorm 之外，我們還需要編碼器、matmul、自注意力、gelu、殘差、softmax 和交叉熵?fù)p失。

「一旦你擁有了所有的層，接下來的工作只是將它們串在一起。講道理，寫起來相當(dāng)乏味和自虐，因?yàn)槟惚仨毚_保所有指針和張量偏移都正確排列，」Karpathy 評論道。

左：我們分配一個(gè) 1D 內(nèi)存數(shù)組，然后將所有模型權(quán)重和激活指向它。右：我們需要非常非常小心地進(jìn)行所有指針運(yùn)算。

一旦你有了前向 / 后向，其余部分（數(shù)據(jù)加載器、Adam 更新等）大多就不足為懼了。

不過，真正的樂趣現(xiàn)在才開始：Karpathy 表示，他現(xiàn)在正在逐層將其移植到 CUDA 上，以便提高效率，甚至期待能在 PyTorch 的合理范圍內(nèi)，但沒有任何嚴(yán)重的依賴關(guān)系 —— 現(xiàn)在工作已經(jīng)完成了幾層。所以這是一個(gè)非常有趣的 CUDA 練習(xí)。

對此，有網(wǎng)友表示：即使頂著指針 ptsd，我也能感受到這些代碼的美。

也有人說，這項(xiàng)目簡直就是完美的機(jī)器學(xué)習(xí)工程師在線面試答案。

從這開始，未來該項(xiàng)目的延伸會(huì)包括將精度從 fp32 降低到 fp16 / 以下，以及增加幾個(gè)層（例如 RoPE）以支持更現(xiàn)代的架構(gòu)，如 llama 2/mistral/gemma/ 等模型。

最后，Andrej Karpathy 表示，一旦項(xiàng)目穩(wěn)定起來，就會(huì)出關(guān)于從頭開始用 C 語言寫大模型的視頻。

llm.c 下一步的目標(biāo)包括：

直接的 CUDA 實(shí)現(xiàn)，讓速度更快，并且可能接近 PyTorch；
使用 SIMD 指令、x86 上的 AVX2 / ARM 上的 NEON（例如蘋果 M 系列芯片的電腦）來加速 CPU 版本；
更多新型架構(gòu)，例如 Llama2、Gemma 等。

看起來，想讓速度更快的目的沒有達(dá)到，這里不得不佩服 PyTorch 如今的效率。對于存儲庫，作者希望維護(hù)干凈、簡單的參考實(shí)現(xiàn)，以及可以接近 PyTorch 的更優(yōu)化版本，但代碼和依賴項(xiàng)只占一小部分。

使用方法

要使用 llm.c，首先要下載并 tokenize 數(shù)據(jù)集。tinyshakespeare 數(shù)據(jù)集的下載和 tokenize 速度最快：

python prepro_tinyshakespeare.py

輸出：

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

.bin 文件是 int32 數(shù)字的原始字節(jié)流，使用 GPT-2 tokenizer 標(biāo)記 token ID，或者也可以使用 prepro_tinystories.py tokenize TinyStories 數(shù)據(jù)集。

原則上，llm.c 到這一步已經(jīng)可以訓(xùn)練模型。然而，基線 CPU/fp32 參考代碼的效率很低，從頭開始訓(xùn)練這些模型不切實(shí)際。因此，這里使用 OpenAI 發(fā)布的 GPT-2 權(quán)重進(jìn)行初始化，然后再進(jìn)行微調(diào)，所以必須下載 GPT-2 權(quán)重并將它們保存為可以在 C 中加載的檢查點(diǎn)：

python train_gpt2.py

該腳本將下載 GPT-2 (124M) 模型，對單批數(shù)據(jù)進(jìn)行 10 次迭代的過擬合，運(yùn)行幾個(gè)生成步驟，最重要的是，它將保存兩個(gè)文件：

gpt2_124M.bin 文件，包含在 C 語言中加載模型所需的權(quán)重；
gpt2_124M_debug_state.bin 文件，包含更多調(diào)試狀態(tài)：輸入、目標(biāo)、logits 和損失。這對于調(diào)試 C 語言代碼、單元測試以及確保 llm.c 與 PyTorch 參考實(shí)現(xiàn)完全可媲美非常重要。

現(xiàn)在，使用 gpt2_124M.bin 中的模型權(quán)重進(jìn)行初始化并使用純 C 語言進(jìn)行訓(xùn)練，首先編譯代碼：

make train_gpt2

這里可以查看 Makefile 及其注釋。它將嘗試自動(dòng)檢測 OpenMP 在當(dāng)前系統(tǒng)上是否可用，這對于以極低的代碼復(fù)雜性成本加速代碼非常有幫助。編譯 train_gpt2 后，運(yùn)行：

OMP_NUM_THREADS=8 ./train_gpt2

這里應(yīng)該根據(jù) CPU 的核心數(shù)量來調(diào)整線程數(shù)量。該程序?qū)⒓虞d模型權(quán)重、token，并使用 Adam 運(yùn)行幾次迭代的微調(diào) loop，然后從模型生成樣本。在 MacBook Pro (Apple Silicon 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 ID，還需要將其解碼回文本。這一點(diǎn)可以很容易地用 C 語言實(shí)現(xiàn)，因?yàn)榻獯a非常簡單，可以使用 tiktoken：

import tiktoken
enc = 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 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

值得注意的是，這里沒有嘗試調(diào)整微調(diào)超參數(shù)，因此很可能還有大幅改進(jìn)的空間，特別是在訓(xùn)練時(shí)間更長的情況下。

附上一個(gè)簡單的單元測試，以確保 C 代碼與 PyTorch 代碼一致。編譯并運(yùn)行：

make test_gpt2
./test_gpt2

這里加載 gpt2_124M_debug_state.bin 文件，運(yùn)行前向傳遞，將 logits 和損失與 PyTorch 參考實(shí)現(xiàn)進(jìn)行比較，然后使用 Adam 進(jìn)行 10 次迭代訓(xùn)練，確保損失可與 PyTorch 參考實(shí)現(xiàn)媲美。