還記得 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」，它是 llm.c 儲存庫的 PyTorch GPT-2 訓練器的改進變體：- 8xH100 上花 45 分鐘訓練 -->8xH100 上花 5 分鐘訓練
- 先進的架構：旋轉嵌入、QK-Norm 和 ReLU^2；
- 架構 shortcut：值殘差和嵌入 shortcut（部分遵循論文《Value Residual Learning For Alleviating Attention Concentration In Transformers》）；
- Tanh soft logit capping（遵循 Gemma 2）；
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
在網路連線良好的 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 install vim tmux python3-pip python-is-python3 -ygit clone https://github.com/KellerJordan/modded-nanogpt.gittmux
pip 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 美元）：除了在前人的肩膀上探索，新專案也使用了 Keller Jordan 自研的最佳化方式。比如這個 Muon 最佳化器，據他所說是目前已知最快的最佳化器，適用於包括 CIFAR-10 和 GPT-2 規模語言建模在內的各種訓練場景。其中 NewtonSchulz5 是 Newton-Schulz 之後的迭代，它近似地用 U @ V.T 替換 G，其中 U, S, V = G.svd ()。def zeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7): assert len (G.shape) == 2 a, b, c = (3.4445, -4.7750, 2.0315) X = G.bfloat16 () / (G.norm () + eps) if G.size (0) > G.size (1): if G.size (0) > G.size (1):
作者表示，生成此最佳化器的許多選擇都是透過追求 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 相比，它具有不同的記憶體和執行時權衡。