ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

Modded-NanoGPT 纪录解读:U-Net 式跳跃连接与学习率翻倍如何将 124M GPT 训练压到 7.23 分钟

Modded-NanoGPT 纪录解读:U-Net 式跳跃连接与学习率翻倍如何将 124M GPT 训练压到 7.23 分钟 人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载本篇文章以 modded-nanogpt 仓库 records/track_1_short/2024-11-10_UNetDoubleLr/ 目录下的纪录文档为核心结合同目录可复现日志与仓库源码完整剖析这条由 Brendan Hogan Rappazzobrendanh0gan创造的 7.23 分钟纪录两项核心改动向 Transformer 中引入 U-Net 风格的编码器-解码器跳跃连接、将优化器学习率整体翻倍究竟改了什么、为什么能同时提升收敛速度与吞吐。读完本文你将掌握该纪录的全部超参、可复现日志的解读方法以及这套 skip connection 思想在今日 modded-nanogpt 架构中的演进脉络。纪录背景在 8 张 H100 上冲刺 3.28 验证损失modded-nanogpt 的 track 1short 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 shortcutx0 注入与 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, v1brendanh0gan 的改进思路是把这种浅层注入升级为完整的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_paramsscalar_params [p for p in params if p.ndim 2] [raw_model.skip_weights]与各层的 lambda、lamb 等标量一起由第四个 Adam 优化器管理与矩阵参数使用的 Muon 优化器分离。数据流剖析从GPT.forward可以还原完整的计算图embedding 之后做 RMSNorm 得到x0v1 None意味着第一个 Block 的注意力层用自身投影的 V 初始化 value residual编码器 6 层依次运行并保存快照随后解码器以x weighted_skip为输入继续运行。最终输出经 RMSNorm 与lm_headCastedLinear权重零初始化并通过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 纪录 lr11/10 本纪录 lr变化Adam (fused)token embeddingwte0.30.6翻倍Adam (fused)lm_head权重0.0020.0084 倍Muon全部 2D 矩阵参数0.020.04翻倍Adam (fused)标量参数 skip_weights0.020.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_bindata/fineweb10B/fineweb_train_*.bin数据input_val_bindata/fineweb10B/fineweb_val_*.bin优化batch_size全局8 × 64 512 条序列优化device_batch_size每卡64 条序列优化sequence_length1024 tokens优化num_iterations3000优化warmup_iters0优化warmdown_iters900三角/梯形 LR 调度优化weight_decay0评估val_loss_every每 125 步评估val_tokens10485760固定以便对比模型层数 / 头数 / 维度12 / 6 / 768head dim 128模型vocab_size5030450257 向上取整到 128 的倍数学习率调度采用线性 warmdown前num_iterations - warmdown_iters即 2100步保持常数 1.0最后 900 步线性衰减到 0。由于warmup_iters 0训练从第一步就以满学习率起步这与 momentum warmup 形成了LR 满、动量渐进的组合。训练管线与系统实现四路优化器分工与 modded-nanogpt 一脉相承该纪录采用Muon 管矩阵、Adam 管其余的参数分工代码中的optimizer1至optimizer4optimizer1Adam 优化 token embeddinglr 0.6optimizer2Adam 优化lm_headlr 0.008embedding 与 lm_head 不共享权重optimizer3Muon 优化所有 2D 矩阵参数[p for p in params if p.ndim 2]lr 0.04momentum 0.95前 500 步从 0.85 升温nesterovoptimizer4Adam 优化标量参数与skip_weightslr 0.04。Muon 优化器的实现在日志中同样完整保留内部先跑 SGD-momentum再用 Newton-Schulz 五次迭代系数a, b, c (3.4445, -4.7750, 2.0315)backend_steps5对每个 2D 参数做正交化得到近似 UV^T 的更新方向并按max(1, g.size(0)/g.size(1))**0.5做尺度补偿。在 8 卡分布式环境中每个矩阵参数按i % WORLD_SIZE RANK分配给对应 GPU 计算更新随后通过一次dist.all_reduce(updates_flat, opdist.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 阈值性能稳态平均每步约145msstep_avg全程训练时间433570ms ≈ 7.23 分钟与纪录文档宣称的 7.23 分钟完全吻合。对比 11 月 6 日纪录的 7.8 分钟本次提速约 30 秒主要来自两方面U-Net 跳跃连接带来的收敛效率提升相同步数下验证损失下降更快、训练更稳以及学习率翻倍后每步参数更新的步子更大。值得注意的是本纪录的吞吐约 145ms/步与前期纪录基本持平说明 U-Net 跳跃连接几乎没有增加计算与通信开销——加权标量乘法与残差相加的成本可忽略不计而它换来的梯度通路价值显著。日志还记录了运行环境PyTorch 2.5.1cu124、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 层的 MUDDMultiway 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-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐AnythingSlider社区贡献指南如何参与这款jQuery轮播插件开发与维护AnythingSlider社区贡献指南如何参与这款jQuery轮播插件开发与维护 AnythingSlider 是一款功能强大的jQuery轮播插件自20SillyTavern3 步部署免费的 AI 角色对话引擎接入 20 主流大模型SillyTavern3 步部署免费的 AI 角色对话引擎接入 20 主流大模型 SillyTavern 是一个开源的 LLM 对话前端面向深度玩家帮人工智能AI 应用交互助手前端Modded-NanoGPT训练日志如何解析与可视化性能数据Modded NanoGPT训练日志如何解析与可视化性能数据 训练日志是优化深度学习模型的关键依据Modded NanoGPT项目通过结构化记录和可视化工具人工智能大模型预训练分布式训练模型优化深度学习上一篇终极Copyparty性能优化指南如何识别和解决性能瓶颈下一篇掌握raylib法线贴图10个技巧让3D游戏细节完美呈现创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表