hyper optimizer

HyperParallel FSDP Muon优化器性能优化 RFC

1. FSDP域分布式Muon更新方案实现
   任务：
   实现Muon优化器在FSDP分片域下的分布式更新流程。Muon的Newton-Schulz正交化计算需要完整矩阵参与，而FSDP将参数切分到不同Rank上，二者存在结构性矛盾。需实现以下方案：
   (1) AllGather基础方案：各节点通过AllGather聚合DP Group内完整梯度动量M_t，执行Newton-Schulz迭代后按Offset截取本地分片更新。该方案存在DP Size倍冗余计算问题。
   (2) AllToAll方案验证：通过AllToAll通信实现零冗余Newton-Schulz计算，每个Rank仅处理分配给自己的完整参数。需验证千卡/万卡规模下AllToAll跨网络超平面全连接通信是否触发QP资源耗尽（Failed to allocate resource[qp]），评估方案可行性。
   (3) AllGather + DP去冗余方案（推荐）：AllGather聚合完整M_t后，对FSDP域内参数精细化分组，各Rank仅保留和处理分配的子矩阵切片，分组并行执行Newton-Schulz迭代，计算完成后通过Broadcast扩散正交化结果，各卡提取对应分片更新本地参数。该方案兼顾万卡网络高可用性与计算效率，计算耗时从485ms降至129ms，性能提升70%+。

参考：图2-图6

2. 高维参数通信计算去冗余（Muon亲和的专家切分）
   任务：
   针对MoE模型中高维专家参数（如3D Expert参数[E, M, N]），利用Newton-Schulz正交化仅依赖后两维（权重矩阵自身结构）、Batch（专家E维度）之间数学完全独立的特性，实现免通信的通算去冗余。
   具体实现：

• FSDP切分后各Rank仅持有[E/shard_size, M, N]的局部切片
• 各Rank直接基于本地切片并行执行独立正交化，免除AllGather与Broadcast通信
• 计算量缩减FSDP Shard Size倍，通信开销清零
  实测效果：专家矩阵绝对耗时从4.9ms降至1.1ms，耗时降低约77.5%。

参考：图7-图8

3. Newton-Schulz矩阵融合优化
   任务：
   解决Muon优化器处理大量小尺寸/小批次矩阵时Host-bound问题。当处理[24, 10240]、[32, 24, 10240]等小矩阵时，NPU Cube实际矩阵乘计算时间极短，绝大部分开销集中于Host端算子下发、小块显存频繁申请/释放与流同步、内核调度延迟。
   具体实现：

• 借鉴高维参数去冗余思想，针对尾部两维形状相同的参数实施算子融合（Kernel Fusion）
• 在Host端构建统一Batch维度，将零散小矩阵重组为连续高维张量[B, M, N]
• 底层硬件触发高效批量矩阵乘（BMM）算子，减少Host侧算子下发次数和显存分配/释放次数
  实测效果：74个[24, 10240]参数的Newton-Schulz迭代总耗时从理论74ms（74*1.0ms）降至2.4ms。

参考：图9-图10

4. HSDP分组去冗余计算
   任务：
   解决HSDP场景下副本间Newton-Schulz冗余计算问题。标准HSDP架构中，各副本组域内维护全量模型参数副本，导致不同副本组执行相同的Newton-Schulz计算。
   具体实现：

• 利用HSDP各副本组域内天然维护全量模型参数副本的特性，实现HSDP组内分组计算
• 各副本间通过Broadcast广播更新后的参数，减少每张卡的冗余计算
• 采用贪心算法在HSDP副本之间均衡分配参数，确保计算负载均衡
  实测效果：整体性能提升1倍，计算时间降低7倍。

参考：图11-图13

5. Muon优化器端到端集成与验证
   任务：
   将上述优化方案集成至HyperParallel FSDP2框架，在昇腾超节点512 Die集群上进行千亿级MoE模型端到端训练验证。
   验证指标：

• Muon优化器单次计算耗时：目标从2700ms优化至450ms（6倍提升）
• 训练吞吐与MFU指标
• 千卡/万卡规模下通信稳定性（尤其AllToAll方案的QP资源限制验证）
• HSDP + Muon联合场景下的正确性与性能

完成情况：
待开发

6. Muon优化器Weight Decay适配
   任务：
   规模化训练时，原始Muon训练的模型权重和层输出的RMS会过度增长，超出bf16精度范围，损害模型性能。需引入标准AdamW风格的权重衰减机制，确保Muon在"过度训练"状态下优于原始Muon和AdamW。
   具体实现：

• 在Muon迭代公式中引入Weight Decay项
• 验证不同Weight Decay系数对训练精度和收敛速度的影响
• 确保与FSDP分片策略兼容

完成情况：
待开发

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 204
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/204

No files, tests, or code entry points are named. Start by locating the HyperParallel FSDP2 Muon optimizer implementation and its distributed-training validation coverage, then determine which of the six proposed workstreams has an agreed scope. Done requires integrated correctness and performance validation against the listed communication, runtime, throughput, MFU, and weight-decay goals.