Modded-NanoGPT:90 秒训练 1.24 亿参数模型,社区协作的 LLM 训练速度挑战

Modded-NanoGPT:90 秒训练 1.24 亿参数模型,社区协作的 LLM 训练速度挑战

Andrej Karpathy 的 llm.c 项目用 8 块 H100 GPU 训练一个 1.24 亿参数的 GPT-2 模型,达到了 FineWeb 验证集 3.28 的交叉熵损失,耗时 45 分钟,消耗 100 亿 token。

Modded-NanoGPT(github.com/KellerJordan/modded-nanogpt)把这个过程压缩到了 90 秒,仅用不到 4 亿 token——训练速度提升了 30 倍,token 效率提升了 25 倍。这不是某家大公司的成果,而是由 Keller Jordan 发起、全球数十名研究者协作完成的开源 speedrun 项目。

它做了什么

这是一个协作式的技术挑战:在 8 块 NVIDIA H100 GPU 上,用最快速度训练一个在 FineWeb 验证集上达到 3.28 cross-entropy loss 的 1.24 亿参数语言模型。

从 llm.c 的 45 分钟到现在的 90 秒,这一跃升来自数十项训练算法优化的叠加:

架构改进

  • 旋转位置编码(Rotary Embeddings)、QK-Norm、ReLU² 激活
  • Flash Attention 3 + 长短滑动窗口注意力模式(受 Gemma 2 启发)
  • 从 embedding 到每个 block 的跳连,以及 block 3 到 6 的跨层跳连
  • 额外 embedding 混入 attention 的 value 层

优化器与训练技巧

  • Muon 优化器(Keller Jordan 团队自研)
  • FP8 矩阵乘法用于 head 层,非对称 rescale 和 softcap logits
  • 投影层初始化为零(类似 muP)
  • 谨慎权重衰减,与学习率调度绑定
  • 残差流的指数衰减
  • Backout 技术,最后 3 个 attention 层使用单 token 输入

数据与调度策略

  • 训练 batch 起始位置与 EoS 对齐
  • 批量大小和序列长度的动态调度
  • 累积 2 步梯度后再更新 embedding 和 lm_head 参数
  • 训练 2/3 处解除 embedding 和 lm_head 的权重绑定
  • 部分 key 偏移(Partial Key Offset)
  • 多 token 预测

每一项单独的贡献可能不大,但叠加在一起产生了数量级的效率提升。

为什么值得关注

这不是一个生产可用的训练框架,而是一个训练算法的试验场。它证明了:

  1. 小模型可以高效训练:不需要堆算力,算法优化的回报巨大
  2. 开放协作的力量:项目贡献者列表中可以看到来自全球的研究者,每个人贡献几行代码或一个想法
  3. 可复现的基准:目标明确(3.28 loss on FineWeb),任何人都可以验证结果

对于想了解 LLM 训练底层细节的研究者和工程师,这个项目是一个极佳的学习资源——每一项技术都可以在代码中找到对应实现。

快速上手

git clone https://github.com/KellerJordan/modded-nanogpt.git && cd modded-nanogpt
pip install -r requirements.txt
# 下载前 900M 训练 token(节省时间)
python data/cached_fineweb10B.py 9
./run.sh

官方验证在 8 块 H100 GPU 上进行(由 PrimeIntellect 赞助)。首次运行 torch.compile 会有约 7 分钟的编译延迟。如果没有 GPU 集群,也可以在 Colab 或单 GPU 上运行较小规模的实验。

适用边界

  • 学习价值 > 生产价值:这是研究性质的 speedrun,不是面向生产环境的训练框架
  • 硬件门槛:要达到 90 秒记录需要 8 块 H100,个人开发者可以用更少 GPU 学习技术思路
  • 特定目标:优化目标是固定 loss 阈值下的最短训练时间,不是通用的”最好模型”

来源