神级项目训练GPT-2仅需5分钟，Andrej Karpathy都点赞

DDD

发布时间：2024-11-21 17:32:02

580人浏览过

来源于机器之心

转载

租用 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 本人都感到惊叹：

☞☞☞AI 智能聊天, 问答助手, AI 智能搜索, 免费无限量使用 DeepSeek R1 模型☜☜☜

在 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。

要进行训练，请运行以下三个命令：

<section><code>pip install -r requirements.txt</code></section><section><code>pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124124 &mdash;upgrade # install torch 2.6.0</code></section><section><code>python data/cached_fineweb10B.py 10 # downloads only the first 1.0B training tokens to save time</code></section><section><code>./run.sh</code></section>

在网络连接良好的 8xH100 上，训练应在 20 分钟内完成。

结果将是一个具有 124124M 活跃参数的 transformer，在 10 亿 Fineweb tokens 上训练了 1875 steps，实现了约 3.278 的验证损失。相比之下，默认的 llm.c PyTorch 训练器在 100 亿 tokens 上训练了 19560 steps 后，验证损失 >3.28。

值得一提的是，要在更少的 GPU 上运行 Modded-NanoGPT，只需修改 run.sh 以获得不同的 --nproc_per_node。如果内存不足，只需在 train_gpt2.py 中将 device_batch_size 缩小到 16 或 32。

这里有一个适用于全新 8xH100 实例的启动脚本：

<section><code>sudo apt-get update</code></section><section><code>sudo apt-get install vim tmux python3-pip python-is-python3 -y</code></section><section><code>git clone <a href="https://www.php.cn/link/e8cb5f581442030021d62fd780fa674d" rel="nofollow" target="_blank" >https://www.php.cn/link/e8cb5f581442030021d62fd780fa674d</a></code></section><section><code>cd modded-nanogpt</code></section><section><code>tmux</code><code></code><code></code><code>pip install numpy==1.23.5 huggingface-hub tqdm</code></section><section><code>pip install --upgrade torch &</code></section><section><code>python data/cached_fineweb10B.py 18</code></section>

如果 CUDA 或 NCCL 版本与你当前的系统设置不兼容，Docker 可以成为一种有用的替代方案。这种方法标准化了 CUDA、NCCL、CUDNN 和 Python 的版本，减少了依赖性问题并简化了设置。注意：系统上必须已安装 NVIDIA 驱动程序。

<section><code>sudo docker build -t modded-nanogpt .</code></section><section><code>sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 18</code></section><section><code>sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh</code></section>

有一个问题在于，NanoGPT 训练很快是很好，但它可能无法扩展，只是过拟合了 val 损失？Keller Jordan 表示，这很难反驳，因为「按规模」是一个无限类别（如果这些方法对 >100T 的模型就不奏效了怎么办？），因此无法完全证明。此外，作者也同意快速运行中使用的一些方法不太可能扩展。但如果读者关心 1.5B 模型，他们可能会被这个结果说服：

直接将快速运行（10/18/24 版本）扩展到 1.5B 参数可以得到一个具有 GPT-2（1.5B）级 HellaSwag 性能的模型，它要比 Karpathy 的基线便宜 2.5 倍（233 美元对比 576 美元）：

Muon optimizer

除了在前人的肩膀上探索，新项目也使用了 Keller Jordan 自研的优化方式。比如这个 Muon 优化器，据他所说是目前已知最快的优化器，适用于包括 CIFAR-10 和 GPT-2 规模语言建模在内的各种训练场景。

Muon 的定义如下：

其中 NewtonSchulz5 是 Newton-Schulz 之后的迭代，它近似地用 U @ V.T 替换 G，其中 U, S, V = G.svd ()。

<section><code>@torch.compile</code></section><section><code>def zeroth_power_via_newtonschulz5 (G, steps=5, eps=1e-7):</code></section><section><code>assert len (G.shape) == 2</code></section><section><code>a, b, c = (3.4445, -4.7750,2.0315)</code></section><section><code>X = G.bfloat16 () / (G.norm () + eps)</code></section><section><code>if G.size (0) > G.size (1):</code></section><section><code>X = X.T</code></section><section><code> </code><code>for _ in range (steps):</code></section><section><code>A = X @ X.T</code></section><section><code>B = b * A + c * A @ A</code></section><section><code>X = a * X + B @ X</code></section><section><code>if G.size (0) > G.size (1):</code></section><section><code>X = X.T</code></section><section><code> </code><code>    return X.to (G.dtype)</code></section>

对于这种训练场景，Muon 具有以下有利特性：

内存使用量比 Adam 低
采样效率提高约 1.5 倍
挂钟开销小于 2%

总结

