尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

modded-nanogpt 速度纪录解读:Untie Embed 解绑词嵌入、Embed 后 RMSNorm 与零初始化输出头的实战剖析

modded-nanogpt 速度纪录解读:Untie Embed 解绑词嵌入、Embed 后 RMSNorm 与零初始化输出头的实战剖析 人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载导读本文围绕 modded-nanogpt 在 2024-11-03 创下的 NanoGPT 训练速度纪录8×H100 上 10.8 分钟达到 FineWeb 3.28 验证损失展开逐项拆解支撑该纪录的三处架构改动——解绑 token embedding 与 lm_head 权重untied embeddings、embedding 之后追加 RMSNorm、lm_head 零初始化。你将了解到这三项改动各自的动机、参数预算与推理吞吐的影响、它们如何与 Muon 优化器体系配合以及这些技巧在当前仓库源码中的实现形态与后续演进。纪录背景10.8 分钟从何而来records/track_1_short/2024-11-03_UntieEmbed/README.md记录了这一时间点新的 NanoGPT 训练速度纪录为3.28 FineWeb val loss / 10.8 分钟8×H100将此前 12.0 分钟的纪录进一步压缩。同目录的训练日志d6b50d71-f419-4d26-bb39-a60d55ae7a04.txt给出了完整证据链最终一步step:4578/4578 val_loss:3.2762 train_time:648063ms即 648 秒 ≈ 10.8 分钟全程平均每步约 141.87ms训练共 4578 步、token 预算 50Bnum_iterations: 4578。README 中的 Changelog 只列了三行却浓缩了本次纪录的全部架构变化untied embed and head weights解绑词嵌入与输出头权重added RMSNorm after embedembedding 之后追加 RMSNorminit head to zero输出头初始化为零。纪录推进由 Grad62304977 主导作者负责超参微调与簿记其中head 零初始化的灵感来自 cloneofsimo 的 scaling guideREADME 明确指出它贡献了纪录中相当大的一部分a significant fraction of the record。解绑词嵌入为什么这是any%纪录README 用一句话给这次纪录定了性这在技术上是某种 any% 纪录——解绑 embedding 和 lm_head 使参数总量增加了 39M。GPT-2 词表 50257项目扩展至 50304嵌入矩阵与输出矩阵各为vocab × n_embd绑定时只有一份权重解绑后变成两份因此多出50304 × 768 ≈ 38.6M参数接近 39M。但关键限定在下一句它不改变活跃参数的数量也不改变推理吞吐。原因是解绑只在训练侧生效正向计算仍是embedding 取词 → 12 层 Transformer → lm_head 投影到 logits两条路径前向 FLOPs 与原始结构一致推理时并不需要两份权重并行存在吞吐不受影响因此后续纪录仍约束在124M 活跃参数即原始 NanoGPT 规模这一口径下比较本次纪录只是阶段性展示了参数预算放宽能带来多少收益。从当前仓库的实现可以印证这一思路的延续与收束。track_1_short/config.py中SPLIT_EMBED_STAGE 4明确标注embed unties from lm_head at the start of this stage, the extension stage (record #360)——即解绑动作被安排进训练阶段表TRAINING_STAGES 的第 4 阶段extension stage由训练流程在特定阶段解锁而不是一开始就解绑。这说明社区最终把解绑从一次性纪录技巧演化为一个可由阶段表控制、在训练中后期生效的调度机制。三项改动的源码级实现1. Head 零初始化当前仓库的继承形态原始纪录中lm_head通过self.lm_head.weight.data.zero_()完成零初始化。当前仓库track_1_short/model/layers.py的CastedLinearT.reset_parameters()保留了这一约定def reset_parameters(self) - None: with torch.no_grad(): nn.init.zeros_(self.weight) # Grad62304977 and others注释中的Grad62304977 and others正是对 11/03 纪录的致谢链。而在track_1_short/model/gpt.py中lm_head 的构造已演进为self.lm_head CastedLinearT(model_dim, self.vocab_size, x_s100/448, w_s2.0/448, grad_s(0.75 / 8) / 448) nn.init.normal_(self.lm_head.weight, mean0, std0.005)即整体零初始化被细化为std0.005的极小方差正态初始化本质仍是近零初始化并配套 fp8 量化所需的静态 scalex_s/w_s/grad_s。其动机保持一致让语言模型在训练初期近似从恒等/均匀分布起步避免随机大权重在 softmax 上制造强烈偏好从而把早期梯度集中在学习结构本身这与 cloneofsimo scaling guide 的建议一脉相承。值得注意零/近零初始化并非只用于 lm_head。同一时代纪录records/track_1_short/2024-10-14_ModernArch/dabaaddd-237c-4ec9-939d-6608a9ed5e27.txt中注意力输出投影与 MLP 下投影同样采用c_proj.weight.data.zero_()注释同为zero init suggested by Grad62304977可见零初始化投影层是当时整套速度纪录的通用组件lm_head 只是其延伸。2. Embed 后 RMSNorm让残差流从规整起点出发README 的第二项改动是added RMSNorm after embed。原始记录中的实现为x self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) x F.rms_norm(x, (x.size(-1),))即在 embedding 查表后立即对嵌入向量做一次 RMSNorm。其作用可以从残差流视角理解token embedding 的幅度随训练漂移若直接进入残差流首层注意力/MLP 的输入尺度不稳定在嵌入后立即归一化等于为整个残差流提供一个尺度受控的起点配合后续逐子层的 RMSNormBlock.forward中每个残差分支前都先F.rms_norm保证各层输入统计稳定。当前仓库track_1_short/model/gpt.py的 forward 中这一操作被保留并融入更复杂的入口处理x self.embed(input_seq) # embed is synced from lm_head during tied phase by optimizer ... x x0 norm(x[None])这里的norm见track_1_short/model/layers.py正是F.rms_norm(x, (x.size(-1),))的封装。值得注意的是当前版本还叠加了 smear前向 token 位置混合与 n-gram/bigram 注入等后验技巧但embed 之后先归一化再进残差流这一来自 11/03 纪录的基本盘没有改变且x0归一化后的嵌入还被后续 X0_INJECT_LAYERS 与 MUDD 门控反复引用成为多层注入的锚点。3. Untie从绑定时同步到阶段化解绑11/03 纪录在代码层面只是不再让 embed 与 lm_head 共享权重但当前仓库给出了一个更精细的工程化版本训练前期保持 tied绑定在特定训练阶段才解绑。绑定阶段track_1_short/model/gpt.py构造时self.embed.weight.copy_(self.lm_head.weight.T)将两者初始化为一致forward 注释写明 embed is synced from lm_head during tied phase by optimizer。优化器协同track_1_short/optim/anvil.py是这套机制的核心。优化器在绑定阶段只对 lm_head 做 Adam 更新与通信Comms and update math are only done on lm_head通过transpose_add(embed_param.grad, param.grad)把 embed 的梯度转置累加进 lm_head 梯度再在 gather 后transpose_copy(lm_param.data, embed_param.data)把 lm_head 权重同步回 embed——即用一次 lm_head 的 Adam 状态同时维护两份权重。解绑点anvil.py提供copy_lm_state_to_embed()注释明确其职责是 Copy the optimizer state from the lm_head to the embed at the untie point并通过 all-gather reshard 处理 lm_head按 model_dim 切分与 embed按 vocab 切分两种不同切分方式的 Adam 状态迁移解绑后 embed 获得独立优化。这种设计平衡了收益与成本绑定阶段 embed 不占额外优化器状态、不增加通信解绑后 embed 独立演进为训练后期提供更灵活的表征空间。参数与优化器分工embed/head 始终交给 Adam解绑后出现了三类参数而优化器分工遵循一个稳定的社区结论records/track_1_short/2024-10-29_Optimizers/README.md已确立embedding 与 lm_head 用 AdamTransformer 主体用 Muon。11/03 纪录的训练代码正是这一分工的典型形态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)三个要点embed 用大学习率0.3embedding 查表是稀疏操作、参数更新相对独立高 lr 允许快速调整 token 表征lm_head 用极小学习率0.002配合零初始化输出头只需微调即可在 softmax 上形成正确分布大幅压缩早期探索成本Muon 只接管 2D 矩阵主体Muon 的文档明确警告不要用于 embedding、最终全连接层及 0/1 维参数Muon(raw_model.transformer.h.parameters(), lr0.02, momentum0.95)因为这些矩阵的行列几何不满足更新近似正交矩阵的假设。解绑之后这一分工自然扩展为两份独立 Adam 状态 一个 Muon 主体与records/track_1_short/2024-11-04_50Bruns/README.md中head and embedding are always optimized by Adam的描述完全吻合说明该分工在更长时程训练50B token中同样适用。从纪录到可复现训练日志能告诉我们什么d6b50d71-f419-4d26-bb39-a60d55ae7a04.txt是一份完整可复现的训练日志头部内嵌全部训练代码随后是nvidia-smi快照8×H100 80GBPyTorch 2.5.1cu124与逐步训练记录。值得关注的工程细节时序纪律代码从 step 10 才开始计时if step 10: training_time_ms 0规避启动与编译导致的慢步验证val_loss也计入墙钟时间确保3.28 / 10.8 分钟口径诚实关键中间检查点step 125 val_loss 4.8310、step 375 3.9903、step 2500 附近 3.31最终 3.2762——衰减曲线平滑未出现发散推理侧微优化inference 时只用最后位置过 lm_headself.lm_head(x[:, [-1], :])配合解绑不改变推理吞吐的论断词表扩展50257 个真实 GPT-2 token 扩展到 50304128 的倍数以适配高效矩阵形状这一约定延续至今gpt.py中next_multiple_of_n(vocab_size, n128)。后续演进三项技巧的归宿11/03 纪录并非终点它留下的三个技巧在后来的纪录中分别演化untie 从一次性开关变成阶段表调度如前述SPLIT_EMBED_STAGE与 anvil.py 的状态迁移机制解绑被精确安排进训练阶段表record #360 体系并衍生出后续对 embed 独立初始化的精细调优RMSNorm after embed 升级为多路注入锚点当前 forward 中归一化后的x0同时服务 X0 注入、bigram 注入与 MUDD 门控成为残差重组的基准信号零初始化 head 演化为近零初始化 fp8 量化CastedLinearT保留零初始化传统同时以std0.005与 fp8 scale 支撑更高吞吐的训练路径训练侧 loss 走 fp8验证侧走 bf16见gpt.py的_loss分支。如果你想复现 11/03 的纪录仓库records/track_1_short/2024-11-03_UntieEmbed/d6b50d71-f419-4d26-bb39-a60d55ae7a04.txt内的完整代码可直接在 8×H100 环境运行若要研究当前形态的解绑实现可依次阅读 config.pySPLIT_EMBED_STAGE、gpt.pyembed/lm_head 构造与 forward 入口、layers.pynorm 与零初始化与 anvil.py绑定/解绑的优化器状态迁移。赞分享人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐modded-nanogpt 新纪录解析Value Embed 门控与 Skip Connection 门控-35 steps-1.3smodded nanogpt 新纪录解析Value Embed 门控与 Skip Connection 门控 35 steps 1.3s 导读 本文围绕人工智能大模型预训练分布式训练模型优化深度学习重新绑定 LM Head 与 Embed 权重modded-nanogpt 的 FP8 尺度重调与步数缩减实战重新绑定 LM Head 与 Embed 权重modded nanogpt 的 FP8 尺度重调与步数缩减实战 导读 本文基于 modded nanogpt人工智能大模型预训练分布式训练模型优化深度学习modded-nanogpt 134.9 秒新纪录Refine Skip 残差连接架构精修与 block lambda 初始化剖析modded nanogpt 134.9 秒新纪录Refine Skip 残差连接架构精修与 block lambda 初始化剖析 本技术记录基于 modde人工智能大模型预训练分布式训练模型优化深度学习上一篇DeepSeek Harness 跨 workspace 会话恢复让 /resume 回到任意项目目录的设计剖析下一篇Vitess v15.0.0 版本深度解析Flags 体系重构、Online DDL 并发与 VDiff2 功能完备创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表