modded-nanogpt 50B-token 长时程训练实验解析：Muon vs Adam 与数据复用策略的实证对比

人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载导读本文基于 modded-nanogpt 仓库中 2024-11-04_50Bruns 实验记录 展开完整还原该项目把 11/03 的 NanoGPT(124M) 速跑记录从短时程约 90 秒 / 数亿 token扩展到 50B FineWeb token 长时程训练的全过程。你将理解两个核心问题的实证答案Muon 优化器相对 Adam 在长时间训练下是否依然保持优势以及10B token 重复 5 个 epoch 与 50B 全新 token 一次性训练 两种数据预算方案为何最终效果几乎相同。文中给出的四组完整运行日志、关键超参配置与源码级原理可直接复用于你自己的长时程 NanoGPT 规模实验设计与优化器选型。实验背景为什么要把速跑记录扩展到 50B tokensmodded-nanogpt 的 track_1_short 赛道专注于在极短预算内训练出高质量 GPT-2(124M)其成名作是 NanoGPT (124M) in 90 seconds 的速跑记录。而2024-11-04_50Bruns这一组实验则完全不同它的目的是检验速跑配方在长时间尺度上是否依然成立尤其是验证被速跑证明有效的 Muon 优化器能否在 50B token 的大预算下继续保持相对于 Adam 的优势。具体来说这次实验回答两个问题优化器对比Transformer 主体body分别使用 Muon 与 Adam而 headlm_head与 embeddingwte在两组中始终由 Adam 优化——这是 Muon 官方使用约束所要求的详见下文源码分析。数据预算对比同样消耗 50B token 的算力是在 10B 数据上循环 5 个 epoch更好还是一次性消费 50B 全新数据更好实验结果出乎意料两者性能几乎相同。该目录下共四组运行日志运行数据方案主体优化器HellaSwagMuon 50B tokens50B 全新 FineWebMuon35.82Adam 50B tokens50B 全新 FineWebAdam34.26Muon 5×10B tokens10B 数据 × 5 epochMuon36.17Adam 5×10B tokens10B 数据 × 5 epochAdam34.05一个关键观察Muon 在两种数据方案下都稳定领先 Adam 约 1.5 个 HellaSwag 点35.82 vs 34.26、36.17 vs 34.05且差距方向一致说明这个差距不太可能由学习率调度等细节偏差造成而数据方案本身循环 vs 一次性带来的差异则小得多Muon 组 0.35 点、Adam 组 0.21 点。README 中作者也谨慎地指出速跑调优的 WSD 学习率调度与零权重衰减在更长训练中可能已非最优但 Muon/Adam 的差距大到不太可能被这类因素抹平。HellaSwag 分数基准125M 规模模型的可比坐标系为了判断 35.82 这个分数好不好README 给出了同一量级模型的公开基准全部为 124M135M 规模基准训练数据HellaSwagKarpathy 的 llm.c 基线10B FineWeb tokens29.9OpenAI GPT-2 (124M)原始 GPT-2 数据29.4OpenAI GPT-3 (124M)300B WebText tokens33.7HuggingFace SmolLM2-135M2T FineWeb/DCLM 等 tokens42.1可以看到50B token 的 Muon 运行35.82已经明显超过 Karpathy 的 llm.c29.9与 OpenAI GPT-2/GPT-3 的 124M 基线29.4 / 33.7但仍显著低于在 2T token约为 40 倍数据量上训练的 SmolLM2-135M42.1。这为读者提供了一个非常直观的定位在 50B token 预算下Muon 配方的 124M 模型已经进入略高于 GPT-3 同等规模的区间而想要追平 SmolLM2 则需要把数据预算再扩大两个数量级。四组运行日志的完整实测数据运行环境四组一致所有运行日志头部都记录了完整环境信息此处以 Muon 50B 运行为例见 530f3ee1 日志 第 542-598 行软件PyTorch 2.5.1cu124编译目标 CUDA 12.4驱动 NVIDIA-SMI 555.42.06CUDA 12.5硬件8× NVIDIA H100 80GB HBM3单卡显存占用约 5.15.4GiB训练时 GPU 利用率在个位数到 18% 之间波动属于短样本间隔下采样的正常表现数据FineWeb 训练集.binshardtrainFineWeb 验证集.binshardval训练曲线与吞吐从日志的前 1000 行可以提取出完整的起步阶段轨迹。以 Muon 50B 运行为例step 0 验证损失 10.8258train_loss 同步从 10.8258 一路下降约 60 步后跌破 5.5约 200 步后进入 4.2 附近前 10 步因编译/预热不计时从 step 11 起计时稳态每步约 139141ms含 8 卡全局 batch 512、序列长 1024每 125 步评估一次验证损失val_loss_every125验证 token 数固定为 10,485,760val_tokens10485760保证组间可比。四组运行的关键结果汇总取自各日志末尾运行最终 train_loss最终 val_loss总时长稳态 step_avgMuon 50B3.11213.050813,525,160 ms约 3.76 小时141.84 msAdam 50B3.14483.078813,262,037 ms约 3.68 小时139.08 msMuon 5×10B3.01503.050013,610,236 ms约 3.78 小时142.73 msAdam 5×10B3.04083.078013,354,902 ms约 3.71 小时140.05 ms几点值得注意Muon 与 Adam 吞吐几乎持平Muon 因额外执行 Newton-Schulz 正交化迭代稳态每步仅比 Adam 慢约 2.52.7ms约 1.8%在 8×H100 上完全可接受数据方案对最终 val_loss 影响极小同一优化器下50B 一次性 vs 10B×5 epoch 的最终 val_loss 差距在 0.03 以内Muon 3.0508 vs 3.0500Adam 3.0788 vs 3.0780与 README 的Surprisingly this does about the same完全吻合优化器差距稳定无论数据方案如何Muon 的最终 val_loss 都低于 Adam 约 0.028对应的 HellaSwag 差距约 1.5 点。超参配置解析四组实验怎么搭出来的四组日志的头部均嵌入了完整训练脚本项目约定每个 run 日志把自身代码全文写入日志便于事后精确复现。核心超参如下dataclass class Hyperparameters: # data hyperparams input_bin : str data/fineweb100B/fineweb_train_*.bin # 50B 组5×10B 组改为 data/fineweb10B/fineweb_train_*.bin input_val_bin : str data/fineweb100B/fineweb_val_*.bin # 验证集两组都用 100B 版本 # optimization hyperparams batch_size : int 8*64 # 全局 batch序列数8 卡 × 64 device_batch_size : int 64 # 单卡 batch序列数 sequence_length : int 1024 # 序列长度 num_iterations : int 95367 # 总迭代数对应 50B token 预算 warmup_iters : int 0 # Muon 组Adam 组为 500 warmdown_iters : int 27247 # WSD 线性退火步数 weight_decay : float 0 # 零权重衰减 val_loss_every : int 125 # 每 125 步评估一次验证损失 val_tokens : int 10485760 # 验证 token 数组间固定以保证可比 save_every : int 0 # 仅末尾保存检查点关于迭代数的计算全局 batch 为 8×64 序列 × 1024 token 524,288 tokens/step因此 50B token 恰好对应 95367 步。5×10B 组的num_iterations同样是 95367区别仅在于input_bin指向 10B 数据集fineweb10B/fineweb_train_*.bin使 10B 数据被完整循环 5 遍——这正是5 epochs of 10B tokens的实现方式。优化器三件套三组共用前两个仅第三个不同# Muon 组 optimizer1 torch.optim.Adam([raw_model.transformer.wte.weight], lr0.3, betas(0.9, 0.95), fusedTrue) optimizer2 torch.optim.Adam([raw_model.lm_head.weight], lr0.002, betas(0.9, 0.95), fusedTrue) optimizer3 Muon(raw_model.transformer.h.parameters(), lr0.02, momentum0.95) # Adam 组仅第 3 行不同 # optimizer3 Muon(raw_model.transformer.h.parameters(), lr0.02, momentum0.95) optimizer3 torch.optim.Adam(raw_model.transformer.h.parameters(), lr0.0018, betas(0.9, 0.95), fusedTrue)关键参数解读wteembeddinglr0.3、lm_head lr0.002token embedding 与输出投影使用完全不同的学习率这是速跑调优沉淀下来的惯例Muon lr0.02, momentum0.95Nesterov 默认开启Adam 对照组的主体 lr0.0018WSD 调度warmup_iters为 0Muon 组或 500Adam 组随后恒定最后 27247 步约 28.6% 的训练线性退火到 0调度函数为def get_lr(it): assert it args.num_iterations if it args.warmup_iters: return (it1) / args.warmup_iters elif it args.num_iterations - args.warmdown_iters: return 1.0 else: decay_ratio (args.num_iterations - it) / args.warmdown_iters return decay_ratio零权重衰减weight_decay 0README 中作者对此表达了保留态度认为 WSD 调度与零 weight decay 都是为速跑时程调优的在 50B 长时程上可能欠调优。模型结构沿用速跑配方GPTConfig(vocab_size50304, n_layer12, n_head6, n_embd768)其中 50304 是 GPT-2 的 50257 词表向上取整到 128 的倍数效率考虑模型经torch.compile编译后包进 DDP并在 PyTorch 2.5.1 下强制启用 cuDNN SDPenable_cudnn_sdp(True)、关闭 Flash/Mem-efficient/Math 后端因为日志注释明确 cuDNN attention 比 Flash 快约 4ms。从源码看 Muon 的分布式正交化实现这一节深挖日志内嵌脚本中的 Muon 实现——这也是本实验最核心的技术点且与当前仓库主线如 optim/anvil.py 中的优化器模块一脉相承。MuonMomentUm Orthogonalized by Newton-schulz的核心思想内部先跑标准 SGD-momentum再用正交化后处理把每个二维参数矩阵的更新替换为最近的近正交矩阵。其文档字符串明确给出使用边界只能处理 2D 参数不得用于 embedding、最终全连接层及任何 0/1 维参数——这些应交给 AdamW 等标准优化器本实验正是如此wte、lm_head走 Adam仅transformer.h走 Muon作者评估小 batch 训练下可能效果不佳、微调预训练模型可能不佳、且尚未在超过 NanoGPT(124M) 的规模上验证。正交化的计算核心是 Newton-Schulz 迭代五阶格式系数(a, b, c) (3.4445, -4.7750, 2.0315)通过把矩阵范数归一化到最大奇异值 ≤ 1 后迭代 510 步近似求解零次幂即正交化。关键技巧是整个迭代可以在 bfloat16 下稳定运行配合torch.compile编译成 GPU 内核而对比后端 SVDzeropower_via_svd虽更精确但速度慢且不易低精度化因此推荐newtonschulz5后端def zeropower_via_newtonschulz5(G, steps10, eps1e-7): a, b, c (3.4445, -4.7750, 2.0315) X G.bfloat16() X / (X.norm() eps) # 保证最大奇异值 1 if G.size(0) G.size(1): X X.T # 转置使迭代在较短维度上进行 for _ in range(steps): A X X.T B A X X a * X b * B c * A B if G.size(0) G.size(1): X X.T return XMuon.step()的分布式流程值得单独拆解这也是 8×H100 上吞吐只比 Adam 慢 ~2ms 的原因按 rank 分片计算更新对第 i 个参数仅当i % WORLD_SIZE RANK时本进程负责其正交化恰好让 12 层 transformer 均匀分布到 8 卡每卡 ~1.5 层随后 flatten 写入updates_flat缓冲一次 all-reduce 同步整个参数组的更新拼成一个 bfloat16 张量后做单次dist.all_reduce(SUM)用求和替代平均——因为每个参数只被一个 rank 计算过SUM 天然等于该值本身比逐参数多轮通信高效得多统一施加更新p.data.add_(g, alpha-lr)动量缓冲buf.mul_(momentum).add_(g)加 Nesterov 修正最后乘以max(1, rows/cols)^0.5的比例因子补偿矩阵非方带来的尺度差异。这套分片计算 单次全量 all-reduce的模式正是 modded-nanogpt 后续 FasterReduce、EvenFasterReduce 等记录持续优化的对象读者可顺藤摸瓜了解该项目的通信优化演进。训练循环的工程细节日志内嵌脚本摘要四组日志共享同一套训练主循环其中若干细节直接关系到长时程训练的稳定性与可复现性计时策略前 10 步不计时编译与缓存预热step 10 重置training_time_ms与t0之后每步输出step_avg验证流程每 125 步及最后一步切到model.eval()在固定 10,485,760 个验证 token 上累计 loss经dist.all_reduce(AVG)得到全卡平均后写入日志注意验证时依然使用 bfloat16 autocast且刻意不套no_grad()——注释说明这会在torch.compile下触发报错梯度累积train_accumulation_steps batch_size // (B * world_size) 1本配置下 8 卡每卡 64 序列恰好等于全局 512 序列因此循环体退化为单次前向/反向代码保留了model.no_sync()分支以支持未来增大累积步数的场景检查点与日志master processrank 0把完整训练脚本源码、nvidia-smi输出、PyTorch/CUDA 版本写入logs/run_id.txt并在save_every指定的步数保存state_step%06d.pt本次实验为 0即仅末尾保存学习率调度同步三个优化器共享同一个get_lr的 LambdaLR 调度器保证 Muon 与两个 Adam 的退火节奏一致。实验结论与设计可复用的要点综合 README、四组日志与源码可以给出以下可被后续实验直接引用的结论与注意事项Muon 在 50B token 长时程上稳定优于 Adam两种数据方案下 HellaSwag 均领先约 1.5 点最终 val_loss 低约 0.028而吞吐代价仅约 2%~140ms vs ~142ms 每步10B×5 epoch 与 50B 一次性几乎等价这与重复数据有害的直觉相悖说明在 10B 规模的 FineWeb 数据上循环 5 遍造成的过拟合影响在此模型规模下可忽略也意味着小数据集反复利用在算力受限时是一个可行的近似方案调度与正则的边际风险WSD 退火时长27247 步 ≈ 28.6%与零 weight decay 是为速跑时程调优的长时程下可能欠优——若进一步延长训练作者暗示值得重新调优这两项可复现的工程底座日志自含完整脚本与环境快照任意读者都可以用相同超参、8×H100 与 FineWeb.bin数据data/cached_fineweb10B.py、data/cached_fineweb100B.py 提供构建脚本复现全部四组运行。如果你想在自己的长时程 NanoGPT(124M) 训练中复用这份配方最稳妥的路径是主体用Muon(lr0.02, momentum0.95)newtonschulz5后端wte/lm_head用 Adamlr 0.3 / 0.002全局 batch 保持 512 序列 × 1024 tokenWSD 退火占比 ~28.6%weight decay 保持 0并优先用data/fineweb100B这样的单遍数据若数据预算不足可参照本实验放心地循环使用 10B 数据 5 遍。赞分享人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐Modded-NanoGPT 的 AdamW 基线训练从 llm.c 日志格式到 5B token 的优化器对照实验Modded NanoGPT 的 AdamW 基线训练从 llm.c 日志格式到 5B token 的优化器对照实验 导读 本文以 Modded NanoGP人工智能大模型预训练分布式训练模型优化深度学习Modded-NanoGPT BatchSize 调优实战训练序列 48K 与验证序列 256K 的权衡Modded NanoGPT BatchSize 调优实战训练序列 48K 与验证序列 256K 的权衡 本文以 Modded NanoGPTNanoGPT人工智能大模型预训练分布式训练模型优化深度学习Modded-NanoGPT训练日志如何解析与可视化性能数据Modded NanoGPT训练日志如何解析与可视化性能数据 训练日志是优化深度学习模型的关键依据Modded NanoGPT项目通过结构化记录和可视化工具人工智能大模型预训练分布式训练模型优化深度学习上一篇5分钟掌握Node.js终端交互神器ora的终极使用指南下一篇3步把AI模型塞进Android AppPaddle-Lite部署实操指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考