KellerJordan/modded-nanogpt

解決的問題

它解決了盡可能快速訓練語言模型的挑戰。具體而言，它旨在尋找一種在 8 個 NVIDIA H100 GPU 上訓練模型的最快演算法，以實現特定的性能目標（在 FineWeb 驗證集上達到 3.28 的交叉熵損失）。

工作原理

該專案是一個協作式的「競速」（speedrun），透過迭代優化訓練流水線來實現。它始於基於 NanoGPT 和 llm.c 的 PyTorch 訓練器，此後集成了數十種架構與系統優化，包括：

• Optimizers: 實作了 Muon 與 NorMuon 優化器。
• Architecture: 使用了 Rotary embeddings、QK-Norm、ReLU² 以及從嵌入到每個區塊的 skip connections。
• Precision: 在 head 與 MLP 前向傳播中使用 FP8，在交叉熵計算中使用 BF16。
• Attention: 集成了 Flash Attention 3、長短滑動窗口模式與 YaRN。
• Systems: 優化了梯度 all-reduce/reduce-scatter 以及計算與通訊的重疊。

適用對象

面向對極端訓練效率、LLM 優化以及挑戰 H100 GPU 硬體利用率極限感興趣的 AI 研究人員、工程師與愛好者。

亮點

• 大幅提速：在達到相同目標損失的情況下，將訓練時間從 45 分鐘（基準）縮短至 2.2 分鐘以內。
• 數據效率：將所需 token 數從 100 億減少到 4 億以下。
• 協作演進：詳細的世界紀錄歷史，記錄了為了縮短幾秒訓練時間而採用的技術演進。