随着大语言模型训练规模持续扩大，FSDP已成为缓解显存压力、支撑千亿至万亿参数模型训练的重要分布式并行技术。与此同时，优化器也从AdamW逐步向Muon演进。但是，Muon优化器的Newton-Schulz计算需要对完整参数矩阵进行正交化更新，而FSDP通常将参数、梯度和优化器状态切分存储在不同设备上，二者在计算粒度和数据布局上存在天然冲突。

对此，MindSpore HyperParallel设计并实现了Muon深度优化方案，在昇腾512卡集群的工程评测中，实测训练千亿参数的 MoE模型，Muon优化器的单次计算耗时由2700 ms 大幅缩减至 450 ms。

本文将带大家一起深度解析基于 HyperParallel-FSDP设计并实现的面向昇腾超节点的Muon融合优化方案 。

Muon优化器在FSDP下的
三大性能挑战

Newton-Schulz正交化迭代需要完整的矩阵参与计算，但是在FSDP策略下，参数按维度0被切分到不同的rank上，每个rank仅持有本地分片，无法直接进行全局正交化操作。因此，必须引入分布式通信来收集完整矩阵，计算完成后再将结果分片发回各rank。这种Muon优化器自身机制与FSDP分片策略之间的天然冲突，为大规模集群下的性能优化带来了三大挑战：

第一，全矩阵收集通信开销巨大。各Rank需通过集合通信操作从所有节点收集完整权重矩阵。在大规模集群中，All-Gather涉及跨节点、跨交换机的海量数据传输，通信延迟与带宽压力显著增加，易成为训练瓶颈。

第二，Newton-Schulz导致Host-bound。Muon优化器内部的Newton-Schulz迭代过程所处理的矩阵规模往往较小（如逐层参数矩阵），计算任务细粒度、短耗时，导致CPU侧下发Kernel的频率过高，Host端调度开销相对突出，难以充分利用NPU计算资源。

第三，Newton-Schulz计算耗时较长。尽管单次迭代计算量适中，但每轮参数更新需在优化器中执行多次矩阵乘法和加法运算（通常5~10次迭代），叠加所有参数层后总体计算延迟明显增长，易成为训练流水线的瓶颈。此外，各节点完成All-Gather全局聚合后独立执行相同的正交化运算，在HSDP场景下DP副本间的冗余计算问题进一步加剧性能劣化。

三重挑战共同制约了Muon优化器在FSDP大规模分布式训练场景下的端到端吞吐效率。 

基于HyperParallel-FSDP的
Muon+优化器实现

Muon参数更新的核心前提是获取完整的全局矩阵（Full Matrix）。在FSDP分片域内，各节点通过All-Gather聚合DP Group内的梯度动量M_t获得完整M_t后，执行Newton-Schulz迭代得到正交矩阵O_t，各节点再依据索引偏移量（Offset）截取对应局部切片（Local Shard）进行分片参数更新。该方案完全遵循FSDP2原生体系，但存在Newton-Schulz被所有DP Rank重复计算的核心痛点，Compute Redundancy等于DP Size，冗余随模型规模增大而愈发严重，如图1所示。

图1. AllGather方案冗余计算说明

在图1中，Rank0-Rank7代表不同的NPU，P0-P10分别代表不同的param参数对应的梯度动量M_t。p0_0..p10_7分别代表不同param进行FSDP切分后在卡上的分片。以P10为例，通过AllGather之后，每张卡上都获得完整的P10，并进行了相同的Newton-Schulz计算，因此存在DP倍的冗余计算。 

All-To-All方案理论上可实现零冗余通信与计算，但是在跨网络超平面（Cross-Hyperplane）时因通信需建立全局全连接，无法在万卡级规模下稳定使用。All-Gather 在底层天然由“机内通信”与“机间通信”两阶段解耦构成，这种分阶段的通信组合有效规避了跨超平面的通信资源限制，保障了大规模训练的网络可用性。 我们针对原始的AllGather方案，设计出一种更理想的去冗余方式，即：

我们对 FSDP 域内的参数进行精细化分组。在DP 组内执行All-Gather聚合完整的梯度动量M_t 后，各 Rank 仅保留和处理分配给当前Rank的子矩阵切片，分组并行开展Newton-Schulz迭代。计算完成后，通过Broadcast通信扩散正交化结果，各卡各自提取对应分片更新本地参数，整个流程如图2所示。

图2. All_Gather DP去冗余方案的Muon优化器实现示意图

该架构兼顾了万卡网络的高可用性与计算效率。相对于AllGather方案，去冗余方案计算耗时从 485 ms 大幅收敛至 129 ms，性能整体提升 70+%， 如图3所示。 

图3. All_Gather DP 去冗余效果

至此，我们完成了HyperParallel-FSDP下Muon优化器的核心架构实现，下面我们将在第3章介绍针对Newton-Schulz计算的性能优化。

在完成AllGather + DP去冗余的Muon优化器方案实现后，我们通过Profiling分析，Newton-Schulz的计算开销成为了系统的主要瓶颈。我们针对Newton-Schulz计算进一步优化，设计了 1）高维参数通信计算去冗余优化； 2）矩阵融合优化；3）HSDP分组去冗余技术 。下面我们将依次介绍三个特性的实现方案。

