租用 H100 的钱只需 233 美元。
还记得 Andrej Karpathy 纯 C 语言复现 GPT-2 大模型的项目吗？
今年 4 月，AI 领域大牛 Karpathy 一个仅用 1000 行代码即可在 CPU/fp32 上实现 GPT-2 训练的项目「llm.c」曾经引发机器学习社区的热烈讨论。
llm.c 旨在大幅简化大模型的训练，ta 使用纯 C 语言 / CUDA，不需要 245MB 的 PyTorch 或 107MB 的 cPython。不过即使是这样的优化，复现 GPT-2 级别的模型也需要在 8 块 H100 上花费 45 分钟进行训练。
没想到几个月过去，业界水平居然有了指数级的提升，让 Karpathy 本人都感到惊叹：
在 GitHub 上出现了一个新项目「Modded-NanoGPT」，对技术进行了大幅度的迭代，现在实现相同的结果只需要 5 分钟。该研究的作者 Keller Jordan 曾在 Hive AI 工作，一直以来的研究方向都着重于模型训练的优化。他在本周三表示，利用具有大序列长度的 FlexAttention，他已把速度的记录从 7.2 分钟提升到了 5 分钟。
现在有了 FlexAttention 和较大的 seqlen，文档的拆分更少了，因此语言建模在训练和验证时都变得更容易。该记录在 HellaSwag 上的准确率略有降低，约为 29%，而之前的记录和 Andrej Karpathy 的原始训练准确率约为 30%。
让我们看看他是怎么做的：
项目链接：https://github.com/KellerJordan/modded-nanogpt/tree/master
Modded-NanoGPT
该项目名为「Modded-NanoGPT」，它是 llm.c 存储库的 PyTorch GPT-2 训练器的改进变体：
10B tokens-->1B tokens
8xH100 上花 45 分钟训练 -->8xH100 上花 5 分钟训练
Modded-NanoGPT 采用如下技术：
先进的架构：旋转嵌入、QK-Norm 和 ReLU^2；
新优化器：Muon；
嵌入中的 Untied Head；
投影和分类层初始化为零（muP-like）；
架构 shortcut：值残差和嵌入 shortcut（部分遵循论文《Value Residual Learning For Alleviating Attention Concentration In Transformers》）；
动量（Momentum）warmup；
Tanh soft logit capping（遵循 Gemma 2）；
FlexAttention。
要进行训练，请运行以下三个命令：
pip install -r requirements.txtpip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124 —upgrade # install torch 2.6.0python data/cached_fineweb10B.py 10 # downloads only the first 1.0B training tokens to save time./run.sh
在网络连接良好的 8xH100 上，训练应在 20 分钟内完成。
sudo apt-get updatesudo apt-get install vim tmux python3-pip python-is-python3 -ygit clone https://github.com/KellerJordan/modded-nanogpt.gitcd modded-nanogpttmuxpip install numpy==1.23.5 huggingface-hub tqdmpip install --upgrade torch &python data/cached_fineweb10B.py 18
sudo docker build -t modded-nanogpt .sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 18sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh
Muon optimizer
@torch.compiledef zeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):assert len (G.shape) == 2b, c = (3.4445, -4.7750, 2.0315)X = G.bfloat16 () / (G.norm () + eps)if G.size (0) > G.size (1):X = X.Tfor _ in range (steps):A = X @ X.TB = b * A + c * A @ AX = a * X + B @ Xif G.size (0) > G.size (1):X = X.Treturn X.to (G.dtype)
内存使用量比 Adam 低 采样效率提高约 1.5 倍 挂钟开销小于 2%
总结
在更新中使用 Nesterov 动量，在动量之后应用正交化。 使用特定的五次 Newton-Schulz 迭代作为正交化方法。 使用五次多项式的非收敛系数以最大化零处的斜率，从而最小化必要的 Newton-Schulz 迭代次数。事实证明，方差实际上并不那么重要，因此我们最终得到一个五次多项式，它在重复应用后（快速）收敛到 0.68、1.13 的范围，而不是到 1。 在 bfloat16 中运行 Newton-Schulz 迭代（而 Shampoo 实现通常依赖于在 fp32 或 fp64 中运行的逆 pth 根）。
扫描二维码添加小助手微信
