Modded-NanoGPT

本仓库托管了 NanoGPT 速度挑战赛，我们（协作/竞争）寻找最快的算法，使用 8 块 NVIDIA H100 GPU 训练一个语言模型，使其在 FineWeb 验证集上达到 3.28 的交叉熵损失。

该目标（FineWeb 验证损失 3.28）源自 Andrej Karpathy 在 llm.c 中对 GPT-2 的复现，该复现在运行 45 分钟后达到了此损失。速度挑战赛的代码同样源自 llm.c 的 PyTorch 训练器，而后者又源自 NanoGPT，因此得名。

得益于众多贡献者的努力，本仓库现已包含一个训练算法，能够在以下条件下达到目标性能：
* 在 8xH100 上 不到 90 秒（llm.c 的 GPT-2 复现需要 45 分钟）
* 不到 4 亿 tokens（llm.c 的 GPT-2 复现需要 100 亿）

训练速度的提升得益于以下技术：
* 现代化架构：旋转位置编码、QK-Norm 和 ReLU²
* Muon 优化器 [文章] [仓库]
* 对输出头使用 FP8 矩阵乘法，以及非对称重新缩放和 logit softcap
* 将投影初始化为零（类似 muP）
* 从嵌入层到每个块的跳跃连接，以及从块 3 到 6 的跳跃连接
* 混入注意力层值（value）的额外嵌入（受 Zhou et al. 2024 启发）
* 带有长短滑动窗口注意力模式（受 Gemma 2 启发）和窗口大小预热（使用 YaRN）的 Flash Attention 3
* 对齐训练批次起始位置到文档结束符，并设置最大文档长度
* 在更新参数前，为嵌入层和语言模型头累积 2 步梯度
* 对最后 3 个注意力层使用单一激活输入的 Backout
* Muon 中的 Polar Express 实现
* 用于实现 1 token 回溯的 Smear 模块
* 稀疏注意力门控
* NorMuon
* 谨慎权重衰减，其调度与学习率绑定
* 残差流的指数衰减
* 批次大小调度
* 最大序列长度调度
* 部分键偏移
* 多 token 预测
* 在训练进行到 2/3 时解绑嵌入和语言模型头
* 对值嵌入和跳跃连接额外添加门控
* 配对头注意力
* 二元组哈希嵌入

以及许多系统层面的优化。

贡献者列表（随每项新纪录不断增长）：@bozavlado; @brendanh0gan; @fernbear.bsky.social; @Grad62304977; @jxbz; @kellerjordan0; @KoszarskyB; @leloykun; @YouJiacheng; @jadenj3o; @KonstantinWilleke, @alexrgilbert, @adricarda, @tuttyfrutyee, @vdlad; @ryanyang0, @vagrawal, @classiclarryd, @byronxu99, @varunneal, @EmelyanenkoK, @bernard24 / https://www.hiverge.ai/, @Gusarich, @li_zichong, @akash5474, @snimu, @roeeshenberg, @ChrisJMcCormick, @dominikkallusky, @acutkosky, @manikbhandari, @andrewbriand, @jrauvola, @soren_dunn_, @photon_mz, @srashedll, @dhrvji, @EmmettBicker, @dualverse-ai, @sisovicm, @moof2x, @samacqua

运行当前记录

要运行当前记录，请执行以下命令。

git clone https://github.com/KellerJordan/modded-nanogpt.git && cd modded-nanogpt
pip install -r requirements.txt
# 仅下载前 9 亿个训练 tokens 以节省时间
python data/cached_fineweb10B.py 9
./run.sh

如果 ./run.sh 报错 torchrun: command not found，请将 torchrun 添加到路径。

注意：首次运行代码时，torch.compile 将增加约 7 分钟的延迟。

官方记录使用来自 https://app.primeintellect.ai/ 的 8 块 NVIDIA H100 GPU 计时。PrimeIntellect 慷慨赞助了近期的验证运行。

备选方案：使用 Docker 运行（推荐用于精确计时）

如果 CUDA 或 NCCL 版本与您当前系统设置不兼容，Docker 是一个有用的备选方案。此方法标准化了 CUDA、NCCL、CUDNN 和 Python 的版本，减少了依赖性问题并简化了设置。
注意：系统上必须已安装 NVIDIA 驱动（当仅有 NVIDIA 驱动和 Docker 可用时很有用）。

git clone https://github.com/KellerJordan/modded-nanogpt.git && cd modded-nanogpt
# 构建 Docker 镜像
sudo docker build -t modded-nanogpt .
# 运行数据下载（例如 8 亿 tokens）
sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt python data/cached_fineweb10B.py 8
# 运行训练脚本
sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt sh run.sh

要获取交互式 Docker 环境，可以使用

sudo docker run -it --rm --gpus all -v $(pwd):/modded-nanogpt modded-nanogpt bash

世界纪录历史

以下是关于以下竞争性任务的世界速度纪录的历史演进：

使用 8 块 NVIDIA H100，将神经网络训练到 ≤3.28 的 FineWeb 验证损失。

注意：3.28 的目标是为了匹配 Andrej Karpathy 的 GPT-2 (small) 复现。