⚠️ 在开始阅读之前，如果你对实时 Agent / 数字人 / 多模态系统 / LiveKit 架构感兴趣，
欢迎先到 GitHub 给项目点一个 ⭐ Star，这是对开源作者最大的支持。

🚀AlphaAvatar 项目地址（强烈建议先收藏，该项目正在持续更新维护）：
👉 https://github.com/AlphaAvatar/AlphaAvatar
🚀AIPapers 项目地址（具有更全的有关LLM/Agent/Speech/Visual/Omni论文分类）：
👉 https://github.com/AlphaAvatar/AIPaperNotes

摘要

优化器的选择对大语言模型（LLM）的训练效率和计算成本有着显著的影响。近年来，Muon 优化器通过正交化参数更新，并利用更好的条件数来改善优化几何结构，展现出了良好的应用前景。尽管 Muon 已成为 Adam 的有力竞争者，但如何将二者的优势结合起来却尚未得到系统性的探索。本文提出了一种名为NorMuon（Neuron-wise Normalized Muon）的优化器，旨在弥补这一研究空白。 NorMuon 巧妙地将正交化与神经元级自适应学习率相结合。我们的分析表明，虽然 Muon 能够有效地降低条件数，但其更新结果却呈现出高度不均匀的神经元范数，导致某些神经元主导了优化过程。NorMuon 通过维护每个神经元的二阶动量统计量，并在正交化之后应用行级归一化来解决这一不平衡问题，从而在保持 Muon 条件数优势的同时，确保参数利用的均衡性。为了实现大规模的实际部署，我们基于 FSDP2 框架开发了一种高效的分布式实现，该实现策略性地将正交化计算分布到各个设备上。在多个模型规模上的实验表明，NorMuon 的性能始终优于 Adam 和 Muon，在 11 亿预训练数据集上，其训练效率比 Adam 提高了 21.74%，比 Muon 提高了 11.31%，同时保持了与 Muon 相当的内存占用。我们的研究结果表明，正交化和自适应学习率是互补而非竞争关系，这为大规模深度学习中的优化器设计开辟了新的途径。我们的实现已开源，地址为 https://github.com/zichongli5/NorMuon.git。

1.Introduction

训练效率仍然是扩展大语言模型（LLM）的核心挑战之一，其中优化器的选择直接影响收敛速度、计算需求，并最终影响大规模训练的可行性。社区标准优化器 Adam 通过逐坐标预处理实现稳健的性能：根据每个参数梯度历史的二阶矩动态调整其学习率。虽然这种逐坐标自适应在计算上高效且通常稳定，但它存在一个根本性的局限性——它将每个参数独立处理，忽略了神经网络层中固有的丰富的几何结构和跨坐标依赖关系。

近年来，人们尝试通过各种方法捕捉跨坐标结构，以克服这一局限。Adam-mini 利用神经网络 Hessian 矩阵近似块对角的结构，将自适应学习率应用于参数块（例如单个神经元所对应的一组参数），而不是单独应用于每个坐标。更为复杂的二阶方法，如 Shampoo 和 SOAP，则借助奇异值分解进行全矩阵预条件化，以捕捉曲率信息以及参数之间的相互依赖关系。然而，这些方法会带来较大的内存和通信开销，同时还具有较强的超参数敏感性，从而限制了其在大规模场景中的实际应用。

近期，Muon 成为一种颇具吸引力的折中方案。它通过有限步的牛顿–舒尔茨迭代，近似计算动量矩阵极分解中的正交因子。该方法得到按矩阵进行正交化的更新，在改善优化条件性的同时，仅引入适度的计算开销，而且内存消耗约为 Adam 的一半，并已在大语言模型训练中展现出良好的效果。

这些优化器在预处理的粒度和目标函数方面存在根本差异。Adam 和 Adam-mini 在不使用指数移动平均 (EMA) 的情况下，分别在每个坐标和每个神经元层面应用 L2 归一化，并在保持更新符号不变的情况下调整学习率。相比之下，理想化的 Shampoo 和 Muon 优化器在每个矩阵层面运行，主动实现参数更新的正交化。

这些优化器采用的各种预处理策略提出了一个重要问题：不同形式的预处理本质上是否相互冲突，或者它们是否可以以产生互补效益的方式结合起来？

为了探究这一现象，我们分析了不同优化器在 1.1B 参数 Transformer 模型预训练过程中更新矩阵的关键属性，考察了奇异值分布和神经元范数。如图 1a 所示，原始动量累积 SGD 的更新具有极高的条件数，表明某些方向占据主导地位，而其他参数则未得到充分利用。AdamW 产生的奇异值分布较为均衡，但改进仍然有限。相比之下，Muon 的近似正交化成功解决了这一条件数问题，在整个频谱范围内产生了均衡的奇异值。然而，考察神经元更新范数（图1b）揭示了另一种视角。与SGD 动量相比，AdamW 在降低神经元更新范数的方差方面表现出更优的性能。相反，尽管 Muon 的正交化有效地改善了矩阵层面的条件数，但神经元更新范数的方差仍然很高，一些神经元获得的更新量相对于其他神经元而言过大。

