神級項目訓(xùn)練GPT-2僅需5分鐘,AndrejKarpathy都點贊
還記得 Andrej Karpathy 純 C 語言復(fù)現(xiàn) GPT-2 大模型的項目嗎?
今年 4 月,AI 領(lǐng)域大牛 Karpathy 一個僅用 1000 行代碼即可在 CPU/fp32 上實現(xiàn) GPT-2 訓(xùn)練的項目「llm.c」曾經(jīng)引發(fā)機(jī)器學(xué)習(xí)社區(qū)的熱烈討論。
llm.c 旨在大幅簡化大模型的訓(xùn)練,ta 使用純 C 語言 / CUDA,不需要 245MB 的 PyTorch 或 107MB 的 cPython。不過即使是這樣的優(yōu)化,復(fù)現(xiàn) GPT-2 級別的模型也需要在 8 塊 H100 上花費 45 分鐘進(jìn)行訓(xùn)練。
沒想到幾個月過去,業(yè)界水平居然有了指數(shù)級的提升,讓 Karpathy 本人都感到驚嘆:
在 GitHub 上出現(xiàn)了一個新項目「Modded-NanoGPT」,對技術(shù)進(jìn)行了大幅度的迭代,現(xiàn)在實現(xiàn)相同的結(jié)果只需要 5 分鐘。該研究的作者 Keller Jordan 曾在 Hive AI 工作,一直以來的研究方向都著重于模型訓(xùn)練的優(yōu)化。他在本周三表示,利用具有大序列長度的 FlexAttention,他已把速度的記錄從 7.2 分鐘提升到了 5 分鐘。
現(xiàn)在有了 FlexAttention 和較大的 seqlen,文檔的拆分更少了,因此語言建模在訓(xùn)練和驗證時都變得更容易。該記錄在 HellaSwag 上的準(zhǔn)確率略有降低,約為 29%,而之前的記錄和 Andrej Karpathy 的原始訓(xùn)練準(zhǔn)確率約為 30%。
讓我們看看他是怎么做的:
項目鏈接:https://github.com/KellerJordan/modded-nanogpt/tree/master
Modded-NanoGPT
該項目名為「Modded-NanoGPT」,它是 llm.c 存儲庫的 PyTorch GPT-2 訓(xùn)練器的改進(jìn)變體:
- 10B tokens-->1B tokens
- 8xH100 上花 45 分鐘訓(xùn)練 -->8xH100 上花 5 分鐘訓(xùn)練
Modded-NanoGPT 采用如下技術(shù):
- 先進(jìn)的架構(gòu):旋轉(zhuǎn)嵌入、QK-Norm 和 ReLU^2;
- 新優(yōu)化器:Muon;
- 嵌入中的 Untied Head;
- 投影和分類層初始化為零(muP-like);
- 架構(gòu) shortcut:值殘差和嵌入 shortcut(部分遵循論文《Value Residual Learning For Alleviating Attention Concentration In Transformers》);
- 動量(Momentum)warmup;
- Tanh soft logit capping(遵循 Gemma 2);
- FlexAttention。
要進(jìn)行訓(xùn)練,請運行以下三個命令:
pip install -r requirements.txtpip install--pre torch --index-url https://download.pytorch.org/whl/nightly/cu124 —upgrade # install torch 2.6.0python data/cached_fineweb10B.py 10# downloads only the first 1.0B training tokens to save time./run.sh
在網(wǎng)絡(luò)連接良好的 8xH100 上,訓(xùn)練應(yīng)在 20 分鐘內(nèi)完成。
結(jié)果將是一個具有 124M 活躍參數(shù)的 transformer,在 10 億 Fineweb tokens 上訓(xùn)練了 1875 steps,實現(xiàn)了約 3.278 的驗證損失。相比之下,默認(rèn)的 llm.c PyTorch 訓(xùn)練器在 100 億 tokens 上訓(xùn)練了 19560 steps 后,驗證損失 >3.28。
值得一提的是,要在更少的 GPU 上運行 Modded-NanoGPT,只需修改 run.sh 以獲得不同的 --nproc_per_node。如果內(nèi)存不足,只需在 train_gpt2.py 中將 device_batch_size 縮小到 16 或 32。
這里有一個適用于全新 8xH100 實例的啟動腳本:
sudo apt-get updatesudo apt-getinstall vim tmux python3-pip python-is-python3 -ygit clone https://github.com/KellerJordan/modded-nanogpt.gitcd modded-nanogpttmuxpip install numpy==1.23.5 huggingface-hub tqdmpip install--upgrade torch &python data/cached_fineweb10B.py 18
如果 CUDA 或 NCCL 版本與你當(dāng)前的系統(tǒng)設(shè)置不兼容,Docker 可以成為一種有用的替代方案。這種方法標(biāo)準(zhǔn)化了 CUDA、NCCL、CUDNN 和 Python 的版本,減少了依賴性問題并簡化了設(shè)置。注意:系統(tǒng)上必須已安裝 NVIDIA 驅(qū)動程序。
sudo docker build -t modded-nanogpt .sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 18sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh
有一個問題在于,NanoGPT 訓(xùn)練很快是很好,但它可能無法擴(kuò)展,只是過擬合了 val 損失?Keller Jordan 表示,這很難反駁,因為「按規(guī)模」是一個無限類別(如果這些方法對 >100T 的模型就不奏效了怎么辦?),因此無法完全證明。此外,作者也同意快速運行中使用的一些方法不太可能擴(kuò)展。但如果讀者關(guān)心 1.5B 模型,他們可能會被這個結(jié)果說服:
直接將快速運行(10/18/24 版本)擴(kuò)展到 1.5B 參數(shù)可以得到一個具有 GPT-2(1.5B)級 HellaSwag 性能的模型,它要比 Karpathy 的基線便宜 2.5 倍(233 美元對比 576 美元):
Muon optimizer
除了在前人的肩膀上探索,新項目也使用了 Keller Jordan 自研的優(yōu)化方式。比如這個 Muon 優(yōu)化器,據(jù)他所說是目前已知最快的優(yōu)化器,適用于包括 CIFAR-10 和 GPT-2 規(guī)模語言建模在內(nèi)的各種訓(xùn)練場景。
Muon 的定義如下:
其中 NewtonSchulz5 是 Newton-Schulz 之后的迭代,它近似地用 U @ V.T 替換 G,其中 U, S, V = G.svd ()。
@torch.compiledefzeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):assertlen (G.shape) == 2a,b, c = (3.4445, -4.7750, 2.0315)X = G.bfloat16 () / (G.norm () + eps)ifG.size (0) > G.size (1):X = X.T for_ in range (steps):A = X @ X.TB = b * A + c * A @ AX = a * X + B @ XifG.size (0) > G.size (1):X = X.T returnX.to (G.dtype)
對于這種訓(xùn)練場景,Muon 具有以下有利特性:
- 內(nèi)存使用量比 Adam 低
- 采樣效率提高約 1.5 倍
- 掛鐘開銷小于 2%
總結(jié)
作者表示,生成此優(yōu)化器的許多選擇都是通過追求 CIFAR-10 快速運行而通過實驗獲得的。其中值得一提的經(jīng)驗包括:
- 在更新中使用 Nesterov 動量,在動量之后應(yīng)用正交化。
- 使用特定的五次 Newton-Schulz 迭代作為正交化方法。
- 使用五次多項式的非收斂系數(shù)以最大化零處的斜率,從而最小化必要的 Newton-Schulz 迭代次數(shù)。事實證明,方差實際上并不那么重要,因此我們最終得到一個五次多項式,它在重復(fù)應(yīng)用后(快速)收斂到 0.68、1.13 的范圍,而不是到 1。
- 在 bfloat16 中運行 Newton-Schulz 迭代(而 Shampoo 實現(xiàn)通常依賴于在 fp32 或 fp64 中運行的逆 pth 根)。
使用 Newton-Schulz 迭代進(jìn)行正交化的方法可以追溯到 Bernstein & Newhouse (2024),他們建議將其作為計算 Shampoo 預(yù)處理器的方法,并從理論上探索了沒有預(yù)處理器累積的 Shampoo。Keller Jordan 特別感謝了論文作者之一 Jeremy Bernstein 的協(xié)助。
如果我們在這里使用 SVD 而不是 Newton-Schulz 迭代,那么這個優(yōu)化器就會因為太慢而無法使用。Bernstein & Newhouse 還指出,沒有預(yù)處理器累積的 Shampoo 相當(dāng)于譜范數(shù)中的最陡下降,因此 Shampoo 可以被認(rèn)為是一種平滑譜最陡下降的方法。所提出的優(yōu)化器可以被認(rèn)為是平滑譜最陡下降的第二種方法,與 Shampoo 相比,它具有不同的內(nèi)存和運行時權(quán)衡。
- 免責(zé)聲明
- 本文所包含的觀點僅代表作者個人看法,不代表新火種的觀點。在新火種上獲取的所有信息均不應(yīng)被視為投資建議。新火種對本文可能提及或鏈接的任何項目不表示認(rèn)可。 交易和投資涉及高風(fēng)險,讀者在采取與本文內(nèi)容相關(guān)的任何行動之前,請務(wù)必進(jìn)行充分的盡職調(diào)查。最終的決策應(yīng)該基于您自己的獨立判斷。新火種不對因依賴本文觀點而產(chǎn)生的任何金錢損失負(fù)任何責(zé)任。