Reuse MoETopkFusion in quantile-balancing routing - #1775
Conversation
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 状态 |
|---|---|---|---|
| 1 | moe_topk_fusion.py:L167 discussion |  |
✅ |
| # After the top-k loop, choice_vals has the top-k positions masked to -inf. | ||
| # The max of the remaining values is the (k+1)-th largest = cutoff alpha. | ||
| if return_alpha: | ||
| alpha_val = tl.max(choice_vals, axis=0) |
There was a problem hiding this comment.
moe_k 次选择后才计算 cutoff；当合法配置 moe_k == n_experts 时，所有候选都已被置为 -inf，alpha 会变成 -inf。旧路径对 min(k + 1, E) 取第 E 个值，得到的是最小 biased score；新值会让 _accumulate_qb_histogram 的 required_bias 全部落入边界 bin，后续 QB callback 更新出错误 bias。请在屏蔽最后一个专家前保存最小值，或对 moe_k >= n_experts 特判使用原始候选的最小值。
There was a problem hiding this comment.
触发条件是 moe_k == n_experts（或 node-limit 下 moe_k == 选中组的候选总数）。这意味着每个 token 路由到全部专家——那是 dense，不是 MoE，而 QB 的唯一目的就是均衡路由负载，全选的时候负载均衡根本没意义。DeepSeek/Kimi 这类配置永远是 moe_k << n_experts（比如 8/256），node-limit 后候选池也还有上百个（topk_group*epg），远大于 moe_k。所以这个分支在任何真实 QB 训练里都不会走到。
moe_k < n_experts 且无 node-limit 时循环后必留有限值，alpha 不可能是 -inf，没有问题。
There was a problem hiding this comment.
你指出的 moe_k == n_experts 全选边界对实际 QB 配置确实不构成问题，我撤回这部分理由。但当前实现仍绕过了 QB 原有的 n_group > 1 拒绝：forward 在 n_group > 1 时把 use_node_limit=True 传给 fusion kernel，kernel 返回的是“选中组内”的 cutoff，而 _accumulate_qb_histogram 把它当成全 E 专家的 (k+1)-th cutoff 使用。这样 histogram 与原 QB 定义不一致，callback 会据此更新出错误的 per-expert bias。请保持 QB fusion 限制为 n_group == 1，或同步实现并验证 node-limit 下的 histogram 语义；在此修复前我保留该 finding。
There was a problem hiding this comment.
感谢 review。这条 finding 描述的状态——「QB + n_group > 1 走 fusion 路径」——在当前实现下无法构造,因此不会发生,理由如下:
- QB 在 init 阶段硬拒绝 n_group > 1,本 PR 未改动该检查。
StandardMoERouter.init 里的守卫一直在:
if self.topk_method == "quantile_balancing":
...
if self.n_group != 1:
raise ValueError(
"Quantile Balancing currently only supports n_group=1. "
"Multi-group routing (n_group>1) is not compatible with QB ..."
)
本 PR 删除的只是另一处「quantile_balancing 与 moe_topk_fusion 不兼容」的检查,这个 n_group 守卫一行都没动(可对照 diff 里对 n_group 的改动,仅出现在 forward 的 use_node_limit = self.n_group > 1 与 kernel 传参处)。也就是说,任何 n_group > 1 的 QB router 在构造期就会抛 ValueError,根本进不到 forward。
- 因此 fusion 路径上 use_node_limit 对 QB 恒为 False。
self.n_group = config.n_group,而 QB 已被强制 n_group == 1,故:
use_node_limit = self.n_group > 1 # QB 分支恒为 False
kernel 不会执行任何 node-limit 组内屏蔽,返回的 alpha 就是全 E 个专家的第 (k+1) 大 cutoff——正是 _accumulate_qb_histogram 所期望的定义。你描述的「kernel 返回选中组内 cutoff、histogram 按全 E 解释」这一步,在 QB 上无法被触发。
- 该不变量已有守卫冗余 + 测试兜底。
第二道守卫在 _topk_quantile_balancing 内,对 n_group != 1 同样抛错。
单测 test_n_group_greater_than_one_rejected 已断言 _build_qb_router(n_group=2) 抛出含 "only supports n_group=1" 的 ValueError。
综上,「保持 QB fusion 限制为 n_group == 1」这一诉求已经被满足,并有测试保证,不是隐式约定。node-limit 下的 histogram 语义无需在本 PR 实现,因为该组合被显式禁止。
为消除「守卫与 fusion 代码有距离」的顾虑,我在 fusion 分支旁补了一行注释,显式标注该不变量,不改变任何行为。
There was a problem hiding this comment.
已核对 StandardMoERouter.__init__ 的守卫和 test_n_group_greater_than_one_rejected：QB + n_group > 1 确实在构造阶段被拒绝，当前 head 不会进入 node-limit fusion。结合你对全选边界的说明，原 finding 的两个触发路径均不成立，我撤回该 finding；当前 head 未发现新的可阻塞问题。
Codecov Report❌ Patch coverage is
❌ Your patch status has failed because the patch coverage (83.33%) 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 #1775 +/- ##
==========================================
Coverage ? 83.33%
==========================================
Files ? 1
Lines ? 12
Branches ? 6
==========================================
Hits ? 10
Misses ? 0
Partials ? 2
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
4 participants
PR Category
User Experience
PR Types
Others
Description
Reuse MoETopkFusion in quantile-balancing routing
是否引起精度变化
否