这一观察结果启发了我们的关键见解：尽管 Muon 的正交化有效地降低了更新的条件数，但神经元范数中剩余的高方差仍然会造成学习动态的不平衡，从而可能导致参数利用效率低下。借鉴 Adam-mini 在神经元自适应学习率方面的成功经验，我们提出引入二阶动量来归一化这些不同的尺度，并确保更均衡的参数更新。我们的方法NorMuon在 Muon 的正交化基础上，利用基于累积二阶统计量计算的神经元级自适应学习率进行增强。正如我们的分析所示，NorMuon 能够实现条件数低（图 1a）且神经元范数均匀（图 1b）的更新，从而结合了 Muon 和 AdamW 的优势，并实现了对网络表征能力的更均衡利用。

除了算法创新之外，基于正交化的优化器的分布式实现仍是文献中相对欠缺的研究方向。为了实现更大规模的训练，我们开发了一个与 FSDP2 框架兼容的分布式 NorMuon 版本。虽然之前关于分布式 Muon 的研究是使用 ZeRO-1 和 Megatron-LM 实现的，但 FSDP2 可以提供更高的灵活性和内存效率。然而，直接将之前的分布式方法应用于 FSDP2 会导致大量的重复计算，因为 FSDP2 将几乎所有参数分片到不同的设备上。我们的实现通过将正交化计算分布到不同的设备上来解决这个问题，从而消除冗余计算并保持负载均衡。此外，我们利用 FSDP2 的行级参数分片机制，实现了高效的神经元级归一化，而无需增加额外的通信开销。

总而言之，我们的贡献体现在三个方面：

• 我们提出了一种名为 NorMuon 的简单高效的优化器，它结合了 Muon 正交化和神经元级自适应学习率。NorMuon 保持神经元范数的一致性，从而确保参数利用率的均衡性，同时保留 Muon 正交化所实现的低条件数。
• 我们在 FSDP2 框架下开发了一种高效的分布式实现。通过精心协调分片优化器的状态，我们收集更新后的动量，并将 Muon 正交化计算均匀地分布在各个GPU上，从而在可控的通信和计算开销下实现了最佳的内存效率。
• 通过在 LLM 预训练的多个尺度上进行广泛的实验，我们证明正交化和分块自适应学习率是互补的而不是冲突的，它们的结合比单独使用任何一种方法都能产生更优越的训练动态。

2. Related Works and Background

2.1 Related Works

Adaptive Gradient Methods。引入逐参数自适应学习率对深度网络的训练至关重要。诸如 AdaGrad、RMSProp、Adam 和 AdamW 等优化器使用一阶矩和二阶矩估计来单独调整每个权重的步长。这种逐坐标预处理提高了异构设置下的稳定性和收敛性，并已成为 LLM 训练的事实标准。然而，独立处理每个权重忽略了神经网络层的底层结构，并且由于每个参数需要存储两个额外的张量而导致高内存开销。这种内存成本促使了 AdaFactor 等技术的出现，它将二阶矩累加器按行和列进行分解以减少内存占用。类似地，Adam-mini 将参数划分为块（例如，每个神经元的权重），并为每个块分配一个学习率，从而在不同模型规模上达到与 AdamW 相同的性能，同时将内存成本减半。GaLore 在由梯度奇异值分解 (SVD) 得到的低秩子空间中保持动量，但其有效性会随着序列长度的增加而降低。 Lion 采用坐标方向的符号更新，放弃二阶矩估计以节省内存。

Second-order Methods。与此同时，其他优化器通过耦合参数更新来捕捉损失曲面的丰富几何信息。K-FAC 及其变体超越了单个坐标，近似捕捉了曲率信息，从而捕捉了参数间的相关性。Shampoo 及其分布式变体采用了克罗内克因子预处理器，并在实践中展现了优异的性能。最近， SOAP 建立了 Shampoo 和 Adafactor 之间的联系，进一步提升了收敛性能。尽管取得了这些进展， Shampoo 和 SOAP 仍然会产生大量的内存开销和计算开销，这限制了它们在LLM规模上的应用。

