KellerJordan/modded-nanogpt

解决的问题

旨在使用 8 张 NVIDIA H100 GPU，找到训练语言模型达到特定性能目标（FineWeb 验证集上 3.28 的交叉熵损失）的最快可能算法和系统配置。

如何工作

该项目以“速度挑战”形式运行，通过迭代应用并测试大量架构、优化和系统级技术，以减少训练时间。它起源于 Andrej Karpathy 的 llm.c 和 NanoGPT，从 45 分钟的基线缩短至 3 分钟以下。关键方法包括：

• 架构变更：实现 Rotary 嵌入、QK-Norm、ReLU²，以及从嵌入到每个模块的跳跃连接。
• 优化器创新：使用 Muon 和 NorMuon 优化器，以及带有双轨动量的 ANVIL 优化器。
• 精度与硬件：在 MLP 前向/反向传播、QKV 投影和头部使用 FP8，同时结合 CUDA 图捕获和 Flash Attention 3。
• 数据与训练逻辑：将批量开始与句子结束（EoS）标记对齐，实现多标记预测，并使用批量大小调度。

适用人群

对极端 LLM 训练效率、硬件加速以及优化器和架构设计前沿感兴趣的科研人员和工程师。

主要亮点

• 在 8xH100 GPU 上将训练时间从 45 分钟缩短至 3 分钟以下。
• 达到目标损失所需的标记数从 100 亿减少至 3.3 亿以下。
• 实现了多种前沿技术，如 Polar Express、稀疏注意力门控和哈希大词元/三词元嵌入表。