Enable optional batched Muon NS iteration by xxyux · Pull Request #79482 · PaddlePaddle/Paddle

xxyux · 2026-07-15T08:49:56Z

PR Category

Performance Optimization

PR Types

Performance

Description

Controlled by the MUON_ENABLE_BATCH environment variable to enable group NS iteration.

When group NS iteration is enabled, parameters with identical shapes and slice functions are grouped together and processed in a single batch NS iteration, reducing optimizer computation time.

是否引起精度变化

否

paddle-bot · 2026-07-15T08:50:02Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

risemeup1111

已复查，当前还有两处需要修复后再合入的问题，细节都在 inline review comments 里。

优先级：P3 非行级：PR 标题 update group_NS 过于笼统，建议改成能直接体现变更的标题，例如 Enable optional batched Muon NS iteration，方便后续回溯和检索。处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

PaddlePaddle-bot · 2026-07-15T10:59:57Z

🤖 Paddle-CI-Agent | ci_status_monitor | 2026-07-17 10:37:21 UTC+08:00

CI报告基于以下代码生成（30分钟更新一次）:
PR commit: 964076b | Merge base: 206a21d (branch: develop)

1 Required任务 : 46/48 通过

当前 required 任务无失败，CI 仍在运行中：失败 0，运行中 2，等待中 0。

总执行（rerun次数）	总任务	✅ 通过	❌ 失败	⏳ 运行中	⏸️ 等待中	跳过
121(39)	82	79	1	2	0	0

任务	错误类型	置信度	日志
无	无	无	无

2 失败详情

无

risemeup1111

已复查新提交。分组 key 相关问题已修复；split_concat_func 的 batch 输入契约问题仍未完全解决，我已在原 inline 线程补充了当前代码和测试覆盖的证据。请继续修复后再提交新的 commit。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

codecov-commenter · 2026-07-15T15:34:41Z

Codecov Report

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

Files with missing lines	Patch %	Lines
python/paddle/optimizer/muon.py	93.50%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop   #79482   +/-   ##
==========================================
  Coverage           ?   93.50%           
==========================================
  Files              ?        1           
  Lines              ?       77           
  Branches           ?        0           
==========================================
  Hits               ?       72           
  Misses             ?        5           
  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.

risemeup1111

已复查最新提交。分组 key 问题和标题已处理，但 split_concat_func 的 batch 输入契约问题仍在；我已在原线程补充了具体复查说明。当前还需要继续修复后再合入。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

xxyux · 2026-07-16T07:31:03Z

/re-run all-failed

risemeup1111

已复查当前提交，仍有阻塞问题；详细说明已经写在 inline review comments 和原线程里。当前这版还不能合入。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

risemeup1111 · 2026-07-16T15:33:12Z

+        # Group parameters by split_concat_func and shape, then update
+        # each group in one batched call.
+        muon_groups = defaultdict(list)


优先级：P1
处理要求：请针对该评论修复并提交新的 commit。

这里把 Muon 参数无条件分组批处理了：上一版的 MUON_ENABLE_BATCH 开关和逐参数 fallback 已被删除，当前仓库里也没有任何 MUON_ENABLE_BATCH 读取，test_muon_sharding_batched_update_group 也随之移除。但 PR 描述仍说明 group NS iteration 由该环境变量控制，因此未设置环境变量的现有训练现在也会默认走新的 batch 路径，破坏原来的可回退行为，并会放大已有 split_concat_func 输入契约问题。

请恢复开关语义：默认继续逐参数更新，仅在 MUON_ENABLE_BATCH=1 时进入分组 batch 路径。实现形态可以类似：

if not MUON_ENABLE_BATCH: for param, grad in muon_params: self._muon_update_group([(param, grad)], ...) else: # existing grouping logic

目前只保留了_muon_update_group函数，删除了_muon_update函数。因为所有case都可以等价替换。不管有没有slice_func，shape是2D还是3D。

这个判断我不同意。当前实现里，多个同 shape 的参数会先被合并成 batched_matrix，再交给 split_concat_func；而仓库里的 _qkv_sep / _ffn_split 仍然是按单个 2D 矩阵、axis=1 写的，不能直接把 grouped batch 当成等价输入。

所以删掉 _muon_update 并不自动成立“所有 case 都可等价替换”。如果要保留这个结论，至少需要把带 split_concat_func 的 group 改成 batch-aware 接口，或者对这类 group 保留逐参数路径，并补上多个同 shape、同 slice 函数的回归测试。

这个判断我不同意。当前 head 里还有一个更直接的语义问题：lr_ratio/effective_lr 和最后的 subtract_ 已经被移到 for 外面，导致 group 里的前几个参数只做了前面的 decay，真正的 Muon step 只会落到最后一个参数上。这个状态下，哪怕不考虑 split_concat_func，_muon_update_group 也不能和逐参数路径视为等价。

