支持AMD GPU的llm.c

anthonix/llm.c: LLM training in simple, raw C/HIP for AMD GPUs (github.com)

llm.c for AMD devices

This is a fork of?Andrej Karpathy's llm.c?with support for AMD devices.

性能

在單個7900 XTX顯卡上使用默認設置，目前的訓練步驟耗時約為79毫秒，相比PyTorch的夜間版本（2.4.0.dev20240513）的約97毫秒，以及tinygrad的約440毫秒來說，表現更優。

對于多GPU訓練，在裝有四個7900 XTX顯卡的機器上，吞吐量達到了每秒約210,000個令牌。

更新（2024年5月28日）：在單個7900 XTX顯卡上，快速注意力分支的訓練步驟時間已經降低到58.340831毫秒，或者在四個7900 XTX顯卡上達到了每秒318,777個令牌的吞吐量。目前正在研究雙緩沖技術以進一步推動性能提升。

狀態

train_gpt2_fp32（基線，最小改動）：使用32位浮點數（FP32）訓練GPT-2模型（基線版本，僅進行最小改動）
train_gpt2 with BF16（基線，最小改動）：使用半精度浮點數（BF16）訓練GPT-2模型（基線版本，僅進行最小改動）
train_gpt2 with BF16 and multiple GPUs：使用半精度浮點數（BF16）并在多個GPU上訓練GPT-2模型
RDNA3 優化內核（進行中）：針對RDNA3架構優化的內核（仍在開發中）
CDNA3 優化內核：針對CDNA3架構優化的內核（具體狀態未提及）

快速入門（AMD目標）

安裝ROCm 6.1.1，檢出倉庫，并執行以下步驟：

pip install -r requirements.txt?（安裝所需的依賴項）
python prepro_tinyshakespeare.py?（預處理tinyshakespeare數據集）
export HF_ENDPOINT=https://hf-mirror.com
python train_gpt2.py?（（此步驟可能用于生成某種訓練數據或配置，但具體細節未在給定指令中明確））
make train_gpt2amd?（編譯AMD特定版本的GPT-2訓練程序）

ROCm 6.0.2報錯：

ld.lld: error: unable to find library -ldevice_gemm_operations
ld.lld: error: unable to find library -ldevice_other_operations
ld.lld: error: unable to find library -lstdc++
clang: error: linker command failed with exit code 1 (use -v to see invocation)
make: *** [Makefile:285：train_gpt2amd] 錯誤 1

不知道ROCm 6.1.1是否能編譯成功。

?6. ./train_gpt2amd?（運行AMD特定版本的GPT-2訓練程序）

[原始README]

karpathy/llm.c

在簡單、純C/CUDA中進行LLM（大型語言模型）訓練。無需245MB的PyTorch或107MB的cPython。在單個文件llm.c/train_gpt2.c中，使用CPU和fp32（32位浮點數）訓練GPT-2大約需要1,000行干凈的代碼，而在llm.c/train_gpt2.cu中使用GPU訓練大約需要3,000行代碼（增加了CUDA內核）。代碼立即編譯并運行，它與PyTorch的參考實現完全匹配，并且當前速度略快于（編譯后的）PyTorch（使用bf16、torch編譯和flash attention）。我選擇GPT-2作為第一個工作示例，因為它是LLM的鼻祖，也是現代堆棧首次結合在一起的例子。

我們當前的目標是復現GPT-2。要了解當前正在進行的工作的概述，請參閱最新的State of the Union帖子。

2024年5月28日更新：一個有用的近期帖子可能是“在llm.c中用90分鐘和20美元復現GPT-2（124M）”，我在其中詳細說明了從零開始復現124M/350M模型的GPT-2微系列的步驟。展示啟動命令的文件本身是run124M.sh和run350M.sh。

我希望這個倉庫只維護C和CUDA代碼。這個倉庫到其他語言的移植非常受歡迎，但應該在單獨的倉庫中完成，然后我很樂意在下面的“顯著分支”部分鏈接到它們，就像我在llama2.c分支中所做的那樣。

快速入門（GPU，速度慢但穩定且適合學習）

對于“我不在乎其他任何事情，我只想訓練，并且我有GPU”的這部分用戶。請運行以下命令：

pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py
make train_gpt2fp32cu
./train_gpt2fp32cu

