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分片策略兼容

完成情况：
待开发