Orthogonal Update Methods。Muon（动量正交化算法）是一项突破性成果，它利用矩阵几何，而无需付出二阶方法的全部代价。Muon 对动量进行近似极分解（通过 Newton-Schulz 迭代），提取其正交分量，并且无需存储二阶动量。因此，与 Adam 相比，Muon 同时提高了收敛速度和内存效率，展现出在模型预训练扩展方面的巨大潜力。我们在第 2.2 节中详细描述了 Muon 算法。最近，Dion 扩展了正交更新范式，使其在分布式环境中具有更高的通信和计算效率。Dion 使用摊销幂迭代的低秩正交化方案，而不是完整的 Newton-Schulz 迭代，并将不同设备上的动量缓冲区解耦，从而避免了完全梯度同步。

2.2 Background: Muon optimizer

Muon 是一款专为神经网络隐藏层中的二维权重矩阵设计的优化器。其关键创新在于在应用参数更新之前对动量进行正交化，从而改善优化轨迹的条件数。形式上，在第ttt次迭代中，给定权重矩阵Wt−1W_{t−1}Wt−1​、学习率ηtη_tηt​和损失函数LLL，Muon 维护一个一阶动量MtM_tMt​，并按如下方式计算更新：

Mt=μMt−1+∇L(Wt−1),(1)M_t=\mu M_{t-1}+\nabla L(W_{t-1}),\tag{1}Mt​=μMt−1​+∇L(Wt−1​),(1)

Ot=NS5(Mt),(2)O_t=NS5(M_t),\tag{2}Ot​=NS5(Mt​),(2)

Wt=Wt−1−ηtOt,(3)W_t=W_{t-1}-η_tO_t,\tag{3}Wt​=Wt−1​−ηt​Ot​,(3)

其中M0=0M_0 = \textbf{0}M0​=0，µµµ为动量系数。关键分量是正交化算子NS5(⋅)NS5(·)NS5(⋅)，其目的是近似动量矩阵的正交投影：

Ortho(M)=argminO{∣∣O−M∣∣F:OTO=I or OOT=I}.(4)Ortho(M)=\mathop{argmin}\limits_{O}\{||O-M||_F:O^TO=I~or~OO^T=I\}.\tag{4}Ortho(M)=Oargmin​{∣∣O−M∣∣F​:OTO=IorOOT=I}.(4)

Muon 算法通过固定次数的牛顿-舒尔茨迭代来近似实现这种正交化。从弗罗贝尼乌斯归一化的动量X0=Mt/∥Mt∥FX_0 = M_t/∥M_t∥_FX0​=Mt​/∥Mt​∥F​开始，该算法执行NNN次迭代（通常N=5N = 5N=5）：

Xk=aXk−1+b(Xk−1Xk−1T)Xk−1+c(Xk−1Xk−1T)2Xk−1,k=1,...,N,(5)X_k=aX_{k-1}+b(X_{k-1}X^T_{k-1})X_{k-1}+c(X_{k-1}X^T_{k-1})^2X_{k-1},\quad k=1,...,N,\tag{5}Xk​=aXk−1​+b(Xk−1​Xk−1T​)Xk−1​+c(Xk−1​Xk−1T​)2Xk−1​,k=1,...,N,(5)

最终的正交化更新Ot=XNO_t = X_NOt​=XN​。系数(a,b,c)(a, b, c)(a,b,c)经过精心选择，使得更新矩阵的奇异值收敛于 1。在实践中，Muon 通常仅应用于隐藏层中的二维权重矩阵，而标量参数、偏置向量、嵌入和反嵌入层则继续使用 Adam 等标准优化器。

3.Method

在本节中，我们介绍 NorMuon，其目的是将 Muon 的正交化与基于分块自适应学习率的观察结果相结合，即近似正交化更新在每个神经元的更新方向范数上可能会经历较高的方差。

3.1 NorMuon

我们在算法 1 中提出了更新规则。该算法维护两个动量状态：Muon 使用的标准一阶动量Mt∈Rm×n\textbf M_t ∈ \mathbb R^{m×n}Mt​∈Rm×n（第 5 行），以及平均二阶动量vt∈Rm\textbf v_t ∈ \mathbb R^mvt​∈Rm，它跟踪每个神经元更新方向的平方幅值（第 7 行）。重要的是，vtv_tvt​所需的额外内存开销极小，与m×nm × nm×n的一阶动量相比，它仅需存储mmm个标量。

在每次迭代中，给定梯度，我们首先遵循 Muon 更新规则更新一阶动量，并应用牛顿-舒尔茨迭代进行正交化（第4-6行），从而得到具有改进条件数的OtO_tOt​。我们并不直接使用这个正交化的更新，而是计算逐行统计量来捕捉每个神经元的更新幅度。具体来说，我们计算OtO_tOt​每一行的列均方值（第7行）。该统计量通过衰减率为β2β_2β2​的指数移动平均累积到我们的平均二阶动量vt\textbf v_tvt​中。然后，我们应用vt\textbf v_tvt​进行逐行归一化（第9行）。这个二阶动量类似于 Adam-mini 的块级降维统计，其中我们将每个神经元（即每一行）视为一个块。如图1所示，这种归一化降低了神经元间更新幅度的方差，同时保留了良好的条件数特性。