以上命令（1）會下載tinyshakespeare數據集，并使用GPT-2的Tokenizer進行分詞處理，（2）下載并保存GPT-2（124M）的權重，（3）在C/CUDA中從這些權重初始化，并在tinyshakespeare數據集上使用AdamW優化器進行一輪（epoch）的訓練（使用批量大小4，上下文長度1024，總共74步），評估驗證損失，并生成一些文本樣本。請注意，在這個快速入門中，我們使用的是CUDA代碼的fp32版本train_gpt2_fp32.cu。在下一節中，我們將介紹當前“主流”的train_gpt2.cu，它使用混合精度，運行速度大約快2倍。

快速入門（GPU，最前沿的優化）

我想看到它運行得更快。在這種情況下，切換到我們最主要的、優化度最高的 train_gpt2.cu。運行如下命令：

pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py
make train_gpt2cu
./train_gpt2cu

如果你額外安裝了cuDNN（請參見下面的CUDA部分），你可以通過flash attention實現更快的速度。調整make命令，如下編譯帶有cudnn/flash attention的版本：

make train_gpt2cu USE_CUDNN=1
./train_gtp2cu

這段話的意思是，如果你已經安裝了cuDNN，那么在編譯和運行一個名為train_gpt2cu的程序時，可以通過設置USE_CUDNN=1來啟用cuDNN的支持，從而利用cuDNN提供的加速功能來提高程序的運行速度。這里的make是一個構建工具，用于自動化編譯過程，而./train_gpt2cu則是運行編譯后的程序。
請注意，默認的批量大小非常小（4）。如果你的GPU有足夠的內存，我建議你將其增加到例如32：

./train_gtp2cu -b 32

我的標準單GPU "生產" 運行（例如使用A100 40GB）不是訓練TinyShakespeare，而是訓練TinyStories，示例如下：

python dev/data/tinystories.py
make train_gtp2cu USE_CUDNN=1
./train_gtp2cu -i dev/data/tinystories/TinyStories_train.bin \-j dev/data/tinystories/TinyStories_val.bin \-v 250 -s 250 -g 144 -o stories.log -b 32

-i 標志是輸入數據的通配符模式，`-j` 是驗證數據。另外，我將驗證損失和采樣的頻率減少到每250步，并且在采樣階段采樣144個tokens（大約能裝下一篇故事），批量大小為32。
如果你想要訓練實際的、真實的預訓練數據，查看最近添加的對 fineweb數據集的支持。與上面的數據集不同，這里的訓練/驗證tokens?不是放在一個.bin文件中，而是現在有多個數據分片。這里有一個示例：

# 將FineWeb數據以1億個標記的分片寫入到dev/data/fineweb10B
python dev/data/fineweb.py -s 100000000
# 編譯并運行
./train_gtp2cu -i "dev/data/fineweb10B/fineweb_train_*.bin" \-j "dev/data/fineweb10B/fineweb_val_*.bin" \-v 250 -s 250 -g 144 -o fineweb.log -b 32

其中，你會注意到使用了通配符 * 來匹配所有的訓練分片。

快速入門（多GPU）

很好，讓我們更進一步。我們將使用MPI和NCCL來進行多GPU訓練。上面部分的內容依然適用，但需要做以下改變：

# 安裝MPI的示例：
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev# 運行命令現在需要以?mpirun?開頭：
mpirun -np <你機器上的GPU數量> ./train_gpt2cu

在最后的命令中，替換為你想要運行的GPU數量。上面部分討論的所有標志在這里也適用。

快速開始（CPU）

“我是如此貧窮以至于連GPU都沒有”的部分。你仍然可以進行訓練！但你不會走得太遠。你仍然可以微調一個GPT-2小模型（1.24億參數模型）以輸出像莎士比亞式的文本，作為示例：

pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py
make train_gpt2
OMP_NUM_THREADS=8 ./train_gpt2

上面的行（1）下載了tinyshakespeare?數據集，使用GPT-2的分詞器對其進行分識，（2）下載并保存GPT-2（124M）的權重，（3）在C中從它們初始化并在tineshakespeare上訓練40步驟使用AdamW（使用批量大小4，上下文長度只有64），評估驗證損失，并抽取一些文本。誠實地講，除非你有一個強大的CPU（并且可以在啟動命令中增加OMP線程的數量），你在CPU上訓練大型語言模型（LLMs）不會走得太遠，但它可能是一個不錯的演示/參考。

