五年後的今天,訓練GPT-2只需不到700刀、24小時,Karpathy又整新活

機器之心報道
編輯:杜偉、澤南
論老黃賣鏟子的技術含量。
2019 年 2 月,OpenAI 釋出了 GPT-2,因為在文字生成上的優異表現,以及對於預訓練 Transformer 架構的充分運用,被認為是如今大預言模型的「始祖」。
五年後的今天,訓練 GPT-2 這樣 15 億引數的大模型,只需要花費 672 美元,在一個 8XH100 的 GPU 節點上跑 24 個小時就可以搞定了。
本週四,前特斯拉 Autopilot 負責人、OpenAI 科學家 Andrej Karpathy 在他純 C 語言復現 GPT-2 大模型的專案「llm.c」的最新進展中分享了他的訓練心得:
令人難以置信的是,由於計算硬體(英偉達 H100 GPU)、軟體(CUDA、cuBLAS、cuDNN、FlashAttention 等)和資料質量(例如 FineWeb-Edu 資料集)的改進,過去 5 年間,大語言模型的訓練成本大幅下降。Karpathy 表示,對於此次實踐,演算法遵循 GPT-2/3 論文基本保持原樣不變。
當年 OpenAI 訓練 GPT-2 花費了多少錢?這是個至今仍然未知的數字。Karpathy 粗略地估算認為是這回成本的 100 倍,大概要到 10 萬美元的量級。
基本相同的任務,執行效率卻有天壤之別,這體現了近幾年來 AI 領域和算力基礎設施的飛速發展。
由於 llm.c 是在 C/CUDA 中 GPT 訓練的直接實現,因此要求其實很少 —— 不需要 conda 環境、Python 直譯器、pip 安裝等。如果你也要嘗試,可以啟動雲 GPU 節點(例如在 Lambda 上),可選擇安裝 NVIDIA cuDNN、NCCL/MPI,下載 .bin 資料分片,編譯並執行,幾分鐘後即可開始。
然後,你就可以等待 24 小時,然後欣賞通用大語言模型的能力了。
「對於 llm.c 專案來說,這是一個非常好的節點。因為整個專案都是從我考慮為教育影片重現 GPT-2 開始的。我遇到一些 PyTorch 的東西時卡住了,然後憤怒地退出,再用 C/CUDA 從頭開始編寫整個專案,」Karpathy 表示。「這讓我踏上了比預想更長的旅程。但它非常有趣,我學到了更多的 CUDA,一路上結交了朋友,現在的 llm.c 真的很棒。它有大約 5000 行程式碼,編譯和步驟非常快,幾乎不需要等待。它具有恆定的記憶體佔用,它以混合精度進行訓練,使用 NNCL 分佈在多節點上。它是按位確定性的,並且徘徊在 MFU 的 50% 左右。所以它很 ok。」
對於 llm.c 專案而言,越做似乎挖得坑越大。Andrej Karpathy 對目前的執行結果仍然不是 100% 滿意 —— 他認為評估應該更好,訓練應該更穩定,尤其是在較長時間執行的較大模型尺寸下。
他還預告了一些有趣的新方向:fp8(即將推出)、推理、微調、多模態(VQVAE 等)、更現代的架構(Llama/Gemma)。llm.c 的目標仍然是為功能齊全的 LLM 智慧體提供簡單、最小、乾淨的訓練堆疊,直接使用 C/CUDA,幷包含配套的教育材料,可以讓許多初學者快速瞭解這個令人敬畏的領域。
說完了這麼多,該看看 24 小時訓練 GPT-2 的成果了:Karpathy 使用更長的 400B token GPT-2 執行(從 33B token 增加),效果良好,直到 330B(達到 61% HellaSwag,遠高於這個大小的 GPT-2 和 GPT-3),然後在這個圖之後不久爆炸了。目前作者還在繼續進行研究。
接下來看詳細專案介紹。
GitHub 地址:https://github.com/karpathy/llm.c/discussions/677
訓練。使用 llm.c 訓練 GPT-2 非常簡單,因為它是用 C/CUDA 編寫的,因此不需要 minconda、Python、PyTorch 等。你只需一個 8XH100 GPU box,Karpathy 建議從 Lambda Labs 購買一個。
不過 llm.c 在計算上很靈活,如果你只有 1 個 GPU,仍然可以訓得 GPT-2,這時你需要等待 8 天而不是 1 天。如果你有 16 個 GPU(例如使用新的 Lambda 1 Click Clusters),則能夠訓練多節點,這時只需等待 12 小時。啟動節點後,以下是訓練 GPT-2 的完整說明:
# 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"