我们观察到，行归一化后得到的方向范数更大。因此，在更新过程中，我们添加了一个学习率缩放，以保持与 Adam 的均方根范数相似的均方根范数（第 10 行）。

我们注意到，在理想情况下，当OtO_tOt​严格正交化（即不使用NS5NS5NS5近似）时，对于m≤nm ≤ nm≤n的满秩矩阵，每个神经元的范数将严格为111。对于这类矩阵，逐神经元归一化并无益处。然而，由于实际应用中正交化是近似的，我们观察到，即使对于m≤nm ≤ nm≤n的矩阵，这种归一化仍然是必要且有益的（已在 4.1.3 节中验证）。

3.2 Distributed NorMuon

随着 LLM 训练规模的扩大，分布式训练对于解决内存限制和提高计算效率都至关重要。我们开发了一个与 FSDP2 框架兼容的分布式 NorMuon 版本，该版本采用 ZeRO-3 式分片技术，将优化器状态、参数和梯度分布到多个设备上。

虽然像 Adam 这样的基于坐标的优化器能够自然地扩展到分布式环境，但 NorMuon 由于 Muon 的正交化步骤需要访问完整的动量矩阵，因此面临着独特的挑战。现有的 Muon 分布式实现会在所有设备上收集完整的动量矩阵并重复执行正交化计算。我们通过近乎均匀地将参数分配给不同的设备来避免这种重复计算的开销。

算法 2 展示了我们的分布式实现方案。与算法 1 相比，主要修改之处在于：

• Efficient Orthogonalization Distribution(line 5-9)：为了确保工作分配的均匀性，我们首先按矩阵大小对参数列表进行排序（第 2 行），而不是让所有设备都计算所有参数的正交化，其中 Numel(·) 用于统计每个矩阵中的元素数量。然后，我们使用轮询机制将每个参数张量分配给一个特定的设备。只有被分配的设备才能通过全集通信收集完整的动量矩阵，并执行牛顿-舒尔茨正交化（第 5-8 行），之后将结果分发回所有设备（第 9 行）。这种方法既消除了冗余计算，又保持了设备间的负载均衡。
• Shard-Local Row Normalization(lines 10-12)：我们设计的关键优势在于，逐行归一化完全在本地分片上进行，无需额外的通信。这得益于 FSDP2 采用逐行分片，确保每个设备都拥有权重矩阵的完整行。因此，行统计信息的计算和归一化过程可以独立进行。

3.3 Overhead Analysis

Memory Overhead。NorMuon 保持了 Muon 的内存效率。对于权重矩阵W∈Rm×nW ∈ \mathbb R^{m×n}W∈Rm×n，各优化器的优化器状态内存消耗分别为：(1) Adam：2mn2mn2mn（一阶和二阶动量）；(2) Muon：mnmnmn（仅一阶动量）；(3) NorMuon：m(n+1)m(n + 1)m(n+1)（一阶动量 + 每个神经元的二阶统计量）。与 Muon 相比，NorMuon 的额外内存开销可以忽略不计（1/n1/n1/n因子），同时内存效率比 Adam 高约 50%。

Communication Overhead。与标准 FSDP 训练相比，NorMuon 引入了适度的额外通信。在采用 AdamW 的 FP32 训练下，每个参数的通信开销为：4 字节（forward all-gather）+ 4 字节（backward all-gather）+ 4 字节（gradient reduce-scatter）= 12 字节。对于以 BF16 精度计算的 NS5 迭代，NorMuon 需要：12 字节（standard FSDP communication）+ 2 字节（momentum gather, BF16）+ 2 字节（update scatter, BF16）= 16 字节。

这相当于通信量增加了 33%。当参数使用 BF16 时，相对开销会增加到 50%。然而，这种通信可以与正交化计算重叠进行，从而最大限度地减少延迟的影响。在我们的实验中（4.1.4 节），我们证明了 NorMuon 的每次迭代延迟仅比 AdamW 高 3%，同时收敛效率却显著更高。

4.Experiments

在本节中，我们针对四种不同模型规模（124M、350M、1,1B 和 5.4B 参数）进行了预训练实验，以验证 NorMuon 的有效性。对于较大的模型（1.1B 和 5.4B 参数），我们采用了先前架构扩展工作中的实验设置，结果和配置见 4.1 节。对于较小的模型（124M 和 350M 参数），我们遵循 Modded-NanoGPT 的实验设置，结果和设置见 4.2 节。我们还进行了大量的消融实验来验证我们的设计选择，并进行了详细的效率分析（4.1.3 节和 4.1.4 节）。