訓練：更多細節

在`/dev/data/(dataset).py`的數據文件負責下載、分詞并將分詞保存到文件中。例如，當你運行：

python dev/data/tinyshakespeare.py

我們下載并分詞tinyshakespeare數據集。這個輸出看起來像這樣：

writing 32,768 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_val.bin
writing 305,260 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_train.bin

.bin文件包含一個短的頭部（1024字節），隨后是一個流式的tokens以uint16格式，表明了用GPT-2分詞器的分詞id。更多數據集可以在`/dev/data`找到。
原則上，一旦我們得到了tokens，我們就準備在這里開始訓練模型。然而，當前的代碼還不能從零開始訓練（很快就會加入），所以我們從OpenAI發布的預訓練模型初始化訓練并進行微調。為此，我們需要下載GPT-2的權重并將其作為我們可以在C語言中加載的檢查點來保存。這就是當你運行以下腳本時發生的事：

python train_gpt2.py

你會認出這段代碼，它來自nanoGPT，是一個簡單的PyTorch中的GPT-2參考實現。這個腳本將下載GPT-2 (124M)模型，對單批數據進行了10次迭代的超過擬合，運行了幾步生成，并且更重要的是它將保存三個文件：1）`gpt2_124M.bin`文件，它包含了原始模型權重以便在C語言中加載，2）`gpt2_124M_debug_state.bin`文件，這也包含了更多的調試狀態：輸入，目標，邏輯和損失（對于調試和單元測試很有用），最后3）`gpt2_tokenizer.bin`文件，它存儲了GPT-2分詞器的詞匯表，將分詞id轉換為UTF-8編碼字符串片段的字節序列。文件還保存了上述的fp32版本，以及它們的bfloat16版本以供混合精度訓練。我們現在可以用這些模型權重來初始化并繼續在原始C語言中訓練。然后我們用`make`命令來編譯訓練程序。目前有三個并行實現：

# 簡單的，CPU，參考代碼版本
make train_gpt2
# 單GPU fp32 CUDA版本
make train_gpt2fp32cu
# 多GPU混合精度CUDA版本
make train_gpt2cu

你可以查閱`Makefile`及其注釋。它會嘗試自動探測很多工具和庫（例如：cuDNN, OpenMP, OpenMPI, nvcc），你要盡量獲得盡可能多的勾號。比如當我在我配置完善的機器上運行`make train_gpt2cu USE_CUDNN=1`，我們看到：

? cuDNN found, will run with flash-attention
? OpenMP found
? OpenMPI found, OK to train with multiple GPUs
? nvcc found, including GPU/CUDA support

有些人在Ubuntu上編譯時遇到問題，請查看Issue 19，簡而言之就是你想修改`CFLAGS`：

# 首先嘗試這個
CFLAGS="-Ofast -fno-finite-math-only -Wno-unused-result -march=native" make train_gpt2
# 其次嘗試這個
CFLAGS="-O3 -Wno-unused-result -march=native" make train_gpt2

一旦編譯好了二進制文件，我們就可以運行它。例如最簡單的CPU參考版本運行如下：

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

我喜歡Netflix的出現，很明顯模型的訓練過往仍在影響它。我沒有嘗試調整微調的超參數，所以這個結果很可能還可以大幅度提高。我還注意到，不同的平臺（例如MacOS/Linux）將會（遺憾地）給出非常微小的不同結果，所以可能不要期望得到上文提供的確切的數字或生成結果。

最后，代碼還在變動中。如果發生任何你沒預料到或之前運行正常的奇怪事情，請嘗試`git pull`，重新運行所有上面的命令，回到這個README文件參考，等等。

測試

我還附上了一個簡單的單元測試，以確保我們的 C 代碼與 PyTorch 代碼一致。以 CPU 為例，編譯并且執行如下：

make test_gpt2
./test_gpt2

這將加載?gpt2_124M_debug_state.bin?文件，執行一個前向傳遞，與 PyTorch 參考實現比較 logits 和 loss，然后進行 10 次迭代的 Adam 訓練確保損失與 PyTorch 匹配。要測試 GPU 版本，我運行：