開始最佳化:

num_parameters: 

1557686400

 => bytes: 

3115372800

allocated 

2971

 MiB 

for

 model parameters


batch_size B=

16

 * seq_len T=

1024

 * num_processes=

8and

 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 步,所以現在我們等待 24 小時左右。在每一步中,訓練執行都會佔用約 100 萬個 FineWeb-EDU token(這些來自網際網路的教育網頁),並更新模型的 15.58 億個權重,使其能夠更好地預測序列中的下一個 token。最後將總共處理 32000 * 1048576 = 33.6B 個 token。隨著更好地預測下一個 token,損失會下降。範數將穩定在 0.1-1 左右,學習率在前面幾步預熱。模型 flops 利用率 (MFU) 約為 50%,非常高效。
等待 24 小時後,就可以使用 dev/vislog.ipynb jupyter 筆記本視覺化 main.log 日誌檔案。為此,你還需要安裝 Python 和 matplotlib。
引數指南。OpenAI 釋出的 GPT-2 包含模型權重,但細節很少;而 GPT-3 版本沒有權重,但細節很多。因此,在許多情況下,我們遵循 GPT-3 論文超引數,因為 GPT-2 論文的資訊非常少。具體參見原專案。
記憶體指南。大多數人可能面臨的最大限制是他們的 GPU 沒有 80GB。沒關係,你仍然可以執行上面的所有內容,只是執行速度會更慢。因此如果模型不適配,你會怎麼做?最重要的是微批大小 – b。嘗試減小它,但將其保持在合適的數字,例如 16 → 8 → 4 → 2 → 1。從那裡開始,嘗試使用重計算設定 -r,即 0(最快且有大量記憶體)、1(稍微慢一點,但節省大量記憶體)或 2(稍微慢一點,節省較少記憶體)。
你可以做的下一件事是停用 fp32 中的主權重,可以使用 – w 0 (預設值 1)來執行此操作。我們不會維護 fp32 引數副本。根據經驗,在之前的幾次執行中,這似乎沒問題,可能是因為使用了隨機舍入。如果還不適合,則可以嘗試使用 -t  來減少最大序列長度,預設值為 1024,你可以將其降低到 512、256 等。但現在你會讓模型變得更糟,因為它的最大注意力跨度正在減少。
程式碼。Karpathy 對 llm.c 略有偏愛,認為它非常漂亮:
  • 它只需要基本的 CUDA 依賴項即可執行。
  • 它是 C/CUDA 中直接、最小且易讀的實現。llm.c 共有約 5,000 行 C/CUDA 程式碼。這裡嘗試主要使用 C,而不是 C++,以保持簡單。神經網路訓練只是對單個浮點陣列進行相同的簡單算術運算(如 +、-、、/)的一個 while 迴圈,它實際上不應該那麼複雜。
  • 它編譯和執行非常快(幾秒鐘),因此可以進行更多步進和更短等待。
  • 它在開始時一次性分配其所有 GPU 記憶體,從那時起在訓練期間具有完全恆定的記憶體佔用。因此,一旦開始步進,就可以在剩餘的執行中表現良好並且不會記憶體用完(OOM)。
  • 它是按位(bitwise)確定的。
  • 它非常高效,略低於~50% 的 MFU。
