optimize muon 2D sharding

[cszdrg]

PR Category

Performance Optimization

PR Types

Improvements

Description

优化muon 2d参数的切分
步骤如下：

1. 通信拿到所有rank和机器node的映射
2. 通信是的每个rank可以看到全局的muon 2D参数
3. 贪心的将数据往rank上放，放的顺序如下
   a. 找到压力最小的机器
   b. 找到拥挤程度最低的机器
   c. 找到该机器上压力最小的rank

【测试结果】：在8机情况下，开启ep=4，pp=0的测试条件下，生成checkpoint

1. 正确性检验：
   a. 其中checkpoint-10、checkpoint-40为老版本代码生成；checkpoint-20、checkpoint-30为新版本代码生成。
   b. 四次训练的loss值如下，正常接续

c. 收敛结果：

3. 负载均衡检验：
   文件大小：
   老版本的所有权重均聚集于排序靠前的rank

新版本权重按照机器均分，高负载机器尽可能远离（1号机器和8号机器）

是否引起精度变化

否

[Paddle-Bot]

Paddle-Bot Review Board （review完成)

序号	位置	优先级	规则来源	状态
1	PR 验证信息缺失		仓库规则：PR 信息与验证	🟡
2	多进程通信路径缺少验证		仓库规则：分布式测试与验证	🚧

^{Powered by Nyanpasu with gpt-5.6-sol high, please check the suggestions carefully.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 72.89720% with 29 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@df1fbe1). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...d/fleet/meta_optimizers/muon_sharding_optimizer.py	72.64%	29 Missing ⚠️

❌ Your patch status has failed because the patch coverage (72.89%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             develop   #79763   +/-   ##
==========================================
  Coverage           ?   72.89%           
==========================================
  Files              ?        2           
  Lines              ?      107           
  Branches           ?        0           
==========================================
  Hits               ?       78           
  Misses             ?       29           
  Partials           ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Paddle-Bot]

Failed CI看板

流水线名称	问题标签	修复建议	日志片段
Require review-bot approval	当前 head 未获得机器人或人工审批	【与当前改动无关】请让 risemeup1111、Paddle-Bot 或指定审批人针对当前 SHA 77465e5 完成批准。	Job日志
Fleet Unit test (multi-card)	test_block_attn_res_vpp.py 失败，首个错误缺失	【无法判断】补充或获取该测试的完整失败堆栈，再判断是否与 Muon 2D 分片改动有关；当前证据不足以修改 PR。	Job日志
Coverage test	Python diff coverage 低于门槛	【与当前改动有关】为新增的机器均衡分片分支、初始化路径及通信信息收集逻辑补充覆盖测试，使 Python diff coverage 达到 90%。	Job日志

日志分析报告

失败的测试 case:

1. Require review-bot approval
   - 执行审批状态检查。
   - 当前 head 上 risemeup1111 和 Paddle-Bot 均未 APPROVED，随后 job 以 exit code 1 结束。
   - 这是审批门禁失败，不是代码测试失败。

2. Fleet Unit test (multi-card)
   - 共执行 96 个测试。
   - test_pp_layer.py 通过。
   - test_block_attn_res_vpp.py 失败，但快照裁剪掉了该测试的首个具体错误。

3. Coverage test
   - 单元测试全部通过。
   - Assert Diff Coverage 输出 no data found。
   - Assert Python Diff Coverage：expected >= 90.0%，actual 72.9%，失败并以 exit code 9 结束。

根本原因分析:

Coverage test 的确定性根因是新增 Python 代码的差异覆盖率只有 72.9%，低于仓库要求的 90%。PR 在 muon_sharding_optimizer.py 中新增了机器拓扑收集、跨 rank 信息聚合、机器均衡 owner 选择及多种边界分支；现有测试主要直接测试分片结果，未能覆盖全部新增路径。因此这是与当前 PR 直接相关的覆盖率门禁失败，不是单元测试执行错误。

Fleet Unit test (multi-card) 仅能确认 test_block_attn_res_vpp.py 失败，日志没有提供首个异常、断言或退出原因。该测试文件不在 PR diff 中，且快照没有证明它调用了新增 Muon 分片逻辑，因此不能可靠归因于当前改动，也不能确认是环境或不稳定问题。

Require review-bot approval 失败原因是当前 commit 没有满足审批状态条件。该 job 的逻辑独立于代码执行和覆盖率结果；快照没有显示任何机器人或审批人的当前 head APPROVED 状态。三个失败 job 之间没有证据表明存在下游产物缺失或级联关系。

修复建议:

1. 【与当前改动无关】Require review-bot approval：让机器人或 sneaxiy、From00 针对当前 SHA 完成审批后重新运行该门禁。
2. 【无法判断】Fleet Unit test (multi-card)：获取 test_block_attn_res_vpp.py 的完整错误输出；若确认是偶发多卡或通信问题，重跑该 multi-card job，否则根据具体堆栈修复对应代码。
3. 【与当前改动有关】Coverage test：为新增机器均衡分片的初始化、rank-to-machine 收集、跨 rank color-group 聚合、单 rank 或多 group、不同 dtype 和 active-rank 分支补充测试，确认 Python diff coverage 达到至少 90%。

Powered by Nyanpasu with gpt-5.6-luna 默认推理级别, please check the suggestions carefully.

[cszdrg]

/re-run all-failed

[Paddle-Bot]

开启该配置后，__init__ 会让全局所有 rank 依次进入两次 all_gather_object，且其结果直接决定后续 reduce/broadcast 的 owner；但本次新增测试只绕过初始化直接调用分区 helper，未验证开关透传、全局 collective 对称参与及实际通信映射。请补充启用 machine_balanced_2d_partition 的代表性多进程测试，至少覆盖 PP/EP 分组下各 rank 完成初始化并执行一次通信。

[Paddle-Bot]

此前指出的多进程通信测试缺口在当前提交中仍未解决。

^{Powered by Nyanpasu with gpt-5.6-sol high, please check the suggestions carefully.}