我已经在对应代码行补了 inline 评论，请先把这个循环结构修正，再讨论是否可以删除 _muon_update。split_concat_func 的 2D 契约问题也仍然需要单独处理。

PaddlePaddle-bot

🤖 Paddle-CI-Agent | pr_review | 2026-07-16 23:33:38

📋 Review 摘要

PR 概述：为 Muon optimizer 引入按形状和 slice 函数聚合的 batched NS 更新路径。
变更范围：python/paddle/optimizer/muon.py、Muon sharding 分布式测试。
影响面 Tag：Performance Optimization Distributed Strategy User Experience

问题

级别	文件	概述
🔴 Bug	`python/paddle/optimizer/muon.py:728`	`MUON_ENABLE_BATCH` 未接入，分组 NS 无条件启用

历史 Findings 修复情况

Finding	问题	状态
F1	`split_concat_func` 接收分组后的 `batched_matrix`	⚠️ 仍存在
F2	测试未对比 batch path 与逐参数旧路径	⚠️ 仍存在

📝 PR 规范检查

标题缺少官方 Tag，描述结构和“是否引起精度变化”字段符合模板。

标题建议（可直接复制）：

[Performance Optimization] Enable batched Muon group NS update

总体评价

当前分组实现和测试都已扩展，但描述里的 opt-in 开关没有落到代码，导致默认行为被直接切换；这个需要先补上 gating 和对应回归测试。历史 F1/F2 我检查后仍未看到修复证据，已在下方状态块标记。

PaddlePaddle-bot · 2026-07-16T15:35:14Z

+
+        # Group parameters by split_concat_func and shape, then update
+        # each group in one batched call.
+        muon_groups = defaultdict(list)


🔴 Bug MUON_ENABLE_BATCH 没有接入，分组 NS 现在会无条件启用。

PR 描述说 group NS iteration 由 MUON_ENABLE_BATCH 控制，但当前文件全仓搜索不到这个环境变量；这里总是构建 muon_groups，随后所有 Muon 参数都会进入 _muon_update_group。这会把本应可选的优化变成默认行为，用户无法关闭新路径，也无法在 split 函数或数值回归时回退到逐参数更新。

建议修复方式：在模块加载时读取 MUON_ENABLE_BATCH，只有为 1 时才按 (split_concat_func, shape) 聚合；未开启时按 muon_params 顺序逐个调用 _muon_update_group([(param, grad)], ...)，并补一组 env 未开启时不聚合的测试。

目前只保留了_muon_update_group函数，删除了_muon_update函数。因为所有case都可以等价替换。不管有没有slice_func，shape是2D还是3D。

xxyux · 2026-07-17T02:28:27Z

/re-run all-failed

sneaxiy

LGTM

xxyux · 2026-07-17T04:43:03Z

/re-run all-failed

risemeup1111

已复查最新提交。当前仍有阻塞问题，新的代码行级问题已写在 inline review comment，关于 _muon_update_group 等价性的回复也已补到原线程。请继续修复后再提交新的 commit。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

risemeup1111 · 2026-07-17T11:02:46Z

+
+                final_step = orthogonal_update * lr

            lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param)


优先级：P1
处理要求：请针对该评论修复并提交新的 commit。

这里 lr_ratio/effective_lr 和后面的减步逻辑已经跑到 for i, (param, _) in enumerate(group_params_grads): 外面了。按当前缩进，循环里的每个参数只会执行前面那段基于 lr 的 decay / final_step，而真正带 effective_lr 的更新只会作用到最后一个参数；这会让 group 内参数更新不一致，甚至把最后一个参数更新两次。即使不考虑 split_concat_func，这也不是和逐参数路径等价的实现。

请把 lr_ratio/effective_lr、decay 和 subtract_ 都放回每个参数的循环内部，形态应类似：

for i, (param, _) in enumerate(group_params_grads): orthogonal_update = orthogonal_updates[i] find_master = find_master_list[i] master_weight = self._master_weights[param.name] if find_master else None lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param) effective_lr = lr * lr_ratio if with_decay and weight_decay > 0: # scale current param/master_weight with effective_lr ... final_step = orthogonal_update * effective_lr # subtract final_step from the current param/master_weight ...

已经修复，这个问题是#79466 引入的，解冲突错误导致。

已复查当前 head，这条已经修复：lr_ratio/effective_lr、weight decay 和最终 subtract_ 都已经回到 for i, (param, _) in enumerate(group_params_grads): 循环内部，因此本线程的问题不再阻塞。

剩余阻塞项仍是另外两个独立问题：split_concat_func 的 batch 输入契约，以及 MUON_ENABLE_BATCH 开关未接入；我已在对应线程和最新 review 里说明。

