日前，AI 大牛 Andrej Karpathy 在社交媒体上力荐了一个超棒开源项目—— Modded-NanoGPT [1] ，该项目是由 KellerJordan 精心打造，对 Karpathy 自己的教学项目 llm.c [2] 中 PyTorch GPT-2 训练器的颠覆性升级，取得的成绩令人震撼：仅仅使用 8 块 H100 GPU，惊人的训练效率使 124M 参数模型的训练时间缩减至 5 分钟，不仅完美复现了原始模型的验证损失，更将机器学习训练效率推向了一个令人难以置信的新高度。并且，按照作者的迭代速度，这一纪录可能还会被打破。
Modded-NanoGPT 的成功源于其巧妙整合了多项尖端技术和优化策略。项目核心采用了现代化的架构，包括旋转嵌入（（Rotary Embeddings））、 QK-Norm 和 ReLU^2 激活函数，这些改进显著提升了模型性能和训练稳定性。
一个关键创新是引入了 Muon 优化器。这种基于牛顿-舒尔茨迭代正交化动量的新型优化算法比传统方法具有显著优势。Muon 不仅降低了内存使用，还提供了约 1.5 倍更好的样本效率，同时计算开销极小。
项目还针对模型架构进行了一些调整来进一步提升性能，包括解绑嵌入和语言模型头部、将投影和分类层初始化为零，以及实现值残差和嵌入快捷方式。这些修改使模型能在训练过程中更有效地学习。
Modded-NanoGPT 还利用了动量预热和 Tanh 软 logit 限制等先进的训练技术。这些方法有助于稳定训练过程并提高模型的泛化能力。此外，FlexAttention 灵活注意力机制的实施在达成最新速度记录中起到了关键作用。
Modded-NanoGPT 的优秀表现不仅限于 124M 参数模型。将这种方法扩展到 1.5B 参数时，也显示出令人鼓舞的结果，在 HellaSwag 基准测试中以更低的成本匹配了 GPT-2(1.5B)的性能。这种可扩展性展示了这些技术对于更大型语言模型的潜力。
该项目提供了 docker 运行的方式，有条件的读者可以尝试：
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
项目支持在不影响训练的情况下，调整有关参数，使得能够在较少的 GPU 和较少内存的情况下运行，具体可以参见项目配置。
参考资料
[1] Modded-NanoGPT: https://github.com/KellerJordan/modded-nanogpt
[2] Modded-NanoGPT: https://github.com/KellerJordan/modded-nanogpt
[3] llm.c: https://github.com/karpathy/llm.c
后台回复“进群”入群讨论。