純C語言手搓GPT-2，前OpenAI、特斯拉高管新專案火了

机器之心發表於2024-04-09
C語言GPTOpenAI

「Real men program in C.」

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

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

圖片

今天凌晨，前特斯拉 Autopilot 負責人、OpenAI 科學家 Andrej Karpathy 釋出了一個僅用 1000 行程式碼即可在 CPU/fp32 上實現 GPT-2 訓練的專案「llm.c」。

GitHub 連結：https://github.com/karpathy/llm.c

訊息一出，立即引發了機器學習社群的熱烈討論，專案的 Star 量不到七個小時就衝上了 2000。有網友表示，大佬從零開始用 C 語言寫大模型只為好玩，我等只能膜拜：

圖片

llm.c 旨在讓大模型（LM）訓練變得簡單 —— 使用純 C 語言 / CUDA，不需要 245MB 的 PyTorch 或 107MB 的 cPython。例如，訓練 GPT-2（CPU、fp32）僅需要單個檔案中的大約 1000 行乾淨程式碼（clean code），可以立即編譯執行，並且完全可以媲美 PyTorch 參考實現。

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

原始訓練的實現在這裡：https://github.com/karpathy/llm.c/blob/master/train_gpt2.c

你會看到，專案在開始時一次性分配所有所需的記憶體，這些記憶體是一大塊 1D 記憶體。然後在訓練過程中，不會建立或銷燬任何記憶體，因此記憶體佔用量保持不變，並且只是動態的，將資料批次流過。這裡的關鍵在於手動實現所有單個層的前向和後向傳遞，然後將它們串聯在一起。

圖片

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

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

圖片 左：我們分配一個 1D 記憶體陣列，然後將所有模型權重和啟用指向它。右：我們需要非常非常小心地進行所有指標運算。

一旦你有了前向 / 後向，其餘部分（資料載入器、Adam 更新等）大多就不足為懼了。

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

對此，有網友表示：即使頂著指標 ptsd，我也能感受到這些程式碼的美。

圖片

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

從這開始，未來該專案的延伸會包括將精度從 fp32 降低到 fp16 / 以下，以及增加幾個層（例如 RoPE）以支援更現代的架構，如 llama 2/mistral/gemma/ 等模型。

最後，Andrej Karpathy 表示，一旦專案穩定起來，就會出關於從頭開始用 C 語言寫大模型的影片。

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

  • 直接的 CUDA 實現，讓速度更快，並且可能接近 PyTorch；

  • 使用 SIMD 指令、x86 上的 AVX2 / ARM 上的 NEON（例如蘋果 M 系列晶片的電腦）來加速 CPU 版本；

  • 更多新型架構，例如 Llama2、Gemma 等。

看起來，想讓速度更快的目的沒有達到，這裡不得不佩服 PyTorch 如今的效率。對於儲存庫，作者希望維護乾淨、簡單的參考實現，以及可以接近 PyTorch 的更最佳化版本，但程式碼和依賴項只佔一小部分。

使用方法

要使用 llm.c，首先要下載並 tokenize 資料集。tinyshakespeare 資料集的下載和 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 數字的原始位元組流，使用 GPT-2 tokenizer 標記 token ID，或者也可以使用 prepro_tinystories.py tokenize TinyStories 資料集。

原則上，llm.c 到這一步已經可以訓練模型。然而，基線 CPU/fp32 參考程式碼的效率很低，從頭開始訓練這些模型不切實際。因此，這裡使用 OpenAI 釋出的 GPT-2 權重進行初始化，然後再進行微調，所以必須下載 GPT-2 權重並將它們儲存為可以在 C 中載入的檢查點：

python train_gpt2.py

該指令碼將下載 GPT-2 (124M) 模型，對單批資料進行 10 次迭代的過擬合，執行幾個生成步驟，最重要的是，它將儲存兩個檔案：

  • gpt2_124M.bin 檔案，包含在 C 語言中載入模型所需的權重；

  • gpt2_124M_debug_state.bin 檔案，包含更多除錯狀態：輸入、目標、logits 和損失。這對於除錯 C 語言程式碼、單元測試以及確保 llm.c 與 PyTorch 參考實現完全可媲美非常重要。

現在，使用 gpt2_124M.bin 中的模型權重進行初始化並使用純 C 語言進行訓練，首先編譯程式碼：

make train_gpt2

這裡可以檢視 Makefile 及其註釋。它將嘗試自動檢測 OpenMP 在當前系統上是否可用，這對於以極低的程式碼複雜性成本加速程式碼非常有幫助。編譯 train_gpt2 後，執行：

OMP_NUM_THREADS=8 ./train_gpt2

這裡應該根據 CPU 的核心數量來調整執行緒數量。該程式將載入模型權重、token，並使用 Adam 執行幾次迭代的微調 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 語言實現，因為解碼非常簡單，可以使用 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 notflourish against the earlocks of
Allay

值得注意的是，這裡沒有嘗試調整微調超引數，因此很可能還有大幅改進的空間，特別是在訓練時間更長的情況下。

附上一個簡單的單元測試，以確保 C 程式碼與 PyTorch 程式碼一致。編譯並執行：

make test_gpt2
./test_gpt2

這裡載入 gpt2_124M_debug_state.bin 檔案，執行前向傳遞，將 logits 和損失與 PyTorch 參考實現進行比較，然後使用 Adam 進行 10 次迭代訓練，確保損失可與 PyTorch 參考實現媲美。

圖片

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

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

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

圖片

這次新專案的配套影片什麼時候出？我們都很期待。

參考內容：

https://news.ycombinator.com/item?id=39973467

https://twitter.com/karpathy/status/1777427944971083809

相關文章