risemeup1111 · 2026-07-17T11:03:30Z

+        # Group parameters by split_concat_func and shape, then update
+        # each group in one batched call.
+        muon_groups = defaultdict(list)


这个判断我不同意。当前 head 里还有一个更直接的语义问题：lr_ratio/effective_lr 和最后的 subtract_ 已经被移到 for 外面，导致 group 里的前几个参数只做了前面的 decay，真正的 Muon step 只会落到最后一个参数上。这个状态下，哪怕不考虑 split_concat_func，_muon_update_group 也不能和逐参数路径视为等价。

我已经在对应代码行补了 inline 评论，请先把这个循环结构修正，再讨论是否可以删除 _muon_update。split_concat_func 的 2D 契约问题也仍然需要单独处理。

risemeup1111

已复查最新提交。lr_ratio/effective_lr 的循环位置问题已经修复；没有新增 inline 评论，我已在已有 split_concat_func 线程补充当前 head 的证据。

当前仍有两个 P1 阻塞项未完成：split_concat_func 仍直接接收 batched_matrix，并且 MUON_ENABLE_BATCH 仍未接入、分组路径仍默认启用。请按已有线程修复并提交新的 commit 后再复查。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

xxyux · 2026-07-18T09:48:12Z

/re-run all-failed

xxyux · 2026-07-19T15:32:24Z

/re-run all-failed

xxyux · 2026-07-20T02:29:29Z

/re-run all-failed

This comment was marked as outdated.

Sign in to view

risemeup1111 suggested changes Jul 15, 2026

View reviewed changes

Comment thread python/paddle/optimizer/muon.py Outdated

Comment thread python/paddle/optimizer/muon.py

xxyux changed the title ~~update group_NS~~ Enable optional batched Muon NS iteration Jul 15, 2026

xxyux force-pushed the batch-compute-muon branch from 8a6e662 to 756ed0b Compare July 15, 2026 11:47

risemeup1111 suggested changes Jul 15, 2026

View reviewed changes

This comment was marked as outdated.

Sign in to view

xxyux force-pushed the batch-compute-muon branch from 756ed0b to c4227a4 Compare July 16, 2026 03:24

This comment was marked as outdated.

Sign in to view

risemeup1111 suggested changes Jul 16, 2026

View reviewed changes

xxyux force-pushed the batch-compute-muon branch from c4227a4 to 964076b Compare July 16, 2026 15:29

risemeup1111 suggested changes Jul 16, 2026

View reviewed changes

PaddlePaddle-bot suggested changes Jul 16, 2026

View reviewed changes

GuoxiaWang previously approved these changes Jul 17, 2026

View reviewed changes

sneaxiy previously approved these changes Jul 17, 2026

View reviewed changes

xxyux mentioned this pull request Jul 17, 2026

Enable optional batched Muon NS iteration #79495

Merged

xxyux dismissed stale reviews from sneaxiy and GuoxiaWang via 36b2604 July 17, 2026 10:56

risemeup1111 suggested changes Jul 17, 2026

View reviewed changes

update group_NS

f447b4f

xxyux force-pushed the batch-compute-muon branch from 36b2604 to f447b4f Compare July 18, 2026 06:54

risemeup1111 suggested changes Jul 18, 2026

View reviewed changes

sneaxiy approved these changes Jul 20, 2026

View reviewed changes

sneaxiy merged commit 80cc5d8 into PaddlePaddle:develop Jul 20, 2026
123 of 136 checks passed


		final_step = orthogonal_update * lr

		lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param)

Uh oh!

Conversation

xxyux commented Jul 15, 2026

PR Category

PR Types

Description

是否引起精度变化

Uh oh!

paddle-bot Bot commented Jul 15, 2026

Uh oh!

This comment was marked as outdated.

Uh oh!

risemeup1111 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

PaddlePaddle-bot commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

1 Required任务 : 46/48 通过

2 失败详情

Uh oh!

risemeup1111 left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

codecov-commenter commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This comment was marked as outdated.

Uh oh!

risemeup1111 left a comment

Choose a reason for hiding this comment

Uh oh!

xxyux commented Jul 16, 2026

Uh oh!

risemeup1111 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

PaddlePaddle-bot left a comment

Choose a reason for hiding this comment

📋 Review 摘要

问题

历史 Findings 修复情况

📝 PR 规范检查

总体评价

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xxyux commented Jul 17, 2026

Uh oh!

sneaxiy left a comment

Choose a reason for hiding this comment

Uh oh!

xxyux commented Jul 17, 2026

Uh oh!

risemeup1111 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PaddlePaddle-bot commented Jul 15, 2026 •

edited

Loading

codecov-commenter commented Jul 15, 2026 •

edited

Loading