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

機器之心報道

編輯：澤南、小舟

還記得 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 訓練器的改進變體：

Modded-NanoGPT 採用如下技術：

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

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

在網絡連接良好的 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 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 版本與你當前的系統設置不兼容，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 18sudo 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.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)

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

總結

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

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

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