# fp32 測試（不支持 cudnn）
make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
# 混合精度 cudnn 測試
make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu

教程

我在這里附上了一個非常小的教程，在?doc/layernorm/layernorm.md。這是一個實現 GPT-2 模型的單層，layernorm 層的簡單分步指導。這是理解 C 中是如何實現層的一個好的起點。

CUDA

整個訓練循環也在一個文件中使用純CUDA實現，但是內核的優化還在進行中。目前，我們的速度略微超過了PyTorch Nightly的速度。我們組織代碼的方式是，在`dev/cuda`文件夾中收集了越來越多的復雜程度遞增的內核，詳見dev/cuda/README.md。然后，我們將最好的內核復制粘貼到單一訓練文件`train_gpt2cu.cu`中的主要訓練循環中。

WIP警告，2024年4月23日。我們合并了第一個版本的混合精度訓練代碼。我將fp32版本的檢查點備份到包含`_fp32`文件名的單獨文件中，并希望保留這個版本在倉庫的根目錄中，因為它1）不需要最新的CUDA，更有可能編譯和更加便于移植，2）它更簡單，并且充當參考。事實上，我們想讓fp32版本朝著純CUDA的方向發展（例如，默認情況下甚至不調用cuBLAS），用作教育參考，甚至可能是CUDA課程的一個內核。從現在開始，與速度有關的"主線"開發將轉移到train_gpt2.cu文件，該文件包含混合精度訓練。

在下面的描述中，我現在默認使用fp32版本，因為它當前更加便攜和穩定，然后在最后我將介紹新的混合精度版本。

正確性。首先，我們可以做10次訓練迭代并驗證我們的代碼是否與PyTorch完全匹配和再現數字：

make test_gpt2fp32cu
./test_gpt2fp32cu

這會打印出`overall okay: 1`。因此，前向激活、后向梯度和10次迭代的各個損失值都完全匹配。

訓練。在單GPU上以fp32進行訓練：

make train_gpt2fp32cu
./train_gpt2fp32cu

這將加載tiny_shakespeare數據集的驗證和訓練劃分。在默認設置B=4，T=1024下，有8個驗證批次和74個訓練批次。該腳本目前配置為以1e-4的學習速率進行單輪微調，并在此過程中評估驗證性能和生成樣本，例如：

step 1/74: train loss 4.367631 (80.639749 ms)
step 2/74: train loss 4.031242 (77.378867 ms)
step 3/74: train loss 4.034144 (77.315861 ms)
step 4/74: train loss 3.859865 (77.357575 ms)
...
step 72/74: train loss 3.085081 (78.850895 ms)
step 73/74: train loss 3.668018 (78.197064 ms)
step 74/74: train loss 3.467508 (78.009975 ms)
val loss 3.516490
generating:
---
?Where will you go?
I take you wherefore I can, myself, and must.
I cast off my beak, that I may look him up on the point;
For on his rock shall he be opencast.<|endoftext|>My little nephew:
Keep on with me, my

這在我的A100上大約運行了~10秒。我們可以這樣與naive PyTorch進行比較，我們開啟了`torch.compile`和使用TensorCores，它使用的是tf32類型：

python train_gpt2.py --write_tensors 0 --sequence_length 1024 --batch_size 4 --compile 1 --tensorcores 1

編譯（第一次迭代）大約需要~27秒，但之后在我的A100上目前的運行速度約為每次迭代~80ms。

混合精度。新的CUDA混合精度版本，未來大部分開發將在此進行，是train_gpt2.cu，以及它的測試test_gpt2.cu。在這里，許多計算以較低精度格式（fp16或bf16）進行，這使我們能夠以非常快的速度運行（約為上面TF32性能的2倍）。注意，我描述的基線實現作為`fp32`，但更精確地說，實際上是`tf32`（TensorFloat32）。訓練和測試的命令都一樣，只需省略fp32部分：

make train_gpt2cu
./train_gpt2cumake test_gpt2cu
./test_gpt2cu

如果您有最新的CUDA，應該期望它可以編譯OK，并且應該看到性能大幅改進。

Flash Attention。截至2024年5月1日，我們現在支持來自cuDNN的Flash Attention。因為cuDNN使編譯時間從幾秒增加到約一分鐘，并且這個代碼路徑現在非常?新穎，目前默認情況下是禁用的。您可以通過以下方式編譯來啟用它：