我们通过Profiling分析，专家权重的Newton-Schulz计算占比耗时长。由于Muon的正交化空间变换仅依赖后两维（即权重矩阵自身结构），对于 dim > 2 的高维专家参数（如 3D Expert 参数 [E, M, N]），其Batch（专家 E 维度）之间在数学上完全独立。因此，无需进行跨卡的全局通信聚合，可直接按本地分片的 Batch维度独立并行执行Muon计算，实现了完全的免通信通算去冗余。具体实现方案如图4所示。

图4. 专家计算去冗余方案的Muon优化器实现示意图

我们以专家数量E=4的梯度动量矩阵M_t为例来说明该机制。在未开启 FSDP 切分时，矩阵完整维度为 [4, 4, 4]；当FSDP Shard Size为2时，Rank0与Rank1仅各自持有 [2, 4, 4] 的局部切片。在 2.1 节的AllGather基础方案中，系统需先通过全局通信将矩阵还原至 [4, 4, 4]，利用 Mask 操作在特定Rank完成全量计算后，再通过Broadcast同步正交化结果，最后由各节点截取（Slice）对应分片更新参数。而本高维参数去冗余方案，则巧妙利用了Newton-Schulz迭代在Batch维度的数学独立性：Rank0与 Rank1彻底免除集合通信步骤，直接基于本地 [2, 4, 4] 切片并行执行独立正交化。该设计将高维专家参数的通信开销完全清零，并使计算量缩减了FSDP Shard Size倍。

图5. 专家计算去冗余方案的优化效果

图5中实际训练的硬件性能Profiling结果分析，专家矩阵的绝对耗时由4.9 ms大幅缩减至1.1 ms，耗时骤降约77.5%。此外，该设计在底层逻辑上有效规避了传统全量All-Gather与Broadcast操作的高昂通信开销，实现了计算效率与通信代价的双重优化。

Profiling结果表明底层内存分配与释放（alloc/free）的耗时占比偏高。进一步分析发现，当 Muon 优化器处理大量小尺寸或小批次矩阵（如 [24, 10240]、[32, 24, 10240]）时，NPU核心（Cube）实际用于矩阵乘的计算时间极短。系统的绝大部分开销集中于Host端的算子下发（Host Dispatch）、小块显存的频繁申请与流同步（Stream Sync），以及内核调度延迟（Kernel Scheduling Overhead）。

如图6所示，借鉴高维参数通信与计算的去冗余思想，我们针对尾部两维形状相同的参数实施了算子融合（Kernel Fusion），通过在Host端构建统一的Batch维度，我们将原本零散的海量小矩阵重组为连续的高维张量 [B, M, N]，从而在底层硬件触发高效的批量矩阵乘（BMM）算子。

图6. 矩阵融合优化方案示意图

该优化显著降低了执行延迟。以74个形状为 [24, 10240] 的参数为例，批量执行Newton-Schulz迭代的总耗时降至2.4 ms，远低于优化前逐个参数执行的理论总耗时（74*1.0  = 74 ms）。

图7. 矩阵融合优化方案示意图

HSDP（Hybrid Sharded Data Parallel，混合分片数据并行）主要是为了解决FSDP在扩展到大规模集群（如万卡级别）时遇到的高昂的通信开销以及高的通信延迟问题。HSDP的切分示意图如图8所示，图示是4张NPU/GPU场景，普通 FSDP 4 张卡会把模型切成4份；HSDP (shard size=2），则在2张卡内切FSDP，模型切2份。在万卡集群训练中，我们可以通过开启HSDP，通过设置shard size < 超节点数，限制高频率的AllGather和 ReduceScatter 操作在超节点内部，仅有权重副本间的梯度同步AllReduce为跨超节点通信。 

图8. FSDP 与 HSDP 参数切分示意图

在HSDP下，AllGather+DP去冗余的Newton-Schulz计算流程如图9 (1)所示。我们可以发现，在Rank0&Rank2，Rank1&Rank3执行相同的Newton-Schulz计算。这里提出新的命题，是否可以在HSDP中实现分组计算减少每张卡的计算冗余。

图9. HSDP下AllGather+DP去冗余的
Newton-Schulz计算流程示意图

针对HSDP下副本间的冗余计算问题，我们设计了一个HSDP分组计算方案，如图9所示。在标准 HSDP 架构中，各副本组域内天然维护着全量模型参数副本，利用这一特性，HSDP 组内分组计算，通过在各副本间广播（Broadcast）更新后的参数。其中，为了更均衡的分配参数，hsdp副本之间通过贪心算法分配参数。

经过HSDP分组优化冗余计算之后，整体性能提升1倍，计算时间降低7倍。

图10. HSDP 分组更新参数优化效果

HyperParallel-Muon优化代码已上库，欢迎体验：

MindSpore HyperParallel开源代码仓：
https://gitcode.com/mindspore/hyper-parallel

我们诚挚邀请各位开发者、研究者加入MindSpore HyperParallel项目。

无论是贡献代码、完善文档，还是提出改进建议、分享实际落地中的踩坑经验——您的每一次参与，都将推动大模型分布式并行技术的边界。