主要入口點和大部分程式碼位於檔案 train_gpt2.cu 中。它包含 GPT-2 模型定義和約 2,000 LOC 的訓練 loop,並從 llmc 目錄匯入了一堆帶有各種實用程式和各個層實現的輔助檔案。最後 cloc llmc  報告了 23 個檔案、3170 LOC,而 cloc train_gpt2.cu 目前是 1353 LOC。
多節點訓練。如果你擁有大量 GPU,並且 llm.c 支援多節點訓練,則不用考慮太多了。Karpathy 見過訓練 llm.c 時最多使用了約 500 個 GPU,他自己迄今為止進行的最大規模執行是在 Lambda 的全新一鍵叢集功能上進行的,在 2 個節點中共使用了 16XH100 GPU。
同時 lambda 團隊提供了有關如何在其一鍵叢集上訓練 llm.c 模型的詳細說明。例如使用 512-GPU H100 叢集,每小時花費 2,300 美元,你或許能夠在約 30 分鐘內訓練 GPT-2。你必須增加總批次大小(例如增加至約 8M),或許還得微調超引數。Karpathy 還沒有嘗試過,但它可能有效,而且會非常酷。
與 PyTorch 比較。Karpathy 認為在 PyTorch 中相當的執行看起來像這樣,使用並行 PyTorch 實現:

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 

1

PyTorch 程式碼僅供測試參考,而非實際實現,因此訓練 loop 在某些地方會略有不同(例如資料載入器不會對分片進行置換等),但這仍可能作為參考點有用。這裡還將預設詞彙大小修改為 50257 → 50304 以提高效率,然後當前的 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)


...


現在不能說完全有信心 PyTorch 指令碼已得到最大程度的調整,但可以得到以下觀察結果。
PyTorch 似乎佔用了更多記憶體(此次執行約為 80GB),而 llm.c 佔用了 57GB(減少了 29%)。記憶體很重要,因為它允許增加批處理大小(例如 llm.c 在此處最多可以增加到 24 個微批處理),這樣速度會更快一些。
其次,每次迭代大約為 3386 毫秒,而非 2750 毫秒,因此 llm.c 的速度提高了約 19%。這裡的一些收益是已知的,例如 llm.c 包括啟動反向傳遞的融合分類器等最佳化,這是 torch.compile 目前無法做到的。
但是也可能存在一種情況,這個指令碼沒有完全進行最大程度的調整。這裡不做贅述。
最終模型。以下幾個連結可能對其他人有幫助:
  • 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)
  • 轉換為 huggingface transformers GPT-2 模型(https://huggingface.co/karpathy/gpt2_1558M_final2_hf)
模型匯出。模型匯出可以按如下方式進行:

python dev/eval/export_hf.py --input log_gpt2_128M/model_00032000.bin --output gpt2_1558M_export


然後就可以執行 Eleuther 評估工具,或者執行 huggingface 取樣 pipeline 來獲取模型樣本:
# 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,以獲取有關如何執行 Eleuther Evaluation Harness、以及 HuggingFace Open LLM 排行榜評估等說明。
400B token 執行。Karpathy 還嘗試將訓練 GPT-2 的時間遠超過 33B token,特別是將 -x 更改為 400,000 以訓練 420B token(甚至比使用 300B token 訓練的 GPT-3 還要多)。這個模型執行看起來很棒,直到大約 330,000 步:
最終,模型在 HellaSwag 上大大超越了同等大小的 GPT-2 和 GPT-3(最高可達約 61%),但遺憾的是,從那時起它就變得不穩定了。在此過程中,還有更多較小的峰值,但程式碼配置為檢測更簡單的瞬時不穩定性並跳過更新(Karpathy 使用了標誌 sl 5.0 -sg 5.0),這有助於緩解和推遲問題。但是,他認為對初始化、啟用範圍和整體模型訓練穩定性還不夠謹慎,並存在更深層次問題,這些問題會逐漸使模型陷入不穩定狀態,較大模型和長時間訓練更是如此。
參考內容:
https://x.com/karpathy/status/1811467135279104217
創意為王安全為先,AIGC的雙線作戰
7月17日,《AIGC體驗派》第五期,邀請到火山引擎內容安全與風控負責人張建洋和NVIDIA企業級開發者社群高階經理何琨,一起聊聊AIGC在營銷領域的創新與安全問題:
  • 如何避免大語言模型不再胡言亂語?

  • 如何確保AIGC創作內容的質量與安全?
  • 如何避免營銷活動成為黑產的提款機?
識別海報二維碼點選閱讀原文,立即報名直播。
© THE END 
轉載請聯絡本公眾號獲得授權
投稿或尋求報道:[email protected]


相關文章