作者表示，生成此优化器的许多选择都是通过追求 CIFAR-10 快速运行而通过实验获得的。其中值得一提的经验包括：

在更新中使用 Nesterov 动量，在动量之后应用正交化。
使用特定的五次 Newton-Schulz 迭代作为正交化方法。
使用五次多项式的非收敛系数以最大化零处的斜率，从而最小化必要的 Newton-Schulz 迭代次数。事实证明，方差实际上并不那么重要，因此我们最终得到一个五次多项式，它在重复应用后（快速）收敛到 0.68、1.13 的范围，而不是到 1。
在 bfloat16 中运行 Newton-Schulz 迭代（而 Shampoo 实现通常依赖于在 fp32 或 fp64 中运行的逆 pth 根）。

使用 Newton-Schulz 迭代进行正交化的方法可以追溯到 Bernstein & Newhouse (2024)，他们建议将其作为计算 Shampoo 预处理器的方法，并从理论上探索了没有预处理器累积的 Shampoo。Keller Jordan 特别感谢了论文作者之一 Jeremy Bernstein 的协助。

如果我们在这里使用 SVD 而不是 Newton-Schulz 迭代，那么这个优化器就会因为太慢而无法使用。Bernstein & Newhouse 还指出，没有预处理器累积的 Shampoo 相当于谱范数中的最陡下降，因此 Shampoo 可以被认为是一种平滑谱最陡下降的方法。所提出的优化器可以被认为是平滑谱最陡下降的第二种方法，与 Shampoo 相比，它具有不同的内存和运行时权衡。

AI生成活动现场的即时新闻稿

贾跃亭：法拉第未来当前股价被严重低估目标价5美元

内存飙升成本暴涨汽车会涨价还是减配？

10家中国车企入围入围财富世界500强比亚迪进前百

关税压力下创纪录！韩国汽车年出口额达720亿美元

相关专题

github中文官网入口 github中文版官网网页进入

github中文官网入口https://docs.github.com/zh/get-started，GitHub 是一种基于云的平台，可在其中存储、共享并与他人一起编写代码。通过将代码存储在GitHub 上的“存储库”中，你可以： “展示或共享”你的工作。持续“跟踪和管理”对代码的更改。

4241

2026.01.21

k8s和docker区别

k8s和docker区别有抽象层次不同、管理范围不同、功能不同、应用程序生命周期管理不同、缩放能力不同、高可用性等等区别。本专题为大家提供k8s和docker区别相关的各种文章、以及下载和课程。

280

2023.07.24

docker进入容器的方法有哪些

docker进入容器的方法：1. Docker exec；2. Docker attach；3. Docker run --interactive --tty；4. Docker ps -a；5. 使用 Docker Compose。本专题为大家提供相关的文章、下载、课程内容，供大家免费下载体验。

516

2024.04.08

docker容器无法访问外部网络怎么办

docker 容器无法访问外部网络的原因和解决方法：配置 nat 端口映射以将容器端口映射到主机端口。根据主机兼容性选择正确的网络驱动（如 host 或 overlay）。允许容器端口通过主机的防火墙。配置容器的正确 dns 服务器。选择正确的容器网络模式。排除主机网络问题，如防火墙或连接问题。本专题为大家提供相关的文章、下载、课程内容，供大家免费下载体验。

416

2024.04.08

docker镜像有什么用

docker 镜像是预构建的软件组件，用途广泛，包括：应用程序部署：简化部署，提高移植性。本专题为大家提供相关的文章、下载、课程内容，供大家免费下载体验。

454

2024.04.08

Docker容器化部署与DevOps实践

本专题面向后端与运维开发者，系统讲解 Docker 容器化技术在实际项目中的应用。内容涵盖 Docker 镜像构建、容器运行机制、Docker Compose 多服务编排，以及在 DevOps 流程中的持续集成与持续部署实践。通过真实场景演示，帮助开发者实现应用的快速部署、环境一致性与运维自动化。

2026.02.11

pytorch是干嘛的

pytorch是一个基于python的深度学习框架，提供以下主要功能：动态图计算，提供灵活性。强大的张量操作，实现高效处理。自动微分，简化梯度计算。预构建的神经网络模块，简化模型构建。各种优化器，用于性能优化。想了解更多pytorch的相关内容，可以阅读本专题下面的文章。

469

2024.05.29

Python AI机器学习PyTorch教程_Python怎么用PyTorch和TensorFlow做机器学习

PyTorch 是一种用于构建深度学习模型的功能完备框架，是一种通常用于图像识别和语言处理等应用程序的机器学习。使用Python 编写，因此对于大多数机器学习开发者而言，学习和使用起来相对简单。 PyTorch 的独特之处在于，它完全支持GPU，并且使用反向模式自动微分技术，因此可以动态修改计算图形。

2025.12.22