make train_gpt2cu USE_CUDNN=1

這將嘗試使用cudnn進行編譯并運行。您必須在您的系統上安裝cuDNN。通過apt-get安裝的cuDNN安裝說明將獲取默認的cuDNN包集。對于最小安裝來說，cuDNN dev包是足夠的，例如在Ubuntu 22.04上為CUDA 12.x：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install libcudnn9-dev-cuda-12

除此之外，您還需要cuDNN前端，但這只是頭文件。只需將倉庫克隆到您的磁盤即可。如果您將其放在其他地方，將`CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include`添加到`make`命令行中。

多GPU訓練。截至2024年4月26日，現在還支持使用MPI和NCCL進行多GPU訓練。確保您安裝了MPI，例如在Linux上：

sudo apt install openmpi-bin openmpi-doc libopenmpi-dev

然后：

make train_gpt2cu
mpirun -np <GPU數量> ./train_gpt2cu

fp32版本的代碼不支持多GPU。這是因為我們希望GPT-2 fp32版本成為CUDA優化課程的一個不錯的教育終點。混合精度版本是我們進行前沿開發的版本，因此這是支持多GPU訓練的版本。

實驗/掃描

現在基本的 argparse 和日志功能已存在于 .cu 腳本中，我們可以開始進行第一輪學習率掃描。目前這還相當手動，但是下面記錄了一個示例過程，展示了在一臺擁有4個GPU的機器上針對TinyStories數據集進行學習率掃描的過程。在你當然已經使用?chmod u+x sweep.sh?命令賦予了sweep.sh腳本執行權限后，運行一個名為?sweep.sh?的shell腳本：

#!/bin/bashlearning_rates=(3e-5 1e-4 3e-4 1e-3)for i in {0..3}; doexport CUDA_VISIBLE_DEVICES=$iscreen -dmS "tr$i" bash -c "./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -l ${learning_rates[$i]} -o stories$i.log"
done# 你可以使用以下命令關閉這些屏幕會話
# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit

這個例子打開了4個screen會話，并使用不同的學習率運行四個命令。這會將所有損失寫入日志文件?stories$i.log，你可以按照自己的意愿在Python中進行繪制。解析和繪制這些日志文件的一個快速例子可以在?dev/vislog.ipynb?中找到。

代碼倉庫理念

關于我希望這個代碼倉庫`llm.c`能夠成為的幾點想法：
首先，我希望`llm.c`能成為一個教學場所。舉個例子，我們的`dev/cuda`文件夾是一個包含了各種手寫、文檔齊全的內核庫，從最簡單的內核開始，到更復雜／更快速的內核。如果您有帶有不同權衡的新內核，請隨時貢獻至此。
話雖如此，我也希望`llm.c`能夠非常快速，甚至在實際中用于訓練網絡。比如，首先，我們應該能復現大型GPT-2（1.6B）的訓練過程。這需要我們整合各種最快的內核，包含使用如cuBLAS、cuBLASLt、CUTLASS、cuDNN等庫。我同樣認為這樣做對于確立一個專家級的上限，并作為一種度量單位，具有教育意義。比如，你可以說你手寫的內核速度達到了cuBLAS的80%等。然后你可以選擇進行一個超快的運行，或者你可以選擇 "拖放 "任何你想使用的手動內核，并用那些運行。
然而，作為一個限制，我希望保持根目錄下的主線`llm.c`簡單且可讀。如果有一個PR能夠改進性能2%，但它“花費了”500行復雜的C代碼，可能還有一些非主流的第三方依賴，我可能會拒絕這個PR，因為復雜性不值得。具體的一個例子 - 讓cuBLAS成為根訓練循環的默認矩陣乘法是明智的：它讓主線代碼快了很多，它是一行易于理解的代碼，而且是一個非常常見的依賴。在這一側，我們可以在`dev/cuda`中有與cuBLAS競爭的手寫實現。
最后，對于項目根目錄中包含主要／默認文件的部分，我會對復雜性更加敏感。相比之下，`dev/`文件夾更像是一個草稿空間，供我們開發內核或類的庫，分享有用或相關的或教育性代碼，其中一些代碼是可以接受的（局部）復雜性。