ARTICLE DETAIL

资讯详情

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

优化器分组粒度下钻:modded-nanogpt 将 K/V 按 Paired-Head 拆分为独立 Muon 组,以更少步数守住 3.28 验证损失

优化器分组粒度下钻:modded-nanogpt 将 K/V 按 Paired-Head 拆分为独立 Muon 组,以更少步数守住 3.28 验证损失 人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载本篇技术指南基于仓库记录 records/track_1_short/2026-04-08_PairedHeadMuon/README.md 展开讲解 modded-nanogpt 在 124M GPT-2「90 秒速跑」赛道上的一个关键优化器改动不再让一层里的所有注意力头共用同一个 Muon 正交化组而是把 K/V 权重按head-pair拆成各自独立的组使每个 head-pair 拥有独立的极分解polar decomposition。读完本文你将掌握这次改动的动机、实现粒度、实验收益每步损失改善、总步数减少 10 步、约 0.4s 的墙钟节省以及该思想在当前仓库 ANVIL 优化器track_1_short/optim/anvil.py中的源码级落地方式。1. 背景90 秒速跑与 3.28 达标线modded-nanogpt 的目标是用约 90 秒的墙钟时间在 8×H100 单机上把 124M 参数GPT-2 规模的语言模型在 FineWeb 数据上训练到指定质量。该赛道有一条明确的验收线最终验证集交叉熵final val CE均值 ≤ 3.28且所有 run 计入统计后续的 ANVIL2 记录在 records/track_1_short/2026-08-30_ANVIL2/README.md 中再次确认了这条 gate。在这样的短跑里每一秒都珍贵训练步数上限直接决定墙钟时间而每减少一步都要靠优化器每步进步更大来弥补。本记录所解决的问题正是这个——在不额外增加训练步数的前提下让每一步的损失下降更多。2. 基线做法每层一个(hdim, hdim)的 Muon 组在介绍改动前先看基线优化器如何处理注意力投影矩阵。Muon 是 modded-nanogpt 用于投影矩阵Q/K/V/O、MLP 权重的主力优化器它对梯度做正交化通过 Newton–Schulz 迭代逼近极分解/极坐标分解后再按学习率更新其余参数embedding、标量、gate 权重等仍交给 Adam。基线实现中Muon 把每层的Q、K、V、O 各当作一个(hdim, hdim)组hdim即该层所有 head 拼起来的宽度也就是说一层只有 4 个 Muon 组Q 组、K 组、V 组、O 组该层全部 head 的行向量被打包进同一个矩阵共享同一次正交化同一组内的 Gram 矩阵是完整的(hdim, hdim)尺寸。这样做的好处是实现简单、kernel 调用次数少坏处是粒度太粗见下节动机。3. 动机head 专化让「一锅烩」的正交化不再最优文档明确指出改动的直觉来源Given heads often specialize, it seemed intuitive that treating eachheadas its own Muon group would further improve optimization.在训练充分的多头注意力中不同 head 往往专化出不同角色有的负责局部语法、有的负责长程依赖、有的负责特定位置模式。当这些行为各异的 head 被拼进同一个矩阵做正交化时它们的更新会互相平均、互相牵制相当于用一个「平均方向」代替了每个 head 各自的最优方向。而本仓库的模型在部分层使用了paired head attention相邻两个 head 的 query 互相对 key 做注意力这更强化了「以 head-pair 为基本优化单元」的合理性——配对的两个 head 在功能上天然绑定。当前仓库在 track_1_short/model/gpt.py 中把PAIRED_HEAD_LAYERS定义为(0, 2, 5)即第 0、2、5 层使用 paired head attention其前向语义可见 track_1_short/model/attention.pyPaired heads: adjacent heads queries attend to each others keys. Two copies of the input stream are interleaved (q, k already are, by the norm/rotary kernel), which doubles each sequences length and halves the effective window.配对 head 的序列被交错、有效窗口减半——这两个 head 共享同一条信息通路把它们作为一个优化组既符合结构也减少了组间干扰。4. 核心改动K/V 按 head-pair 独立成组基于上述动机本次改动即本记录对应的 PR把K/V 权重按 head-pair 拆成独立的 Muon 组拆分对象K 和 V文档原文I split K/V into per-head-pair groups (because of paired head attention in some layers)分组单位相邻的两个 head 组成一个 pair一个 pair 对应一个独立的 Muon 组结果每个 head-pair 获得独立的极分解polar decomposition正交化不再跨 head 共享。为什么只拆 K/V 而不是全部四个矩阵文档的解释是 paired head attention 只存在于部分层且 head-pair 是信息处理的天然单元。此外拆分 K/V 还有一个附带收益——文档在括号里点了一句splitting Muon groups further should reduce flops。这一点可以从计算上理解正交化/白化迭代的主要成本是组内 Gram 矩阵的乘法一组d×d矩阵的乘法复杂度约为O(d³)若把一组拆成g个(d/g)大小的组总成本约为g·(d/g)³ d³/g²组数越多、单组越小总 FLOPs 反而越省代价是 kernel 启动次数变多这正是下节时间开销的来源。5. 实验结果10 步差距与约 0.4s 收益文档给出了三组对照实验每组多 run报告均值和标准差这是理解收益最核心的证据完整继承如下Runs Steps Time μ Time σ Time /- Time p Loss μ Loss σ Loss /- Loss p baseline-1490 4 1490 93.0205 0.1629 0.0000 nan 3.2782 0.0010 -0.0018 0.0187 baseline-1480 8 1480 92.5491 0.0588 -0.4714 0.0040 3.2800 0.0017 0.0000 0.5000 this PR 8 1480 92.6259 0.2129 -0.3946 0.0038 3.2788 0.0009 -0.0012 0.0029解读这张表能精确还原这个 PR 的论证逻辑baseline-1490对照组4 runs1490 步、均值耗时 93.02s、Loss μ 3.2782。它达标 3.28是「用满步数」的基线。baseline-1480对照组8 runs在未做任何优化器改动的前提下单纯砍掉 10 步Loss μ 升到 3.2800恰好越过 3.28 达标线不达标Loss p 0.5000 说明它与 1490 步基线相比没有统计意义上的损失改善。this PR8 runs同样只跑 1480 步但配合 per-head-pair Muon 分组后Loss μ 3.2788重新压回 3.28 以内且 Loss p 0.0029远小于 0.05改善是统计显著的与baseline-1490相比 Loss 差异 -0.0012、p 0.0187说明它甚至接近或打平多跑 10 步的完整基线。结论链非常清晰per-head-pair 分组提升了每步的损失下降幅度loss/step足以补偿砍掉的 10 步若不同时引入分组砍 10 步必然超标。最终收益是墙钟时间1480 步版本约 92.6s比 1490 步的 93.0s 少约 0.4s同时验证损失还稳在达标线内方差 σ 0.0009是三组中最低的之一。6. 成本账约 0.1% 的每步开销文档对代价给出了诚实的量化Treating each head-pair in K and V as its own Muon group improved loss per step while increasing time/step by ~.1%即每步耗时增加约 0.1%。原因正是上一节分析的组变多、每组变小虽然总 FLOPs 下降但正交化相关 kernel 的启动次数变多、调度开销上升。作者同时指出这个 overhead 大概率还能继续压this overhead can probably be reduced因为把 Muon 组进一步拆分本身就会减少 FLOPs——也就是说当前实现的 kernel 层还没完全吃到「拆组省算力」的红利。值得注意的是作者对时间收益的判断是稳健的加速不来自单步变快而来自步数变少。即便单步变慢 0.1%只要总步数从 1490 降到 1480墙钟依然净省。7. 从 PR 到当前仓库分组思想在 ANVIL 中的源码级落地本 PR 目录下只有 README但分组思想随后被后续记录继承并演化。当前仓库的优化器已从「Muon Newton–Schulz」演进为ANVILtwin-rail momentum 白化级联见 track_1_short/optim/anvil.py而「按 head 粒度拆组」的设计被保留并扩展到了 QK 侧。对照源码可以看清落地细节7.1 QK bankper-head-pair 的 ANVIL 组track_1_short/model/gpt.py 中 QK bank 的注释明确写道QK bank: per-head-pair ANVIL groups for Q, K weights. Each pair of adjacent heads gets its own independent whitening: a slots 2 * num_heads heads (Q heads, then K heads) are num_heads groups of two full-width heads.即每个 bank slot对应一个注意力层里Q 的 heads 和 K 的 heads 被切成num_heads个「两个全宽 head」的组每组独立做白化whitening即 ANVIL 版的正交化。代码上qk_groups_per_slot num_heads7 个注意力层共num_slots * qk_groups_per_slot 42个真实组由于参数在 8 卡间按world_size均分组数需要对齐到 8 的倍数next_multiple_of_n(num_qk_groups, nself.world_size)补齐到48最终qk_bank形状为(48, 2 * head_dim, model_dim)head_dim128即(48, 256, 768)VO bank 则保留 per-layer 粒度num_vo_real 2 * num_slots 14补齐到(16, hdim, model_dim)gpt.py。这恰好说明「分组粒度」是可独立调节的旋钮Q/K 沿用 PR 的 head-pair 思路扩展到 Q 侧V/O 在后续记录中回到了逐层粒度——不同矩阵的最优粒度可以不同。7.2 组如何变成优化器里的 bank 与 chunkANVIL 优化器track_1_short/optim/anvil.py通过reshape属性把一组矩阵声明为一个 bank再按世界大小切成 chunkANVIL 参数必须带.reshape且reshape[0]必须能被world_size整除每个 rank 的chunk_size reshape[0] // world_sizeanvil.py每个 bank 在_init_banks中分配一个持久的AnvilBankanvil.py包含 reduce-scatter 落地的grad、twin-rail 的velocityfp32[2, chunk, rows, cols]、逐 lane 能量均衡器的lane_energy、以及低 16 位mantissa缓冲梯度先reduce_scatter到本 rank 的 chunk更新完成后all_gather回完整矩阵anvil.py。组数越多每个 chunk 越小通信与白化都更细粒度。7.3 参数表哪些矩阵走 ANVIL、哪些走 Adamtrack_1_short/training.py 的param_table是运行时的配置入口qk_bank、vo_bank、mlp_bank三个投影矩阵 bank 全部走optim: anvilcomms: sharded其余scalars、smear_gate、ve_gate_bank、lm_head、embed、value_embeds、MUDD 参数族等走 Adam。ANVIL 的默认超参在同一文件可见training.pylr0.023、momentum0.95Nesterov 前瞻每步由get_rail_beta重写、beta20.9lane 能量 EMA 衰减、weight_decay2.25。通信调度上scatter_order把qk_bank、vo_bank、mlp_bank排在最先training.py保证它们的 all-gather 是下一步 fp8 权重刷新最早等到的两个结果。7.4 后续记录的印证在 2026-08-30 的 ANVIL2 基线代码里records/track_1_short/2026-08-30_ANVIL2/baseline/ 内各.txt的GPT类定义同样保留着注释QK bank: per-head-pair Muon groups for Q, K weights. Each pair of adjacent heads gets its own independent polar express orthogonalization.可见本 PR 奠定的「head-pair 独立正交化」是 ANVIL2 及其后续架构的直接前身——只是正交化算法从 Newton–Schulz 换成了 Polar Express / ANVIL 白化级联参见 records/track_1_short/2026-08-30_ANVIL2/README.md分组粒度本身被完整继承。8. 时间备注与复现注意文档末尾附了一段诚实的 caveat值得原样保留给复现者Note on timing:not sure why my 8xh100 is so slow -- ~8% slower per step for the baseline. However, given that the speedup comes from reducing steps, I think this will hold in a more controlled setup.即作者所在 8×H100 环境下 baseline 每步比预期慢约 8%存在环境层面的不确定性但由于本改动的收益来自减少步数1480 步达成原本 1490 步的质量而非单步提速这个结论在更受控的机器上应当依然成立。若要在当前仓库复现该记录链路的运行方式仓库根目录提供了入口python data/cached_fineweb10B.py 9准备数据后执行bash run.sh入口脚本 run.sh、训练主程序 train_gpt.py完整环境CUDA 13 base image、torch 2.10.0cu128、kernels 0.16.1 等由仓库根目录的 Dockerfile 与 requirements.txt 固定更细的运行注意驱动版本 ≥580、离线加载 patched FA3 内核、libcudart.so.13等可参考 ANVIL2 记录的 Requirements 一节。小结Paired Head Muon 分组是一个典型的「优化器粒度调优」案例把注意力投影矩阵的 Muon 正交化从「每层一锅烩」下沉到「每个 head-pair 独立成组」利用 head 的专化结构换取每步更大的损失下降从而在总步数减少 10 步约 0.4s的前提下仍守住 3.28 验证损失目标且损失改善统计显著p 0.0029。该思想随后沉淀为当前仓库 ANVIL 优化器中 QK bank 的 per-head-pair 分组设计成为 124M 速跑链路的一部分——对任何关心「优化器分组粒度与训练效率」的读者这都是一个从动机、实验到源码落地都完整可查的范本。赞分享人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐Modded-NanoGPT 中的 SOAP 优化器实战3.15B Tokens 达成 3.28 验证损失的样本效率记录Modded NanoGPT 中的 SOAP 优化器实战3.15B Tokens 达成 3.28 验证损失的样本效率记录 本文基于 Modded NanoGP人工智能大模型预训练分布式训练模型优化深度学习modded-nanogpt SimplifyHC缓存层 7 激活作为统一注意力输入124M 训练提速并稳定压入 3.28 损失以下modded nanogpt SimplifyHC缓存层 7 激活作为统一注意力输入124M 训练提速并稳定压入 3.28 损失以下 本篇记录来自 reco人工智能大模型预训练分布式训练模型优化深度学习IBAnimatable模块化设计将动画组件拆分为独立Pod的实践方法IBAnimatable模块化设计将动画组件拆分为独立Pod的实践方法 模块化拆分的必要性与目标 随着iOS应用复杂度提升单一动画库往往面临体积膨胀、编译耗移动开发UI组件上一篇YOLOv7多任务学习目标检测与语义分割联合训练下一篇Easy Peasy 终极指南从零构建完整生产级应用的最佳实践创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表