- 人工智能
- 大模型
- 预训练
- 分布式训练
- 模型优化
- 深度学习
【免费下载链接】modded-nanogpt
NanoGPT (124M) in 90 seconds
本篇文章以 modded-nanogpt 仓库 records/track_1_short/2024-11-10_UNetDoubleLr/ 目录下的纪录文档为核心,结合同目录可复现日志与仓库源码,完整剖析这条由 Brendan Hogan Rappazzo(@brendanh0gan)创造的 7.23 分钟纪录:两项核心改动(向 Transformer 中引入 U-Net 风格的编码器-解码器跳跃连接、将优化器学习率整体翻倍)究竟改了什么、为什么能同时提升收敛速度与吞吐。读完本文,你将掌握该纪录的全部超参、可复现日志的解读方法,以及这套 skip connection 思想在今日 modded-nanogpt 架构中的演进脉络。
纪录背景:在 8 张 H100 上冲刺 3.28 验证损失
modded-nanogpt 的 track 1(short track)是一场"速度竞赛":用 8 张 NVIDIA H100 GPU 训练一个约 1.24 亿参数(124M)的 GPT-2 规模模型,使其在 FineWeb 验证集上达到 3.28 的交叉熵损失,比拼的是完成这一目标所需的墙钟时间。3.28 这一阈值源自 Andrej Karpathy 在 llm.c 中用 45 分钟达到的 GPT-2 复现成绩,而 modded-nanogpt 通过架构与优化器的持续改进,把这个时间不断压缩(详见仓库根 README.md)。
2024-11-10 这一天,@brendanh0gan 提交的新纪录将时间从 11 月 6 日纪录的7.8 分钟(见 2024-11-06_ShortcutsTweaks 纪录)进一步推进到7.23 分钟。该纪录文档记录的两项核心改动是:
- 向 Transformer 中加入 U-Net 风格的跳跃连接(U-net-like skip connections)
- 将学习率翻倍(Doubled the learning rate)
纪录文档明确说明,可复现日志就存放在同目录下,代码可以直接从日志中提取,这为后人复现与验证提供了完整依据。
核心改动一:U-Net 式跳跃连接
从"embed shortcut"到 U-Net 的动机
在 11 月 6 日的纪录中,模型已经引入了"embed shortcut"(x0 注入)与 value residual 等跳跃结构——每个 Block 的 forward 接收x0并用可学习参数lambdas混合:
class Block(nn.Module): def forward(self, x, v1, x0): x = self.lambdas[0] * x + self.lambdas[1] * x0 x1, v1 = self.attn(F.rms_norm(x, (x.size(-1),)), v1) x = x + x1 x = x + self.mlp(F.rms_norm(x, (x.size(-1),))) return x, v1@brendanh0gan 的改进思路是把这种"浅层注入"升级为完整的U-Net 编码器-解码器结构:将 12 层 Transformer 一分为二,前半段作为编码器逐层下采样特征,后半段作为解码器,将编码器对应层的输出通过可学习的加权跳跃连接重新引入解码器输入,从而为深层提供更直接的梯度与信息通路(代码见可复现日志 c87bb826-797b-4f37-98c7-d3a5dad2de74.txt)。
实现:encoder/decoder 划分与 learnable skip weights
在GPT.__init__中,模型将 12 层划分为 6 层编码器与 6 层解码器,并为每条跳跃连接配备一个独立的可学习标量权重:
# U-net design by @brendanh0gan self.encoder_layers = config.n_layer // 2 # Half of the layers for encoder self.decoder_layers = config.n_layer - self.encoder_layers # Remaining for decoder # Add learnable skip connection weights for decoder layers self.skip_weights = nn.Parameter(torch.ones(self.decoder_layers))前向传播中,编码器阶段把每层的输出暂存进skip_connections列表;解码器阶段则从列表尾部(pop())取出对应的编码器输出,乘上可学习权重后与当前残差流相加,再送入解码器 Block:
# Encoder pass - process only the first half of the blocks for i in range(self.encoder_layers): x, v1 = self.transformer.hi skip_connections.append(x) # Store the output for skip connections # Decoder pass - process the remaining blocks with weighted skip connections for i in range(self.decoder_layers): skip_connection = skip_connections.pop() # Get the corresponding encoder output # Apply learnable weight to skip connection weighted_skip = self.skip_weights[i] * skip_connection x, v1 = self.transformer.hself.encoder_layers + i关键设计点:
- 对称连接:编码器第 k 层的输出(k = 0..5)注入解码器第 5-k 层,形成 0↔5、1↔4、2↔3 的镜像配对,与 U-Net 的对称拓扑一致。
- 可学习权重:
skip_weights初始化为全 1,随训练由 Adam 优化器更新,模型可以自行决定每条短路的强度,而不是硬编码相加。 - 参数归属:
skip_weights是 1 维标量参数,被归入scalar_params(scalar_params = [p for p in params if p.ndim < 2] + [raw_model.skip_weights]),与各层的 lambda、lamb 等标量一起由第四个 Adam 优化器管理,与矩阵参数使用的 Muon 优化器分离。
数据流剖析
从GPT.forward可以还原完整的计算图:embedding 之后做 RMSNorm 得到x0;v1 = None意味着第一个 Block 的注意力层用自身投影的 V 初始化 value residual;编码器 6 层依次运行并保存快照;随后解码器以x + weighted_skip为输入继续运行。最终输出经 RMSNorm 与lm_head(CastedLinear,权重零初始化),并通过30 * torch.tanh(logits / 30)的 tanh logit 软封顶后计算交叉熵损失。
与 11 月 6 日纪录相比,CausalSelfAttention中仍保留了lamb(默认 0.5)控制的 value residual 混合v = (1 - self.lamb) * v + self.lamb * v1.view_as(v),以及 QK-Norm、Rotary、ReLU² MLP、c_proj零初始化等既有技术——U-Net 跳跃连接是在这套基础上叠加的增量改进,而不是推倒重来。
核心改动二:学习率翻倍
纪录的第二项改动是"学习率翻倍"。将本纪录代码(c87bb826…txt)与 11 月 6 日纪录代码(见 2024-11-06_ShortcutsTweaks 目录)中四路优化器的学习率逐一对比:
| 优化器 | 优化对象 | 11/06 纪录 lr | 11/10 本纪录 lr | 变化 |
|---|---|---|---|---|
| Adam (fused) | token embeddingwte | 0.3 | 0.6 | 翻倍 |
| Adam (fused) | lm_head权重 | 0.002 | 0.008 | 4 倍 |
| Muon | 全部 2D 矩阵参数 | 0.02 | 0.04 | 翻倍 |
| Adam (fused) | 标量参数 +skip_weights | 0.02 | 0.04 | 翻倍 |
其中承载绝大多数参数更新的 Muon 优化器学习率从 0.02 提高到 0.04,这正是"Doubled the learning rate"的主体含义;embedding 与标量参数同步翻倍,lm_head则抬升得更高。
一个值得注意的配套细节:本纪录将 Muon 的momentum warmup从 11/06 纪录的固定 momentum 升级为前 500 步线性插值:
# momentum warmup for Muon frac = min(step/500, 1) optimizer3.param_groups[0]['momentum'] = (1 - frac) * 0.85 + frac * 0.95即 momentum 从前 500 步的 0.85 逐渐升温到 0.95。在整体学习率翻倍、训练早期梯度噪声较大的情况下,较低的初始 momentum 有助于抑制早期震荡,待模型进入平稳收敛区间后再加大动量——这是让高学习率"接得住"的关键工程配套,也是该纪录能够在 3000 步内稳定收敛到 3.28 阈值以下的重要保障。
完整超参与可复现配置
该纪录并非独立的训练脚本,而是 modded-nanogpt 家族中的一员:训练代码、超参数与完整 3000 步日志都被原样写入可复现日志文件 c87bb826-797b-4f37-98c7-d3a5dad2de74.txt(文件首尾各有一行===分隔符,中间即完整 Python 源码与日志输出)。核心超参数如下:
| 类别 | 参数 | 值 |
|---|---|---|
| 数据 | input_bin | data/fineweb10B/fineweb_train_*.bin |
| 数据 | input_val_bin | data/fineweb10B/fineweb_val_*.bin |
| 优化 | batch_size(全局) | 8 × 64 = 512 条序列 |
| 优化 | device_batch_size(每卡) | 64 条序列 |
| 优化 | sequence_length | 1024 tokens |
| 优化 | num_iterations | 3000 |
| 优化 | warmup_iters | 0 |
| 优化 | warmdown_iters | 900(三角/梯形 LR 调度) |
| 优化 | weight_decay | 0 |
| 评估 | val_loss_every | 每 125 步 |
| 评估 | val_tokens | 10485760(固定以便对比) |
| 模型 | 层数 / 头数 / 维度 | 12 / 6 / 768(head dim 128) |
| 模型 | vocab_size | 50304(50257 向上取整到 128 的倍数) |
学习率调度采用线性 warmdown:前num_iterations - warmdown_iters(即 2100)步保持常数 1.0,最后 900 步线性衰减到 0。由于warmup_iters = 0,训练从第一步就以满学习率起步,这与 momentum warmup 形成了"LR 满、动量渐进"的组合。
训练管线与系统实现
四路优化器分工
与 modded-nanogpt 一脉相承,该纪录采用"Muon 管矩阵、Adam 管其余"的参数分工(代码中的optimizer1至optimizer4):
optimizer1:Adam 优化 token embedding,lr 0.6;optimizer2:Adam 优化lm_head,lr 0.008(embedding 与 lm_head 不共享权重);optimizer3:Muon 优化所有 2D 矩阵参数([p for p in params if p.ndim == 2]),lr 0.04,momentum 0.95(前 500 步从 0.85 升温),nesterov;optimizer4:Adam 优化标量参数与skip_weights,lr 0.04。
Muon 优化器的实现在日志中同样完整保留:内部先跑 SGD-momentum,再用 Newton-Schulz 五次迭代(系数a, b, c = (3.4445, -4.7750, 2.0315),backend_steps=5)对每个 2D 参数做正交化,得到近似 UV^T 的更新方向,并按max(1, g.size(0)/g.size(1))**0.5做尺度补偿。在 8 卡分布式环境中,每个矩阵参数按i % WORLD_SIZE == RANK分配给对应 GPU 计算更新,随后通过一次dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM)同步全部更新向量,再反序列化应用到参数上——这是该仓库早期"分片 Muon"实现,注释里也点明了其约束:"多 4 层整除的 Transformer 可以恰好完美地分配到 8 张 GPU"。
分布式数据加载与训练循环
数据加载使用自定义的DistributedDataLoader:读取带魔数头(20240520)的.bin数据分片,每个进程从process_rank * B * T偏移处开始消费数据,next_batch()每次返回长度为 B×T 的输入/目标对。训练采用梯度累积:全局 batch 512 条序列、每卡 64 条,故train_accumulation_steps = 512 // (64 × 8) = 1,即每步恰好一张卡一个 batch;累积逻辑仅在 batch 较大时才生效(model.no_sync()跳过中间步的梯度同步)。
计时方面,代码刻意忽略前 10 步的编译与冷启动开销,从第 11 步开始统计training_time_ms,并在每 125 步与最后一步做验证集评估,保证val_tokens = 10485760固定不变,从而使不同纪录之间的验证损失可比。
结果与验证:7.23 分钟的构成
日志末尾给出了完整 3000 步的训练轨迹,几个关键数据点:
- 起步:
step:0/3000 val_loss:10.8258,训练损失从 10.8 量级开始; - 中期:
step:125/3000 val_loss:4.5262、step:250/3000 val_loss:4.0672、step:2750/3000 val_loss:3.3033; - 收尾:
step:3000/3000 val_loss:3.2753,稳稳低于 3.28 阈值; - 性能:稳态平均每步约145ms(
step_avg),全程训练时间433570ms ≈ 7.23 分钟,与纪录文档宣称的 7.23 分钟完全吻合。
对比 11 月 6 日纪录的 7.8 分钟,本次提速约 30 秒,主要来自两方面:U-Net 跳跃连接带来的收敛效率提升(相同步数下验证损失下降更快、训练更稳),以及学习率翻倍后每步参数更新的"步子"更大。值得注意的是,本纪录的吞吐(约 145ms/步)与前期纪录基本持平,说明 U-Net 跳跃连接几乎没有增加计算与通信开销——加权标量乘法与残差相加的成本可忽略不计,而它换来的梯度通路价值显著。
日志还记录了运行环境:PyTorch 2.5.1+cu124、8 张 H100 80GB(每卡显存占用约 5GB)、NCCL 后端、CUDNN attention(日志中显式关闭 Flash SDP、启用 CUDNN SDP,注释说明"在 PyTorch 2.5.1 中 CUDNN attention 比 Flash 快约 4ms"),以及torch._inductor.config.coordinate_descent_tuning = True等编译调优开关。
从 7.23 分钟纪录到今日架构的演进
7.23 分钟纪录在 speedrun 时间轴上只是一个节点,但其"数据依赖、可学习的跳跃连接"思想在此后的纪录中不断深化。翻阅当前仓库源码 track_1_short/model/gpt.py,可以看到今天的主干代码已经演进为 11 层的 MUDD(Multiway Dynamic Dense Connections)架构:不仅保留 x0 注入与中间层缓存(CACHE_LAYERS = (3, 7)),更将跳跃连接升级为逐 token 动态计算的系数——forward_mudd通过一个 64 维隐层的 MLP 为每个位置生成跳跃系数,最后一层甚至用 14 个系数同时混合 cache[0]、cache[7]、当前 x 等多个来源(LAST_LAYER_MUDD_COEFS = 14),并在残差流与 attention 的 V 上同时施加跳跃(v_mudd)。此外,embedding 哈希 n-gram 表(bigram/trigram 注入)、value embedding 等"旁路信息注入"也都是对"给深层模型铺直路"这一思路的延伸。
换句话说,7.23 分钟纪录里那对简单的skip_weights[i] * skip_connection标量权重,正是后来 MUDD 中"每条跳跃连接都拥有独立、可学习、甚至逐 token 变化系数"这一家族设计的早期雏形,也为今天 README 中"Skip connections from embedding to every block as well as from block 3 to 6"与"MUDD skip connections to residual stream and attention values"的技术清单提供了谱系背景。
小结
2024-11-10 的 7.23 分钟纪录,是 modded-nanogpt 速度竞赛史上一次典型而优雅的增量胜利:没有引入新的算子或复杂的调度,仅仅是把 Transformer 组织成 U-Net 式的编码器-解码器对称拓扑、为每条捷径配备可学习权重,并把四路优化器的学习率整体翻倍、配合 Muon momentum warmup,就在不牺牲吞吐的前提下把收敛速度推进了一个台阶。它同时印证了两条在后续演进中被反复验证的经验:深度网络中的短路径连接是低成本高回报的结构性收益,以及优化器超参数应与新架构协同调整,而非孤立看待。
如果你想亲手复现或深挖,可复现日志 c87bb826-797b-4f37-98c7-d3a5dad2de74.txt 内含全部代码与 3000 步输出;相邻纪录(11/06 ShortcutsTweaks、11/09 Replicateleloykun)则提供了前后对照的基准。
- 人工智能
- 大模型
- 预训练
- 分布式训练
- 模型优化
- 深度学习
【免费下载链接】modded-nanogpt
NanoGPT (124M) in 90 seconds
相关推荐
AnythingSlider社区贡献指南:如何参与这款jQuery轮播插件开发与维护
AnythingSlider社区贡献指南:如何参与这款jQuery轮播插件开发与维护 AnythingSlider 是一款功能强大的jQuery轮播插件,自20
SillyTavern:3 步部署免费的 AI 角色对话引擎,接入 20+ 主流大模型
SillyTavern:3 步部署免费的 AI 角色对话引擎,接入 20+ 主流大模型 SillyTavern 是一个开源的 LLM 对话前端,面向深度玩家,帮
人工智能AI 应用交互助手前端Modded-NanoGPT训练日志:如何解析与可视化性能数据
Modded NanoGPT训练日志:如何解析与可视化性能数据 训练日志是优化深度学习模型的关键依据,Modded NanoGPT项目通过结构化记录和可视化工具
人工智能大模型预训练分布式训练模型优化深度学习
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考