将 3.8B LLM 培训至 0.384 CORE，仅需 998 美元 – Hugo Vergnes

在“nanoGPT 玩具”和“你需要一个研究实验室”之间，有一个很大的、未被充分描述的区域，一个人只要花几千美元就可以训练一个有意义的模型。我希望看到语言和理解能力从我自己的随机权重中产生，并学习那些只能从头开始学习的部分。这个项目是晚上写的，在 5090 上调试，在租用的 B200 上完成。它很大程度上受到 Andrej Karpathy 的 nanochat 的启发。结果是 3.8B 参数模型在 CORE 上得分为 0.384，在 43 小时内对 65B 代币进行了训练，花费了 998 美元。接下来是什么有效，什么无效，以及我仍然不知道的。模型 参数 代币 硬件时间成本 CORE GPT-2 (OpenAI) 1.5B — — — — 0.2565 nanochat d26 ~561M 11.2B 8× H100 ~3h — ~0.258 nanochat d32 ~1B — 8× H100 ~33h ~$1000 0.310 Little-lm 3.8B (1024 ctx) 3.848B 57.3B 8× B200 35.9h $820 0.338 little-lm 3.8B (2048 ctx) 3.848B 65.3B 8× B200 43h $998 0.384 我的模型比 nanochat d32 大，并且花费了相似的挂钟时间。 B200 的每单位工作价值高于 H100。但与 nanochat 的 1,000 美元配置大致相同的价格，它明显领先于它。这是一个令人鼓舞的数据点，说明在实验室或拥有数百万计算预算的大型公司之外可以实现什么。随着边界的移动，1,000 美元会带你走得越来越远。设置 我构建了little-lm作为一个配置驱动的框架，用于训练小型仅解码器的LLM。每次运行都由 YAML 文件完全指定：模型、数据集、优化器、计划、回调。组件自行注册到全局注册表中并通过名称进行解析，因此交换优化器或数据集只需一行配置更改。良好的基础设施几乎可以立即收回成本。普通的软件工程学科（例如关注点分离、干净的界面、可以交换的组件）在人工智能工作中非常重要。一开始我花了一点钱，后来又花了几次来修复糟糕的合同或次优的问题。但这一次，当您遇到第一个收敛问题时，投资就会收回成本。我发现一个伟大的基础设施是几乎从不需要你手动编辑代码的基础设施。如果您可以阅读配置并准确理解发生了什么，并且没有隐藏的机制，则意味着您做得很好。以下报告是能够将实验表达为三行 YAML diff 而不是分支的结果。最终模型是 Llama 风格的：RMSNorm、RoPE、GQA（24 个查询头、8 个 KV 头）、relu² MLP、QK-norm、logit softcap、每层可学习残差标量和 ResFormer 风格的值嵌入。组件参数 令牌嵌入 154.5M LM 头（未绑定） 154.5M 28 个解码器层 2,818.7M 值嵌入（14 个表） 721.2M 总计 3.848B 值得注意的是，值嵌入占参数计数的 19%。 14 个 vocab × kv_dim 表，每隔一层一个。结果 早期实验 在良好的运行之前，有很多糟糕的运行。我在 FineWeb-Edu 上训练了 858M Llama，花费了 16.4B 代币，在单个 A100 上训练了 5.8 天。 AdamW 为 2.5e-4，余弦衰减为零，5% 预热，通过梯度累积批量 256，2048 个上下文。结果：PIQA 60.45%。 GPT-2 124M 得分约为 63%。从 2019 年开始，我花了六天的时间来构建比模型小七倍的模型还要糟糕的东西。一代又一代是重复的，而且近乎荒谬。损失曲线说明了这一点。余弦衰减为零。大约 70% 的台阶后，曲线完全平坦。最后 30% 的计算预算基本上没有产生任何结果，因为学习率可能太低。线性冷却时间在很晚之后才保持有用的速率。巅峰LR太保守了。 2.5e-4 对于 858M 参数来说较低。对于那些小模型，你可以非常激进。 AdamW 无所不包。对于这个规模的矩阵参数来说，μon 应该在每个标记上表现得更好。事实上，这在消融运行中很快得到了证明。数据。 FineWeb-Edu 还不错。它不是最好的。事后剖析得出了五处变化。它们共同构成了上述运行与大幅击败 GPT-2 的模型之间的区别。梯形 LR 时间表。预热5%。保持平坦，并在峰值的最后 50% 至 5% 内以线性冷却结束。关键是模型会一直学习直到最后，而不是滑行到尾部。在 3.8B 运行中，最后一步的 eval 损失仍在下降，这正是 858M 运行未能产生的行为。 Muon 代表矩阵参数，AdamW 代表其他一切。 Muon 每步速度较慢（牛顿-舒尔茨正交化不是免费的，在浅累积基准中大约为 25%），但每个优化器步骤需要支付一次成本：在 7 个梯度累积步骤中，它会稀释到约 4%。根据总运行时间来衡量，总体收敛速度要快得多。 ClimbMix 而不是 FineWeb-Edu。这是收敛速度的巨大飞跃。正如卡帕蒂所发现的那样。 FP8 + 词汇填充。通过 torch._scaled_mm 使用动态张量尺度进行 FP8 训练