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 亿以下。
• 协作演进：详细的世界纪录历史，记录了为了缩短几秒训练时间而采用的技术演进。