ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

modded-nanogpt Track 3 优化实录:Aurora 行均衡 polar + Contra-Muon + u/w-floor 将 3.28 验证损失推进到 3175 步

modded-nanogpt Track 3 优化实录:Aurora 行均衡 polar + Contra-Muon + u/w-floor 将 3.28 验证损失推进到 3175 步 人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载本篇技术指南以 modded-nanogpt 仓库 Track 3 优化 benchmark 的 Aurora 提交result #17为核心深入讲解 Aurora 如何通过**对角均衡diagonal equilibration**改造 Muon 的 polar 因子更新使其输出矩阵具备均匀的行范数leverage 均衡并与 Contra-Muon 扰动、u/w-floor 无权重衰减策略组合在 20 个种子、3200 步训练内以 3175 步达到 3.28 验证损失目标。读完本文你将掌握 Aurora 的数学动机与伪代码、其在 训练脚本 中的完整源码实现、整套优化器的超参配置以及 benchmark 的统计显著性判定规则。Aurora 是什么给 Muon 的 polar 加行均衡从 leverage 分数说起Muon 优化器的核心步骤是对动量混合后的梯度矩阵G做 polar 分解即计算polar(G)。但标准 polar 存在一个隐患对于非方阵G其结果会继承G的左奇异向量行范数的不均匀性——这些行范数就是所谓的leverage 分数leverage scores。行范数差异过大意味着更新在不同输出维度上的力度失衡。Aurora来自 tilde-research/aurora-release正是针对这一点提出改进不直接计算polar(G)而是先乘一个正对角矩阵D计算polar(D · G)其中D被选为使得结果的行范数均匀分布。具体来说对于 tall 矩阵Gm 行 n 列m n返回polar(D · G)对于 wide 矩阵Gm n返回polar(G · D)等效于先转置处理再转回均匀行范数的目标值为sqrt(min(m, n) / max(m, n))对于方阵Aurora 自动退化为标准 polar——因为此时没有任何 leverage 自由度可以利用。Aurora 的迭代伪代码给定形状为 m x nm n的动量混合梯度GAurora 按如下流程计算D 1 / ||G_i:|| # initial per-row inverse norms for k in range(pp_iterations): U polar(D * G) # standard Newton-Schulz polar if k pp_iterations - 1: row_sq sum(U_i:^2) D D * (n/m / row_sq)^pp_beta # damped fixed-point update of D return U要点解释初始化D初始化为每行 L2 范数的倒数1 / ||G_i:||即对每行做一次归一化内层 polarpolar(D · G)采用标准的Newton-Schulz 迭代项目代码中的zeropower_via_newtonschulz512 次迭代阻尼不动点更新在非最后一次迭代中根据当前U的行平方范数row_sq与目标值n/m的比值以指数pp_beta更新D。这是一个阻尼不动点迭代逐步把行范数拉向均匀。作者在文档中给出的关键经验值是pp_iterations2, pp_beta0.5此时经过 Newton-Schulz polar 后行范数与均匀分布的偏差在2% 以内。而pp_iterations1则退化为Muon EQ基线——即先做一次逐行 L2 归一化再 polar。源码实现从伪代码到可运行的 Aurora 正交化本次提交把 Aurora 完整实现在 日志文件内的训练脚本 中该文件遵循 benchmark 规则将全部代码嵌入日志以便复现。核心函数如下def aurora_orthogonalize(G: Tensor, pp_iterations: int 2, pp_beta: float 0.5, eps: float 1e-7) - Tensor: Aurora: leverage-uniform polar via diagonal preconditioning. Reference: tilde-research/aurora-release. m, n G.size(-2), G.size(-1) if m n: return zeropower_via_newtonschulz5(G) transposed m n if transposed: G G.mT m, n n, m G32 G.to(torch.float32) target_row_sq n / m row_norm G32.norm(dim-1, keepdimTrue).clamp_(mineps) D 1.0 / row_norm for k in range(pp_iterations): U zeropower_via_newtonschulz5(D * G32) if k pp_iterations - 1: row_sq U.to(torch.float32).pow(2).sum(dim-1, keepdimTrue).clamp_(mineps * eps) D D * (target_row_sq / row_sq).pow(pp_beta) return U.mT if transposed else U实现细节与伪代码一一对应方阵短路m n时直接返回标准 Newton-Schulz polar避免无意义的均衡开销wide 矩阵处理m n时先转置按 tall 情形处理后再U.mT转回数值精度G先转成float32计算G32row_norm与row_sq均做clamp防止除零/数值下溢目标行平方范数target_row_sq n / m与文档中的sqrt(min/max)目标平方一致。底层的 Newton-Schulz 迭代zeropower_via_newtonschulz5是本仓库 Muon 系实现的通用底层原语12 次迭代系数 a2, b-1.5, c0.5先对矩阵做谱范数归一化除以norm 1e-7再反复迭代X a*X BX。Aurora 只是把输入从G换成D · G复用了这套标准迭代因此新增代码量极小。组合优化栈Aurora Contra-Muon u/w-floor本次提交README 中标注为 result #17对应 benchmark 历史表 的第 17 行并非单独使用 Aurora而是把它叠加在此前的最优结果 #11Contra-Muon u/w-floor之上。README 明确指出与 #11 相比唯一的算法改动就是把 Newton-Schulz 步骤中的polar(G)替换为polar(D · G)其余全部继承。Contra-Muon减去算子范数方向的扰动Contra-Muon 修改 Muon 的动量更新从更新中减去CONTRA_MUON / 2倍的算子范数归一化后的动量梯度详见 Contra-Muon 结果文档。在本次代码中muon_update的流程是torch.compile def muon_update(grad, momentum, mu0.95, nesterovTrue, pp_iterations2, pp_beta0.5): momentum.lerp_(grad, 1 - mu) update grad.lerp_(momentum, mu) if nesterov else momentum normalized_grad scale_to_unit_operator_norm(update.clone()) update aurora_orthogonalize(update, pp_iterationspp_iterations, pp_betapp_beta) polar_fro update.norm() update update - CONTRA_MUON / 2 * normalized_grad update update * polar_fro / torch.clamp(update.norm(), min1e-10) update * max(1, grad.size(-2) / grad.size(-1))**0.5 return update其中scale_to_unit_operator_norm用 5 次幂迭代估计矩阵的算子范数方向update - CONTRA_MUON / 2 * normalized_grad即 Contra-Muon 扰动随后用polar_fro / update.norm()保持扰动前后 Frobenius 范数不变最后乘以宽高比的平方根max(1, m/n)^0.5完成宽高比缩放。CONTRA_MUON0.4为锁定超参。u/w-floor替代权重衰减的范数地板权重衰减被weight_decay0显式关闭转而使用u/w-floor策略每次更新后检查||u||_F / ||w||_F更新范数与权重范数之比若低于目标TARGET_UW就把更新放大到该地板# u/w-floor: scale update so ||u||_F / ||w||_F TARGET_UW. p_fro p.float().norm().clamp_min(1e-8) u_fro update.float().norm().clamp_min(1e-8) cur_uw u_fro / p_fro scale torch.where(cur_uw TARGET_UW, TARGET_UW * p_fro / u_fro, torch.ones_like(p_fro)) update update * scale.to(update.dtype) p.add_(update, alpha-group[lr])本次锁定TARGET_UW0.35。该机制源自 result #9NorMuon u/w-floor作用是给更新一个相对于权重规模的最小幅度下限替代传统解耦权重衰减的正则化角色。完整的超参配置表README 给出的配置表如下全部为锁定值字段值optimizerAurora Contra-Muon u/w-floorlr(Aurora)0.0375weight_decay0由 u/w-floor 取代pp_iterations2pp_beta0.5CONTRA_MUON0.4TARGET_UW0.35train_steps3200reported step3175均值曲线首次越过 n20 显著性门槛的最小步数n20种子 1..20在 源码实现 中这些超参的落位为Muon优化器以lr0.0375, weight_decay0, pp_iterations2, pp_beta0.5作用于model.blocks中所有ndim 2的参数而 token embedding、输出投影与 1D 参数bias/gains交给辅助的 AdamWbetas(0.8, 0.95), eps1e-10, weight_decay0, fusedTrueembed lr0.3、proj lr1/320、1D 参数 lr0.01。学习率采用先稳定后衰减的 WSD 风格调度前1 - cooldown_frac默认 0.7比例保持eta1.0之后线性下降到 0。实验结果20 个种子的统计显著性判定标准根据 benchmark 规则提交必须满足训练步数内平均验证损失低于 3.28且通过单侧 z 检验假设σ0.0013要求p 0.001即 3.09σ 0.004 的差距。精确条件为(3.28 - avg_loss) * num_runs^0.5 0.004。同时规则允许训练接近结束时每 25 步打印一次验证损失并跨所有试验统一选择最早达到统计显著的步数作为报告步数禁止基于验证损失做逐 trial 的 early stopping。报告步数的选择20 个非挑拣non-cherry-picked种子在训练末尾的验证损失均值与显著性判定如下step 3150: mean3.28039, (3.28-mean)·√20 -0.00173 FAIL (mean above 3.28) step 3175: mean3.27885, (3.28-mean)·√20 0.00517 PASS ← reported step 3200: mean3.27811, (3.28-mean)·√20 0.00843 PASS (final, comfortable margin)即 3150 步时均值仍在 3.28 之上未通过3175 步首次越过门槛并满足条件(3.28 - mu) · √n ≥ 0.0043200 步训练终点则以更宽裕的余量通过。README 还给出了参考性的单侧 z 检验在 σ0.0013 下3175 步的z 0.886, p 0.188即在此 n 下要到 3225 步才能达到p 0.001——但判定基准是 README 指定的精度条件而非 z 检验本身。20 个种子的完整数据| seed | step 3150 | step 3175 | step 3200 | | -: | -: | -: | -: | | 1 | 3.27943 | 3.27786 | 3.27712 | | 2 | 3.28162 | 3.28008 | 3.27935 | | 3 | 3.28003 | 3.27844 | 3.27771 | | 4 | 3.28084 | 3.27927 | 3.27853 | | 5 | 3.28178 | 3.28024 | 3.27951 | | 6 | 3.28089 | 3.27936 | 3.27862 | | 7 | 3.28077 | 3.27928 | 3.27855 | | 8 | 3.28085 | 3.27934 | 3.27861 | | 9 | 3.27929 | 3.27777 | 3.27704 | | 10 | 3.28069 | 3.27911 | 3.27837 | | 11 | 3.28018 | 3.27863 | 3.27791 | | 12 | 3.27854 | 3.27701 | 3.27628 | | 13 | 3.27826 | 3.27670 | 3.27601 | | 14 | 3.27892 | 3.27741 | 3.27668 | | 15 | 3.28116 | 3.27959 | 3.27886 | | 16 | 3.28064 | 3.27911 | 3.27838 | | 17 | 3.27979 | 3.27828 | 3.27751 | | 18 | 3.28090 | 3.27935 | 3.27863 | | 19 | 3.28274 | 3.28119 | 3.28046 | | 20 | 3.28041 | 3.27888 | 3.27815 | |mean|3.28039|3.27885|3.27811| |std|0.00111|0.00111|0.00111|全部 20 个种子的原始逐 step 验证损失记录可查看 完整运行日志单次运行约 503 秒 / 3200 步8 卡 H100每步约 157 ms。与历史记录 #10、#11 的对比README 的对比图中同时绘制了 4 条曲线即上文两张配图#10: NorMuonlr0.035 wd0.0253250 步达到 3.2789n20#11: ContraNorMuon with update clamp-min3225 步达到 3.2785n16pure Aurora不使用 Contra 与 u/w-floor 的纯 Aurora 对照3200 步达到 3.2790n20本次提交: Aurora Contra-Muon u/w-floor3175 步达到 3.2789n20。对比结论均以各曲线自身报告步数为准本次 3175、#11 3225、#10 3250在所有存在多方案数据的共同后端步数上本次提交的均值都是最低的——比 #11 提前 50 步、比 #10 提前 75 步首次越过通过线。图中灰色虚线标注目标target3.28zoom 图则将横轴聚焦到 3000~3300 步、纵轴聚焦到 3.275~3.310用于精确观察收敛阶段的差异。值得注意的额外信息本次的纯 Aurora 对照无 Contra、无 u/w-floor也在 3200 步达到 3.2790n20说明 Aurora 本身单独使用就已具备与当时记录相当的水准而叠加 Contra-Muon 与 u/w-floor 后进一步把步数压到 3175。提交者为 Aurora 原作者PR #284见 benchmark 历史表。复现与扩展建议如需复现或在此基础上迭代可参考 benchmark Quickstartgit clone https://github.com/KellerJordan/modded-nanogpt.git cd modded-nanogpt pip install torch2.11 huggingface_hub python data/cached_fineweb10B.py 20 # 2B tokens足够 4000 步 torchrun --standalone --nproc_per_node$(nvidia-smi -L | wc -l) records/track_3_optimization/train_gpt_simple.py注意该命令运行的是基准基线脚本本次提交的完整可复现代码内嵌在 运行日志 中分隔线之前的部分按 benchmark 规则 的做法把日志中之前的 Python 代码提取为train_gpt_simple.py再执行 Quickstart 即可复现含 1/2/4/8 卡任意组合脚本自动适配 world_size。基准架构为 12 层、768 维的 GPT-2 规模模型vocab 50304上下文 1024batch size 8 x 64K tokens每步只允许一次前向-反向传播。从 Track 3 历史 还可以看到 Aurora 思路的后续演化result #30 把 Aurora 行均衡 polar 应用到宽矩阵mlp.proj上并进一步压缩到 2930 步说明行均衡 polar这一机制在后续记录链中持续扮演重要角色。对研究型读者建议在调试 Aurora 时重点观察两个超参的影响pp_iterations1 为 Muon EQ 基线、2 为推荐值与pp_beta阻尼指数越大收敛到均匀行范数越快但可能引入振荡。小结本次 Aurora 提交展示了一条清晰的优化器研究路径在不改变架构、数据集、batch size 与每步前向-反向次数的前提下仅靠优化算法层面的三项叠加Aurora 行均衡 polar Contra-Muon 扰动 u/w-floor 地板就把达到 3.28 验证损失的步数从 #10 的 3250 步降到 3175 步。其中 Aurora 的贡献在于通过polar(D · G)的对角均衡机制消除标准 polar 继承的 leverage 分数不均匀性使隐藏矩阵的更新在行维度上获得一致的力度配合pp_iterations2, pp_beta0.5可在两次 Newton-Schulz 迭代内把行范数偏差压到约 2% 以内。所有超参、源码与 20 个种子的完整日志均可在 提交目录 中核验。赞分享人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐优化器分组粒度下钻modded-nanogpt 将 K/V 按 Paired-Head 拆分为独立 Muon 组以更少步数守住 3.28 验证损失优化器分组粒度下钻modded nanogpt 将 K/V 按 Paired Head 拆分为独立 Muon 组以更少步数守住 3.28 验证损失 本篇技术人工智能大模型预训练分布式训练模型优化深度学习Modded-NanoGPT 中的 SOAP 优化器实战3.15B Tokens 达成 3.28 验证损失的样本效率记录Modded NanoGPT 中的 SOAP 优化器实战3.15B Tokens 达成 3.28 验证损失的样本效率记录 本文基于 Modded NanoGP人工智能大模型预训练分布式训练模型优化深度学习Podman 镜像签名中的 --sign-passphrase-file 选项passphrase 文件读取机制与实战详解Podman 镜像签名中的 sign passphrase file 选项passphrase 文件读取机制与实战详解 本文基于 Podman 仓库中 sig人工智能大模型预训练分布式训练模型优化深度学习上一篇Nix 实验特性Experimental Features机制全解析功能开关、生命周期与源码实现下一篇探索免费大语言模型API的完整路径从零开始构建AI应用生态创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表