⚠️ 在开始阅读之前如果你对实时 Agent / 数字人 / 多模态系统 / LiveKit 架构感兴趣欢迎先到 GitHub 给项目点一个 ⭐ Star这是对开源作者最大的支持。AlphaAvatar 项目地址强烈建议先收藏该项目正在持续更新维护 https://github.com/AlphaAvatar/AlphaAvatarAIPapers 项目地址具有更全的有关LLM/Agent/Speech/Visual/Omni论文分类 https://github.com/AlphaAvatar/AIPaperNotes摘要优化器的选择对大语言模型LLM的训练效率和计算成本有着显著的影响。近年来Muon 优化器通过正交化参数更新并利用更好的条件数来改善优化几何结构展现出了良好的应用前景。尽管 Muon 已成为 Adam 的有力竞争者但如何将二者的优势结合起来却尚未得到系统性的探索。本文提出了一种名为NorMuonNeuron-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 Background2.1 Related WorksAdaptive 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 optimizerMuon 是一款专为神经网络隐藏层中的二维权重矩阵设计的优化器。其关键创新在于在应用参数更新之前对动量进行正交化从而改善优化轨迹的条件数。形式上在第ttt次迭代中给定权重矩阵Wt−1W_{t−1}Wt−1​、学习率ηtη_tηt​和损失函数LLLMuon 维护一个一阶动量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)OtNS5(Mt),(2)O_tNS5(M_t),\tag{2}Ot​NS5(Mt​),(2)WtWt−1−ηtOt,(3)W_tW_{t-1}-η_tO_t,\tag{3}Wt​Wt−1​−ηt​Ot​,(3)其中M00M_0 \textbf{0}M0​0µµµ为动量系数。关键分量是正交化算子NS5(⋅)NS5(·)NS5(⋅)其目的是近似动量矩阵的正交投影Ortho(M)argminO{∣∣O−M∣∣F:OTOI or OOTI}.(4)Ortho(M)\mathop{argmin}\limits_{O}\{||O-M||_F:O^TOI~or~OO^TI\}.\tag{4}Ortho(M)Oargmin​{∣∣O−M∣∣F​:OTOIorOOTI}.(4)Muon 算法通过固定次数的牛顿-舒尔茨迭代来近似实现这种正交化。从弗罗贝尼乌斯归一化的动量X0Mt/∥Mt∥FX_0 M_t/∥M_t∥_FX0​Mt​/∥Mt​∥F​开始该算法执行NNN次迭代通常N5N 5N5XkaXk−1b(Xk−1Xk−1T)Xk−1c(Xk−1Xk−1T)2Xk−1,k1,...,N,(5)X_kaX_{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 k1,...,N,\tag{5}Xk​aXk−1​b(Xk−1​Xk−1T​)Xk−1​c(Xk−1​Xk−1T​)2Xk−1​,k1,...,N,(5)最终的正交化更新OtXNO_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 AnalysisMemory Overhead。NorMuon 保持了 Muon 的内存效率。对于权重矩阵W∈Rm×nW ∈ \mathbb R^{m×n}W∈Rm×n各优化器的优化器状态内存消耗分别为(1) Adam2mn2mn2mn一阶和二阶动量(2) Muonmnmnmn仅一阶动量(3) NorMuonm(n1)m(n 1)m(n1)一阶动量 每个神经元的二阶统计量。与 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 节。