神級項目訓練GPT-2僅需5分鐘，Andrej Karpathy都點讚

機器之心報導

編輯：澤南、小舟

租用 H100 的錢只需 233 美元。

還記得 Andrej Karpathy 純 C 語言複現 GPT-2 大模型的項目嗎？

今年 4 月，AI 領域大牛 Karpathy 一個僅用 1000 行代碼即可在 CPU/fp32 上實現 GPT-2 訓練的項目「llm.c」曾經引發機器學習社區的熱烈討論。

llm.c 旨在大幅簡化大模型的訓練，ta 使用純 C 語言 / CUDA，不需要 245MB 的 PyTorch 或 107MB 的 cPython。不過即使是這樣的優化，複現 GPT-2 級別的模型也需要在 8 塊 H100 上花費 45 分鐘進行訓練。

沒想到幾個月過去，業界水平居然有了指數級的提升，讓 Karpathy 本人都感到驚歎：

在 GitHub 上出現了一個新項目「Modded-NanoGPT」，對技術進行了大幅度的迭代，現在實現相同的結果只需要 5 分鐘。該研究的作者 Keller Jordan 曾在 Hive AI 工作，一直以來的研究方向都著重於模型訓練的優化。他在本週三表示，利用具有大序列長度的 FlexAttention，他已把速度的記錄從 7.2 分鐘提升到了 5 分鐘。

現在有了 FlexAttention 和較大的 seqlen，文檔的拆分更少了，因此語言建模在訓練和驗證時都變得更容易。該記錄在 HellaSwag 上的準確率略有降低，約為 29%，而之前的記錄和 Andrej Karpathy 的原始訓練準確率約為 30%。

讓我們看看他是怎麼做的：

項目鏈接：https://github.com/KellerJordan/modded-nanogpt/tree/master

Modded-NanoGPT

該項目名為「Modded-NanoGPT」，它是 llm.c 存儲庫的 PyTorch GPT-2 訓練器的改進變體：

10B tokens–>1B tokens
8xH100 上花 45 分鐘訓練 –>8xH100 上花 5 分鐘訓練

Modded-NanoGPT 採用如下技術：

先進的架構：旋轉嵌入、QK-Norm 和 ReLU^2；
新優化器：Muon；
嵌入中的 Untied Head；
投影和分類層初始化為零（muP-like）；
架構 shortcut：值殘差和嵌入 shortcut（部分遵循論文《Value Residual Learning For Alleviating Attention Concentration In Transformers》）；
動量（Momentum）warmup；
Tanh soft logit capping（遵循 Gemma 2）；
FlexAttention。

要進行訓練，請運行以下三個命令：

pip install -r requirements.txt

pip install–pre torch –index-url https://download.pytorch.org/whl/nightly/cu124 —upgrade # install torch 2.6.0

python data/cached_fineweb10B.py 10# downloads only the first 1.0B training tokens to save time

./run.sh

在網絡連接良好的 8xH100 上，訓練應在 20 分鐘內完成。

結果將是一個具有 124M 活躍參數的 transformer，在 10 億 Fineweb tokens 上訓練了 1875 steps，實現了約 3.278 的驗證損失。相比之下，預設的 llm.c PyTorch 訓練器在 100 億 tokens 上訓練了 19560 steps 後，驗證損失 >3.28。

值得一提的是，要在更少的 GPU 上運行 Modded-NanoGPT，只需修改 run.sh 以獲得不同的 –nproc_per_node。如果內存不足，只需在 train_gpt2.py 中將 device_batch_size 縮小到 16 或 32。

這裏有一個適用於全新 8xH100 實例的啟動腳本：

sudo apt-get update

sudo apt-getinstall vim tmux python3-pip python-is-python3 -y

git clone https://github.com/KellerJordan/modded-nanogpt.git

cd modded-nanogpt

tmux

pip install numpy==1.23.5 huggingface-hub tqdm

pip install–upgrade torch &

python data/cached_fineweb10B.py 18

如果 CUDA 或 NCCL 版本與你當前的系統設置不兼容，Docker 可以成為一種有用的替代方案。這種方法標準化了 CUDA、NCCL、CUDNN 和 Python 的版本，減少了依賴性問題並簡化了設置。注意：系統上必須已安裝 NVIDIA 驅動程序。

sudo docker build -t modded-nanogpt .

sudo docker run -it –rm –gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 18

sudo docker run -it –rm –gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh

有一個問題在於，NanoGPT 訓練很快是很好，但它可能無法擴展，只是過擬合了 val 損失？Keller Jordan 表示，這很難反駁，因為「按規模」是一個無限類別（如果這些方法對 >100T 的模型就不奏效了怎麼辦？），因此無法完全證明。此外，作者也同意快速運行中使用的一些方法不太可能擴展。但如果讀者關心 1.5B 模型，他們可能會被這個結果說服：

直接將快速運行（10/18/24 版本）擴展到 1.5B 參數可以得到一個具有 GPT-2（1.5B）級 HellaSwag 性能的模型，它要比 Karpathy 的基線便宜 2.5 倍（233 美元對比 576 美元）：

Muon optimizer

除了在前人的肩膀上探索，新項目也使用了 Keller Jordan 自研的優化方式。比如這個 Muon 優化器，據他所說是目前已知最快的優化器，適用於包括 CIFAR-10 和 GPT-2 規模語言建模在內的各種訓練場景。

Muon 的定義如下：

其中 NewtonSchulz5 是 Newton-Schulz 之後的迭代，它近似地用 U @ V.T 替換 G，其中 U, S, V = G.svd ()。

@torch.compile

defzeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):

assertlen (G.shape) == 2

a,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.T

B = b * A + c * A @ A

X = a * X + B @ X

ifG.size (0) > G.size (1):

X = X.T

returnX.to (G.dtype)

對於這種訓練場景，Muon 具有以下有利特性：

內存使用量比 Adam 低
采樣效率提高約 1.5 倍
掛鍾開銷小於 2%

總結

作者表示，生成此優化器的許多選擇都是通過追求 CIFAR-10 快速運行而通過實驗獲得的。其中值得一提的經驗包括：

在更新中使用 Nesterov 動量，在動量之後應用正交化。
使用特定的五次 Newton-Schulz 迭代作為正交化方法。
使用五次多項式的非收斂係數以最大化零處的斜率，從而最小化必要的 Newton-Schulz 迭代次數。事實證明，方差實際上並不那麼重要，因此我們最終得到一個五次多項式，它在重覆應用後（快速）收斂到 0.68、1.13 的範圍，而不是到 1。
在 bfloat16 中運行 Newton-Schulz 迭代（而 Shampoo 實現通常依賴於在 fp32 或 fp64 中運行的逆 pth 根）。

使用 Newton-Schulz 迭代進行正交化的方法可以追溯到 Bernstein & Newhouse (2024)，他們建議將其作為計算 Shampoo 預處理器的方法，並從理論上探索了沒有預處理器累積的 Shampoo。Keller Jordan 特別感謝了論文作者之一 Jeremy Bernstein 的協助。

如果我們在這裏使用 SVD 而不是 Newton-Schulz 迭代，那麼這個優化器就會因為太慢而無法使用。Bernstein & Newhouse 還指出，沒有預處理器累積的 Shampoo 相當於譜範數中的最陡下降，因此 Shampoo 可以被認為是一種平滑譜最陡下降的方法。所提出的優化器可以被認為是平滑譜最陡下降的第二種方法，與 Shampoo 相比，它具有不同的內存和運行時權衡。

你可能喜歡