五年后的今天，訓(xùn)練GPT-2只需不到700刀、24小時(shí)，Karpathy又整新活
共 15971字，需瀏覽 32分鐘
·
2024-07-12 22:00
極市導(dǎo)讀
論老黃賣鏟子的技術(shù)含量。 >>加入極市CV技術(shù)交流群，走在計(jì)算機(jī)視覺(jué)的最前沿
2019 年 2 月，OpenAI 發(fā)布了 GPT-2，因?yàn)樵谖谋旧缮系膬?yōu)異表現(xiàn)，以及對(duì)于預(yù)訓(xùn)練 Transformer 架構(gòu)的充分運(yùn)用，被認(rèn)為是如今大預(yù)言模型的「始祖」。
五年后的今天，訓(xùn)練 GPT-2 這樣 15 億參數(shù)的大模型，只需要花費(fèi) 672 美元，在一個(gè) 8XH100 的 GPU 節(jié)點(diǎn)上跑 24 個(gè)小時(shí)就可以搞定了。
本周四，前特斯拉 Autopilot 負(fù)責(zé)人、OpenAI 科學(xué)家 Andrej Karpathy 在他純 C 語(yǔ)言復(fù)現(xiàn) GPT-2 大模型的項(xiàng)目「llm.c」的最新進(jìn)展中分享了他的訓(xùn)練心得：
令人難以置信的是，由于計(jì)算硬件（英偉達(dá) H100 GPU）、軟件（CUDA、cuBLAS、cuDNN、FlashAttention 等）和數(shù)據(jù)質(zhì)量（例如 FineWeb-Edu 數(shù)據(jù)集）的改進(jìn)，過(guò)去 5 年間，大語(yǔ)言模型的訓(xùn)練成本大幅下降。Karpathy 表示，對(duì)于此次實(shí)踐，算法遵循 GPT-2/3 論文基本保持原樣不變。
當(dāng)年 OpenAI 訓(xùn)練 GPT-2 花費(fèi)了多少錢？這是個(gè)至今仍然未知的數(shù)字。Karpathy 粗略地估算認(rèn)為是這回成本的 100 倍，大概要到 10 萬(wàn)美元的量級(jí)。
基本相同的任務(wù)，運(yùn)行效率卻有天壤之別，這體現(xiàn)了近幾年來(lái) AI 領(lǐng)域和算力基礎(chǔ)設(shè)施的飛速發(fā)展。
由于 llm.c 是在 C/CUDA 中 GPT 訓(xùn)練的直接實(shí)現(xiàn)，因此要求其實(shí)很少 —— 不需要 conda 環(huán)境、Python 解釋器、pip 安裝等。如果你也要嘗試，可以啟動(dòng)云 GPU 節(jié)點(diǎn)（例如在 Lambda 上），可選擇安裝 NVIDIA cuDNN、NCCL/MPI，下載 .bin 數(shù)據(jù)分片，編譯并運(yùn)行，幾分鐘后即可開始。
然后，你就可以等待 24 小時(shí)，然后欣賞通用大語(yǔ)言模型的能力了。
「對(duì)于 llm.c 項(xiàng)目來(lái)說(shuō)，這是一個(gè)非常好的節(jié)點(diǎn)。因?yàn)檎麄€(gè)項(xiàng)目都是從我考慮為教育視頻重現(xiàn) GPT-2 開始的。我遇到一些 PyTorch 的東西時(shí)卡住了，然后憤怒地退出，再用 C/CUDA 從頭開始編寫整個(gè)項(xiàng)目，」Karpathy 表示。「這讓我踏上了比預(yù)想更長(zhǎng)的旅程。但它非常有趣，我學(xué)到了更多的 CUDA，一路上結(jié)交了朋友，現(xiàn)在的 llm.c 真的很棒。它有大約 5000 行代碼，編譯和步驟非常快，幾乎不需要等待。它具有恒定的內(nèi)存占用，它以混合精度進(jìn)行訓(xùn)練，使用 NNCL 分布在多節(jié)點(diǎn)上。它是按位確定性的，并且徘徊在 MFU 的 50% 左右。所以它很 ok。」
對(duì)于 llm.c 項(xiàng)目而言，越做似乎挖得坑越大。Andrej Karpathy 對(duì)目前的運(yùn)行結(jié)果仍然不是 100% 滿意 —— 他認(rèn)為評(píng)估應(yīng)該更好，訓(xùn)練應(yīng)該更穩(wěn)定，尤其是在較長(zhǎng)時(shí)間運(yùn)行的較大模型尺寸下。
他還預(yù)告了一些有趣的新方向：fp8（即將推出）、推理、微調(diào)、多模態(tài)（VQVAE 等）、更現(xiàn)代的架構(gòu)（Llama/Gemma）。llm.c 的目標(biāo)仍然是為功能齊全的 LLM 智能體提供簡(jiǎn)單、最小、干凈的訓(xùn)練堆棧，直接使用 C/CUDA，并包含配套的教育材料，可以讓許多初學(xué)者快速了解這個(gè)令人敬畏的領(lǐng)域。
說(shuō)完了這么多，該看看 24 小時(shí)訓(xùn)練 GPT-2 的成果了：Karpathy 使用更長(zhǎng)的 400B token GPT-2 運(yùn)行（從 33B token 增加），效果良好，直到 330B（達(dá)到 61% HellaSwag，遠(yuǎn)高于這個(gè)大小的 GPT-2 和 GPT-3），然后在這個(gè)圖之后不久爆炸了。目前作者還在繼續(xù)進(jìn)行研究。
接下來(lái)看詳細(xì)項(xiàng)目介紹。
GitHub 地址：https://github.com/karpathy/llm.c/discussions/677
訓(xùn)練。使用 llm.c 訓(xùn)練 GPT-2 非常簡(jiǎn)單，因?yàn)樗怯?C/CUDA 編寫的，因此不需要 minconda、Python、PyTorch 等。你只需一個(gè) 8XH100 GPU box，Karpathy 建議從 Lambda Labs 購(gòu)買一個(gè)。
不過(guò) llm.c 在計(jì)算上很靈活，如果你只有 1 個(gè) GPU，仍然可以訓(xùn)得 GPT-2，這時(shí)你需要等待 8 天而不是 1 天。如果你有 16 個(gè) GPU（例如使用新的 Lambda 1 Click Clusters），則能夠訓(xùn)練多節(jié)點(diǎn)，這時(shí)只需等待 12 小時(shí)。啟動(dòng)節(jié)點(diǎn)后，以下是訓(xùn)練 GPT-2 的完整說(shuō)明：
# install cudnn so we can use FlashAttention and run fast (optional)
# https://developer.nvidia.com/cudnn-downloads
# for me, CUDA 12 (run `nvcc --version`) running on Linux x86_64 Ubuntu 22.04
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
# "install" cudnn-frontend to ~/
git clone https://github.com/NVIDIA/cudnn-frontend.git
# install MPI (optional, if you intend to use multiple GPUs)
# (you might also have to install NVIDIA NCCL if it doesn't come with your setup)
sudo apt -y install openmpi-bin openmpi-doc libopenmpi-dev
# download and enter llm.c repo
git clone https://github.com/karpathy/llm.c.gitcd llm.c
# download the "starter pack" (~1GB download)
# contains GPT2-124M weights (used in tests), tokenizer, eval data .bin s
./dev/download_starter_pack.sh
# download the training dataset (FineWeb-Edu 100B token) .bin data shards
# note: this is a total of 1001 data shards. If you only want to test things
# out and don't want to do an actual run, feel free to append the number of
# training shards to download (e.g. for just 10 shards: ./edu_fineweb.sh 10)
# the full dataset is ~200GB, we can store it here in dev/data directory.
cd dev/data
./edu_fineweb.sh
# compile (~1 min 1st time for cuDNN mostly, few sec from then on)
cd ../../
make train_gpt2cu USE_CUDNN=1
# and train! (wait 24 hours here)
mpirun -np 8 ./train_gpt2cu \
-i "dev/data/edu_fineweb100B/edu_fineweb_train_*.bin" \
-j "dev/data/edu_fineweb100B/edu_fineweb_val_*.bin" \
-o "log_gpt2_1558M" \
-v 250 -s 300000 -g 384 \
-h 1 \
-b 16 -t 1024 \
-d 1048576 \
-r 0 \
-z 1 \
-c 0.1 \
-k "cosine" \
-l 0.0006 \
-q 0.1 \
-u 700 \
-n 2000 \
-x 32000 \
-ge 1 \
-y 1 \
-e "d48"
開始優(yōu)化：
num_parameters: 1557686400 => bytes: 3115372800
allocated 2971 MiB for model parameters
batch_size B=16 * seq_len T=1024 * num_processes=8 and total_batch_size=1048576
=> setting grad_accum_steps=8
created directory: log_gpt2_1558M
allocating 40409 MiB for activations
val loss 11.129390
allocating 2971 MiB for parameter gradients
allocating 742 MiB for AdamW optimizer state m
allocating 742 MiB for AdamW optimizer state v
allocating 742 MiB for master copy of params
step 1/32000 | loss 11.133732 (+nanz)| norm 52.9732 (+nanz)| lr 8.57e-07 | 3056.36 ms | 42.6% bf16 MFU | 343080 tok/s
step 2/32000 | loss 10.539388 (+nanz)| norm 43.5996 (+nanz)| lr 1.71e-06 | 2747.19 ms | 47.4% bf16 MFU | 381690 tok/s
step 3/32000 | loss 9.894109 (+nanz)| norm 23.2229 (+nanz)| lr 2.57e-06 | 2753.25 ms | 47.3% bf16 MFU | 381259 tok/s
step 4/32000 | loss 9.566241 (+nanz)| norm 28.4920 (+nanz)| lr 3.43e-06 | 2741.47 ms | 47.5% bf16 MFU | 381690 tok/s
step 5/32000 | loss 9.482848 (+nanz)| norm 23.7817 (+nanz)| lr 4.29e-06 | 2752.07 ms | 47.3% bf16 MFU | 381507 tok/s
step 6/32000 | loss 9.332832 (+nanz)| norm 15.9113 (+nanz)| lr 5.14e-06 | 2751.01 ms | 47.3% bf16 MFU | 381431 tok/s
step 7/32000 | loss 9.165650 (+nanz)| norm 10.5941 (+nanz)| lr 6.00e-06 | 2753.03 ms | 47.3% bf16 MFU | 381327 tok/s
step 8/32000 | loss 9.132234 (+nanz)| norm 16.2733 (+nanz)| lr 6.86e-06 | 2748.91 ms | 47.3% bf16 MFU | 381348 tok/s
step 9/32000 | loss 9.097384 (+nanz)| norm 12.1342 (+nanz)| lr 7.71e-06 | 2748.73 ms | 47.3% bf16 MFU | 381367 tok/s
step 10/32000 | loss 9.072879 (+nanz)| norm 10.5923 (+nanz)| lr 8.57e-06 | 2749.40 ms | 47.3% bf16 MFU | 381369 tok/s
...
每一步大約需要 2.75 秒，共有 32000 步，所以現(xiàn)在我們等待 24 小時(shí)左右。在每一步中，訓(xùn)練運(yùn)行都會(huì)占用約 100 萬(wàn)個(gè) FineWeb-EDU token（這些來(lái)自互聯(lián)網(wǎng)的教育網(wǎng)頁(yè)），并更新模型的 15.58 億個(gè)權(quán)重，使其能夠更好地預(yù)測(cè)序列中的下一個(gè) token。最后將總共處理 32000 * 1048576 = 33.6B 個(gè) token。隨著更好地預(yù)測(cè)下一個(gè) token，損失會(huì)下降。范數(shù)將穩(wěn)定在 0.1-1 左右，學(xué)習(xí)率在前面幾步預(yù)熱。模型 flops 利用率 (MFU) 約為 50%，非常高效。
等待 24 小時(shí)后，就可以使用 dev/vislog.ipynb jupyter 筆記本可視化 main.log 日志文件。為此，你還需要安裝 Python 和 matplotlib。
參數(shù)指南。OpenAI 發(fā)布的 GPT-2 包含模型權(quán)重，但細(xì)節(jié)很少；而 GPT-3 版本沒(méi)有權(quán)重，但細(xì)節(jié)很多。因此，在許多情況下，我們遵循 GPT-3 論文超參數(shù)，因?yàn)?GPT-2 論文的信息非常少。具體參見(jiàn)原項(xiàng)目。
內(nèi)存指南。大多數(shù)人可能面臨的最大限制是他們的 GPU 沒(méi)有 80GB。沒(méi)關(guān)系，你仍然可以運(yùn)行上面的所有內(nèi)容，只是運(yùn)行速度會(huì)更慢。因此如果模型不適配，你會(huì)怎么做？最重要的是微批大小 - b。嘗試減小它，但將其保持在合適的數(shù)字，例如 16 → 8 → 4 → 2 → 1。從那里開始，嘗試使用重計(jì)算設(shè)置 -r，即 0（最快且有大量?jī)?nèi)存）、1（稍微慢一點(diǎn)，但節(jié)省大量?jī)?nèi)存）或 2（稍微慢一點(diǎn)，節(jié)省較少內(nèi)存）。
你可以做的下一件事是禁用 fp32 中的主權(quán)重，可以使用 - w 0 （默認(rèn)值 1）來(lái)執(zhí)行此操作。我們不會(huì)維護(hù) fp32 參數(shù)副本。根據(jù)經(jīng)驗(yàn)，在之前的幾次運(yùn)行中，這似乎沒(méi)問(wèn)題，可能是因?yàn)槭褂昧穗S機(jī)舍入。如果還不適合，則可以嘗試使用 -t 來(lái)減少最大序列長(zhǎng)度，默認(rèn)值為 1024，你可以將其降低到 512、256 等。但現(xiàn)在你會(huì)讓模型變得更糟，因?yàn)樗淖畲笞⒁饬缍日跍p少。
代碼。Karpathy 對(duì) llm.c 略有偏愛(ài)，認(rèn)為它非常漂亮：
它只需要基本的 CUDA 依賴項(xiàng)即可運(yùn)行。
它是 C/CUDA 中直接、最小且易讀的實(shí)現(xiàn)。llm.c 共有約 5,000 行 C/CUDA 代碼。這里嘗試主要使用 C，而不是 C++，以保持簡(jiǎn)單。神經(jīng)網(wǎng)絡(luò)訓(xùn)練只是對(duì)單個(gè)浮點(diǎn)數(shù)組進(jìn)行相同的簡(jiǎn)單算術(shù)運(yùn)算（如 +、-、、/）的一個(gè) while 循環(huán)，它實(shí)際上不應(yīng)該那么復(fù)雜。
它編譯和運(yùn)行非常快（幾秒鐘），因此可以進(jìn)行更多步進(jìn)和更短等待。
它在開始時(shí)一次性分配其所有 GPU 內(nèi)存，從那時(shí)起在訓(xùn)練期間具有完全恒定的內(nèi)存占用。因此，一旦開始步進(jìn)，就可以在剩余的運(yùn)行中表現(xiàn)良好并且不會(huì)內(nèi)存用完（OOM）。
它是按位（bitwise）確定的。
它非常高效，略低于～50% 的 MFU。
主要入口點(diǎn)和大部分代碼位于文件 train_gpt2.cu 中。它包含 GPT-2 模型定義和約 2,000 LOC 的訓(xùn)練 loop，并從 llmc 目錄導(dǎo)入了一堆帶有各種實(shí)用程序和各個(gè)層實(shí)現(xiàn)的輔助文件。最后 cloc llmc 報(bào)告了 23 個(gè)文件、3170 LOC，而 cloc train_gpt2.cu 目前是 1353 LOC。
多節(jié)點(diǎn)訓(xùn)練。如果你擁有大量 GPU，并且 llm.c 支持多節(jié)點(diǎn)訓(xùn)練，則不用考慮太多了。Karpathy 見(jiàn)過(guò)訓(xùn)練 llm.c 時(shí)最多使用了約 500 個(gè) GPU，他自己迄今為止進(jìn)行的最大規(guī)模運(yùn)行是在 Lambda 的全新一鍵集群功能上進(jìn)行的，在 2 個(gè)節(jié)點(diǎn)中共使用了 16XH100 GPU。
同時(shí) lambda 團(tuán)隊(duì)提供了有關(guān)如何在其一鍵集群上訓(xùn)練 llm.c 模型的詳細(xì)說(shuō)明。例如使用 512-GPU H100 集群，每小時(shí)花費(fèi) 2,300 美元，你或許能夠在約 30 分鐘內(nèi)訓(xùn)練 GPT-2。你必須增加總批量大小（例如增加至約 8M），或許還得微調(diào)超參數(shù)。Karpathy 還沒(méi)有嘗試過(guò)，但它可能有效，而且會(huì)非常酷。
與 PyTorch 比較。Karpathy 認(rèn)為在 PyTorch 中相當(dāng)?shù)倪\(yùn)行看起來(lái)像這樣，使用并行 PyTorch 實(shí)現(xiàn)：
torchrun --standalone --nproc_per_node=8 train_gpt2.py \
--input_bin "dev/data/edu_fineweb100B/edu_fineweb_train_*.bin" \
--input_val_bin "dev/data/edu_fineweb100B/edu_fineweb_val_*.bin" \
--write_tensors 0 \
--model d48 \
--batch_size 8 --sequence_length 1024 --total_batch_size 1048576 \
--dtype bfloat16 \
--compile 1 \
--tensorcores 1 \
--flash 1 \
--num_iterations 32000 \
--warmup_iters 700 \
--weight_decay 0.1 \
--overfit_single_batch 0 \
--learning_rate 0.0006 \
--zero_stage 1PyTorch 代碼僅供測(cè)試參考，而非實(shí)際實(shí)現(xiàn)，因此訓(xùn)練 loop 在某些地方會(huì)略有不同（例如數(shù)據(jù)加載器不會(huì)對(duì)分片進(jìn)行置換等），但這仍可能作為參考點(diǎn)有用。這里還將默認(rèn)詞匯大小修改為 50257 → 50304 以提高效率，然后當(dāng)前的 PyTorch 夜間給出：
step 16/32000 | train loss 8.903997 | norm 8.3474 | lr 1.37e-05 | (3381.88 ms | 310057 tok/s)
step 17/32000 | train loss 8.870140 | norm 3.7936 | lr 1.46e-05 | (3381.95 ms | 310051 tok/s)
step 18/32000 | train loss 8.875732 | norm 9.4993 | lr 1.54e-05 | (3393.09 ms | 309033 tok/s)
step 19/32000 | train loss 8.817432 | norm 2.8345 | lr 1.63e-05 | (3379.75 ms | 310253 tok/s)
step 20/32000 | train loss 8.798056 | norm 4.1234 | lr 1.71e-05 | (3386.53 ms | 309631 tok/s)
step 21/32000 | train loss 8.777574 | norm 2.8010 | lr 1.80e-05 | (3386.05 ms | 309675 tok/s)
...
現(xiàn)在不能說(shuō)完全有信心 PyTorch 腳本已得到最大程度的調(diào)整，但可以得到以下觀察結(jié)果。
PyTorch 似乎占用了更多內(nèi)存（此次運(yùn)行約為 80GB），而 llm.c 占用了 57GB（減少了 29%）。內(nèi)存很重要，因?yàn)樗试S增加批處理大小（例如 llm.c 在此處最多可以增加到 24 個(gè)微批處理），這樣速度會(huì)更快一些。
其次，每次迭代大約為 3386 毫秒，而非 2750 毫秒，因此 llm.c 的速度提高了約 19%。這里的一些收益是已知的，例如 llm.c 包括啟動(dòng)反向傳遞的融合分類器等優(yōu)化，這是 torch.compile 目前無(wú)法做到的。
但是也可能存在一種情況，這個(gè)腳本沒(méi)有完全進(jìn)行最大程度的調(diào)整。這里不做贅述。
最終模型。以下幾個(gè)鏈接可能對(duì)其他人有幫助：
main.log 文件（http://llmc.s3-us-west-2.amazonaws.com/gpt2_1558M/main.log）
model_00032000.bin llm.c bin 模型文件（http://llmc.s3-us-west-2.amazonaws.com/gpt2_1558M/model_00032000.bin）
轉(zhuǎn)換為 huggingface transformers GPT-2 模型（https://huggingface.co/karpathy/gpt2_1558M_final2_hf）
模型導(dǎo)出。模型導(dǎo)出可以按如下方式進(jìn)行：
python dev/eval/export_hf.py --input log_gpt2_128M/model_00032000.bin --output gpt2_1558M_export
然后就可以運(yùn)行 Eleuther 評(píng)估工具，或者運(yùn)行 huggingface 采樣 pipeline 來(lái)獲取模型樣本：
# take model for spin
import torch
output = "./gpt2_1558M_final2_hf"
# set pytorch seeds
torch.manual_seed(42)torch.cuda.manual_seed(42)
prompt = "In a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English."
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(output)model = AutoModelForCausalLM.from_pretrained(output, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map='cuda')model.eval()tokens = tokenizer.encode(prompt, return_tensors="pt")tokens = tokens.to('cuda')
output = model.generate(tokens, max_new_tokens=500, pad_token_id=tokenizer.eos_token_id, do_sample=True, top_k=50, num_return_sequences=4)samples = tokenizer.batch_decode(output)for sample in samples:
print('-'*30)
print(sample)
你還可以查看 dev/eval，以獲取有關(guān)如何運(yùn)行 Eleuther Evaluation Harness、以及 HuggingFace Open LLM 排行榜評(píng)估等說(shuō)明。
400B token 運(yùn)行。Karpathy 還嘗試將訓(xùn)練 GPT-2 的時(shí)間遠(yuǎn)超過(guò) 33B token，特別是將 -x 更改為 400,000 以訓(xùn)練 420B token（甚至比使用 300B token 訓(xùn)練的 GPT-3 還要多）。這個(gè)模型運(yùn)行看起來(lái)很棒，直到大約 330,000 步：
最終，模型在 HellaSwag 上大大超越了同等大小的 GPT-2 和 GPT-3（最高可達(dá)約 61%），但遺憾的是，從那時(shí)起它就變得不穩(wěn)定了。在此過(guò)程中，還有更多較小的峰值，但代碼配置為檢測(cè)更簡(jiǎn)單的瞬時(shí)不穩(wěn)定性并跳過(guò)更新（Karpathy 使用了標(biāo)志 sl 5.0 -sg 5.0），這有助于緩解和推遲問(wèn)題。但是，他認(rèn)為對(duì)初始化、激活范圍和整體模型訓(xùn)練穩(wěn)定性還不夠謹(jǐn)慎，并存在更深層次問(wèn)題，這些問(wèn)題會(huì)逐漸使模型陷入不穩(wěn)定狀態(tài)，較大模型和長(zhǎng)時(shí)間訓(xùn)練更是如此。
參考內(nèi)容：
https://x.com/karpathy/status/1811467135279104217
公眾號(hào)后臺(tái)回復(fù)“數(shù)據(jù)集”獲取100+深度學(xué)習(xí)各方向資源整理
極市干貨
點(diǎn)擊閱讀原文進(jìn)入CV社區(qū)
收獲更多技術(shù)干貨