dev muon sharding optimizer comm overlap#79428
Conversation
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
发现两处需要先处理的兼容性问题，详情已放在行级评论里：主要是 Muon 的新 backward-hook overlap 路径缺少与 pipeline parallel 的互斥保护，以及 AMP GradScaler 仍按 _comm_buffers 等待 overlap 通信，当前 Muon 没有暴露这个集合。CI 目前仍有检查在运行中，建议修复后补充 Muon + comm_overlap 覆盖，尤其包含 GradScaler 场景。
| fires once all params of a buffer have checked in, and reduce / | ||
| reduce-scatter over the fused buffer are order-independent. | ||
| """ | ||
| overlap_buffers = [] |
There was a problem hiding this comment.
新 hook 让 Muon 也走 sharding_configs.comm_overlap=True，但 distributed_scaler() 的 overlap 分支会先断言 hasattr(optimizer, "_comm_buffers")，并遍历 optimizer._comm_buffers 等待通信完成后再做 finite/unscale。HybridParallelOptimizer.__getattr__ 会把这个访问转发到内部的 MuonShardingOptimizer，而这里新建的只有 comm_buffer_2d 和 _comm_buffer_list，没有 _comm_buffers；因此 Muon + comm_overlap + AMP GradScaler 在 scaler.step(optimizer) 时会直接断言失败，或即使绕过断言也无法等待新启动的 2D/1D 通信任务。
请把 Muon overlap 暴露成 scaler 期望的缓冲区集合，或同步更新 distributed_scaler 识别 Muon。一个局部修复形态可以是构建完 1D/2D buffer 后设置：
self._comm_buffers = []
if self._use_fuse_gradients and not self.use_group_call_opt:
self._comm_buffers.extend(self.comm_buffer_2d)
self._comm_buffers.extend(self._comm_buffer_list)There was a problem hiding this comment.
这版已经补了 _comm_buffers，普通 overlap 路径的断言问题解决了；但 comm_group_call_opt=True 仍会漏掉 2D 梯度通信。现在 _comm_buffers 在 use_group_call_opt 时只包含 _comm_buffer_list，而 2D buffer 仍只在 reduce_gradients() 的 group-call 分支里同步 reduce。AMP distributed_scaler() 进入 sharding_configs.comm_overlap 分支后会先 buffer.scale_grads()，随后把 optimizer._sharding_enable = False，这样 HybridParallelOptimizer.step() 不再调用 reduce_gradients()，group-call 的 2D reduce 分支就完全不会执行，2D Muon 参数会用本 rank 未 reduce 的梯度更新。
建议不要只把 hook 驱动的 buffer 暴露给 scaler，而是确保 scaler 等待/触发了所有会被禁用的 sharding 通信。可以选择在 _comm_buffers 中也包含 comm_buffer_2d 并让 scaler 之后能完成 group-call 2D reduce，或在 distributed_scaler/Muon 中把 group-call 2D reduce 抽成可在 unscale 前调用的同步路径。例如修复形态需要覆盖：
# scaler disables optimizer-side reduce_gradients(), so group_call_opt 2D reduce
# must be completed before optimizer._sharding_enable is set False.
if self._use_fuse_gradients:
self._comm_buffers.extend(self.comm_buffer_2d)
self._comm_buffers.extend(self._comm_buffer_list)
# and make group_call_opt 2D buffers use a scaler-safe wait/reduce path instead
# of relying solely on reduce_gradients().
CI报告基于以下代码生成（30分钟更新一次）: 1 Required任务 : 47/48 通过
2 失败详情🔴 Coverage test — PR问题：Python Diff Coverage 未达标（置信度: 高）失败用例: Python Diff Coverage 门禁
关键日志:
PR 只修改了 新增路径包括 PP+ 修复建议:
关联变更: |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
复查了新提交：pipeline parallel 与 sharding comm_overlap 的互斥保护已经补上，这个问题看起来已解决。_comm_buffers 也已补充，但 AMP GradScaler 场景在 comm_group_call_opt=True 时仍可能跳过 2D Muon 梯度 reduce，我已在原有 scaler 线程里补充了新的证据和建议。
当前 CI 仍在运行中，尚未形成最终结果。
PaddlePaddle-bot
left a comment
PaddlePaddle-bot
left a comment
There was a problem hiding this comment.
🤖 Paddle-CI-Agent | pr_review |
2026-07-07 20:08:00
📋 Review 摘要
PR 概述：为 MuonShardingOptimizer 增加 sharding stage1 communication overlap hook，并暴露 comm buffer 以兼容 AMP GradScaler。
变更范围：python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py
影响面 Tag：Distributed Strategy Communication Library Performance Optimization
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🔴 Bug | python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py:352 |
AMP GradScaler 会跳过仍依赖同步路径的 Muon 2D 梯度通信 |
📝 PR 规范检查
标题缺少 Paddle PR 模板要求的 [Tag]；描述结构和精度变化字段符合规范。
标题建议（可直接复制）：
[Performance Optimization] Support Muon sharding optimizer comm overlap
总体评价
本 PR 的 1D overlap hook 方向与 DygraphShardingOptimizerV2 逻辑基本一致，但当前 AMP GradScaler 兼容分支会把部分 2D 同步通信路径绕开，需要先补齐或禁用对应组合后再合入。
| # the overlap hook drives. group_call_opt reduces 2D buffers | ||
| # synchronously in reduce_gradients, so they are excluded here. | ||
| self._comm_buffers = [] | ||
| if self._use_fuse_gradients and not self.use_group_call_opt: |
There was a problem hiding this comment.
🔴 Bug 这里把 use_group_call_opt 下的 2D comm buffer 排除在 _comm_buffers 外，但 fleet.scaler.distributed_scaler() 在 comm_overlap=True 时会先遍历 optimizer._comm_buffers 等待通信，然后把 HybridParallelOptimizer._sharding_enable 置为 False，导致后续 reduce_gradients() 整体被跳过。
因此在 AMP GradScaler + Muon comm overlap 下，只要 2D 梯度仍依赖 reduce_gradients() 的同步路径（例如 use_group_call_opt=True，以及 _use_fuse_gradients=False 时没有 2D buffer），这些 2D 参数不会执行 reduce 就进入 finite/unscale/step，各 sharding rank 会用未同步梯度更新。
建议修复方式：确保 scaler 关闭 _sharding_enable 前所有 Muon 2D 梯度也已完成通信。可选方案是对这类组合禁用/报错，或把 2D reduce 改成可由 _comm_buffers 驱动并在 scaler 中等待；不要让仍依赖 reduce_gradients() 的 2D 同步路径被 _sharding_enable=False 跳过。
Codecov Report❌ Patch coverage is
❌ Your patch status has failed because the patch coverage (50.00%) 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 #79428 +/- ##
==========================================
Coverage ? 50.00%
==========================================
Files ? 1
Lines ? 22
Branches ? 0
==========================================
Hits ? 11
Misses ? 11
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/re-run all-failed |
1 similar comment
|
/re-run all-failed |
CI报告基于以下代码生成（30分钟更新一次）: 1 Required任务 : 47/48 通过
2 失败详情🔴 Coverage test — PR问题：Python Diff Coverage 未达标（置信度: 高）分析器: 通用分析(fallback)
关键日志:
PR 只修改了 新增路径包括 PP+ 修复建议:
关联变更: |
4 participants
PR Category
Performance Optimization
PR Types
New features
Description
开发muon支持stage1 overlap的功能
是否引起精度变化
否