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

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

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

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

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

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

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

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

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

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

「一旦你擁有了所有的層，接下來的工作只是將它們串在一起。講道理，寫起來相當乏味和自虐，因為你必須確保所有指針和張量偏移都正確排列， 」Karpathy 評論道。

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

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

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

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

也有人說，這項目簡直就是完美的機器學習工程師在線面試答案。

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

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

llm.c 下一步的目標包括：

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

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

使用方法

要使用 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 標記 token ID，或者也可以使用 prepro_tinystories.py tokenize TinyStories 數(shù)據(jù)集。

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

python train_gpt2.py

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

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

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

make train_gpt2

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

OMP_NUM_THREADS=8 ./train_gpt2

這里應該根據(jù) CPU 的核心數(shù)量來調(diào)整線程數(shù)量。該程序?qū)⒓虞d模型權(quán)重、token，并使用 Adam 運行幾次迭代的微調(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，還需要將其解碼回文本。這一點可以很容易地用 C 語言實現(xiàn)，因為解碼非常簡單，可以使用 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ù)，因此很可能還有大幅改進的空間，特別是在訓練時間更長的情況下。

附上一個簡單的單元測試，以確保 C 代碼與 PyTorch 代碼一致。編譯并運行：

make test_gpt2
./test_gpt2

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

最后，Karpathy 還附上了一個簡單的教程。這是一個簡單的分步指南，用于實現(xiàn) GPT-2 模型的單層（layernorm 層），可以幫助你理解如何用 C 語言實現(xiàn)語言模型。

教程地址：doc/layernorm/layernorm.md

我們知道，最近 Andrej Karpathy 沉迷于制作教程。去年 11 月，他錄制的《大語言模型入門》在 YouTube 上吸引了很多人觀看。

這次新項目的配套視頻什么時候出？我們都很期待。