企业官网建设流程全解析

• 人工智能
• 大模型
• 预训练
• 分布式训练
• 模型优化
• 深度学习

导读

本文围绕 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 预算 50B（num_iterations: 4578）。

README 中的 Changelog 只列了三行，却浓缩了本次纪录的全部架构变化：

• untied embed and head weights：解绑词嵌入与输出头权重；
• added RMSNorm after embed：embedding 之后追加 RMSNorm；
• init head to zero：输出头初始化为零。

纪录推进由 @Grad62304977 主导，作者负责超参微调与簿记；其中"head 零初始化"的灵感来自 @cloneofsimo 的 scaling guide，README 明确指出它"贡献了纪录中相当大的一部分"（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。

但关键限定在下一句：它不改变活跃参数的数量，也不改变推理吞吐。原因是：

1. 解绑只在训练侧生效，正向计算仍是"embedding 取词 → 12 层 Transformer → lm_head 投影到 logits"两条路径，前向 FLOPs 与原始结构一致；
2. 推理时并不需要两份权重并行存在，吞吐不受影响；
3. 因此后续纪录仍约束在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_s=100/448, w_s=2.0/448, grad_s=(0.75 / 8) / 448) nn.init.normal_(self.lm_head.weight, mean=0, std=0.005)

即"整体零初始化"被细化为std=0.005的极小方差正态初始化（本质仍是近零初始化），并配套 fp8 量化所需的静态 scale（x_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 的输入尺度不稳定；在嵌入后立即归一化，等于为整个残差流提供一个尺度受控的起点，配合后续逐子层的 RMSNorm（Block.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 用 Adam，Transformer 主体用 Muon。11/03 纪录的训练代码正是这一分工的典型形态：

optimizer1 = torch.optim.Adam([raw_model.transformer.wte.weight], lr=0.3, betas=(0.9, 0.95), fused=True) optimizer2 = torch.optim.Adam([raw_model.lm_head.weight], lr=0.002, betas=(0.9, 0.95), fused=True) optimizer3 = Muon(raw_model.transformer.h.parameters(), lr=0.02, momentum=0.95)

三个要点：

1. embed 用大学习率（0.3）：embedding 查表是稀疏操作、参数更新相对独立，高 lr 允许快速调整 token 表征；
2. lm_head 用极小学习率（0.002）：配合零初始化，输出头只需微调即可在 softmax 上形成正确分布，大幅压缩早期探索成本；
3. Muon 只接管 2D 矩阵主体：Muon 的文档明确警告不要用于 embedding、最终全连接层及 0/1 维参数（Muon(raw_model.transformer.h.parameters(), lr=0.02, momentum=0.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 80GB，PyTorch 2.5.1+cu124）与逐步训练记录。值得关注的工程细节：

• 时序纪律：代码从 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_head（self.lm_head(x[:, [-1], :])），配合"解绑不改变推理吞吐"的论断；
• 词表扩展：50257 个真实 GPT-2 token 扩展到 50304（128 的倍数）以适配高效矩阵形状，这一约定延续至今（gpt.py中next_multiple_of_n(vocab_size, n=128)）。

后续演进：三项技巧的归宿

11/03 纪录并非终点，它留下的三个技巧在后来的纪录中分别演化：

1. untie 从"一次性开关"变成"阶段表调度"：如前述SPLIT_EMBED_STAGE与 anvil.py 的状态迁移机制，解绑被精确安排进训练阶段表（record #360 体系），并衍生出后续对 embed 独立初始化的精细调优；
2. RMSNorm after embed 升级为多路注入锚点：当前 forward 中归一化后的x0同时服务 X0 注入、bigram 注入与 MUDD 门控，成为残差重组的基准信号；
3. 零初始化 head 演化为近零初始化 + fp8 量化：CastedLinearT保留零初始化传统，同时以std=0.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.py（SPLIT_EMBED_STAGE）、gpt.py（embed/lm_head 构造与 forward 入口）、layers.py（norm 与零初始化）与 anvil.py（绑定/解绑的优化器状态迁移）。

• 人工智能
• 大模型
• 预训练
• 分布式训练
• 模型优化
• 深度学习