Supports Sonicmoe bf16 z and wgrad, optimize sonicmoe performance

[pkuzyc]

PR Category

Performance Optimization

PR Types

Performance

Description

Supoorts high-precision sonicmoe: PFCCLab/supersonic-moe#56, and add config in SonicMoEExpert to set the inside config in sonicmoe.

1. Fix the bugs when save_z_fp8=False
2. Supports recompute_z to save memory.
3. Supports BF16-wgrad.
4. Optimize sonicmoe performance, integrate [Performance] Integrate SonicMoE hotpath optimizations #1400

是否引起精度变化

否

[Paddle-CI-Bot]

PaddleFleet Log Analysis

Run #29107430427 · Attempt 1

日志分析报告

PR #1348 | branch: sonicmoe_bf16_z_wgrad | run attempt: 3 | commit: 4ababba

流水线名称	问题标签	修复建议	日志片段
Unit test (single card)	CI 超时（60min）	单测套件总耗时超过60分钟硬限制，最后一个正常完成的测试是 test_gated_attn_recompute.py，后续测试未来得及运行；需排查是否本次PR新增测试拉长了耗时，或已有测试在新代码下变慢	报错代码
Coverage Upload And Check	增量覆盖率不达标（20% < 90%）	在 fusion_layer_utils.py 和 moe_expert.py 中补充对应行的单测覆盖	报错代码

失败的测试case:

Unit test (single card):
  - 未发生测试用例级失败，所有已运行用例均 PASSED
  - 失败原因：整个 "Single card test" step 触发 60 分钟超时硬限制，
    最后执行至 tests/single_card_tests/transformer/test_gated_attn_recompute.py（1 passed），
    超时前仍有未执行的测试文件

Coverage Upload And Check:
  - src/paddlefleet/transformer/moe/fusion_layer_utils.py: 14.6% 覆盖率
      Missing lines: 51-52,55-56,96-97,100-101,105-106,109-110,115-116,119-126,130-132,138,141,3261-3263,3285,3308-3309,3313,3325
  - src/paddlefleet/transformer/moe/moe_expert.py: 21.4% 覆盖率
      Missing lines: 477-478,480-484,488-489,491-495,501-502,504,538-542,546,574,612,624-629,642-643
  - src/paddlefleet/transformer/moe/moe_layer.py: 100% ✅
  - 增量合计：85 行变更，68 行未覆盖，综合覆盖率 20%（阈值 90%）

根本原因分析:

本 PR (sonicmoe_bf16_z_wgrad) 改动集中在 MoE 层（fusion_layer_utils.py、moe_expert.py、moe_layer.py），新增了约 205 行变更代码，但仅为 moe_layer.py 补充了单测；fusion_layer_utils.py 和 moe_expert.py 中大量新增行（BF16 z-wgrad 相关分支路径）无单测覆盖，导致增量覆盖率仅 20%；与此同时，新增测试文件（如 test_gated_attn_recompute.py）加入测试套件后，整体运行时间超过了 60 分钟的流水线限制。

修复建议:

1. 覆盖率不达标：

   • 在 tests/single_card_tests/transformer/ 下为 fusion_layer_utils.py 新增测试，重点覆盖 lines 51-52、55-56、96-141（BF16 权重梯度分支）以及 lines 3261-3325（对应新增的融合算子路径）
   • 为 moe_expert.py 补充测试，覆盖 lines 477-504（新增的 z-wgrad 计算入口）和 lines 538-643（梯度累积/recompute 分支）
   • 参考 test_moe_layer_precision.py 中已有的 run_test_z_bf16_recompute_z 测试模式进行扩展

2. 单测超时：

   • 检查本次 PR 新增/改动的测试用例（尤其是 test_gated_attn_recompute.py 等新文件）的运行耗时
   • 若个别测试用例耗时过长，考虑拆分或添加 @pytest.mark.slow 标记并在超时场景下跳过
   • 覆盖率问题修复后，超时问题可能因测试用例数量增加而加剧，建议同步优化测试时长

---

🔍 准确性记录：请点击评论底部 😊 图标，选择 👍（准确）或 👎（有误），将自动记录到 CI 监控系统

_{🔄 每次 Re-run 后自动更新}

[liuhao2638]

发现需要先处理的阻塞问题，细节已放在行内评论中，主要涉及 FP8 权重缓存刷新和 allgather SonicMoE 路径兼容性。

• 优先级：P3 非行级：PR 标题/描述无法附到 diff 行。建议将标题调整为 [Performance] Support SonicMoE BF16 z/recompute_z and BF16 wgrad，并修正文案中的 Supoorts 拼写，便于维护者快速对应本次变更范围。

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

[liuhao2638]

优先级：P1

这里把 FP8 权重量化缓存的刷新条件收窄成只看 need_quant_weight()，训练中第一次 forward 后 weight.fp8 / weight.transposed_fp8 会一直存在；优化器更新 BF16 参数后，下一步 run_sonic_moe(..., use_fp8=True) 仍会使用旧缓存，实际参与训练的专家权重会落后一轮。当前 SonicMoEExpert.step() 只做 layout flush，没有清空这些缓存。请在训练态重新量化，或接入明确的 optimizer-step 缓存失效/版本检查，并补一个包含 optimizer.step() 后第二次 FP8 forward 的回归测试。一个保守的本地修复是：

Suggested change

	if self.sonic_moe_config.enabled is True and self.need_quant_weight():
	self.quant_weight()
	if self.sonic_moe_config.enabled is True and (
	self.training or self.need_quant_weight()
	):
	self.quant_weight()

[liuhao2638]

优先级：P1

这里绕过了上面 get_dispatched_routing() 的公共返回值，改成无条件读取 self.token_dispatcher._comm_manager；但 moe_token_dispatcher_type="allgather" 初始化的是 AllGatherTokenDispatcher，该类没有 _comm_manager，同时 _validate_allgather_config() 会强制 moe_use_fusion_node=True，所以 EP>1 allgather SonicMoE 会进入 fusion_moe_forward() 并在这里触发 AttributeError。请保留公共 API 返回的 dispatched_indices / dispatched_probs，不要对所有 dispatcher 访问私有 manager。

Suggested change

	dispatched_indices = (
	self.token_dispatcher._comm_manager.dispatched_indices
	)
	dispatched_probs = self.token_dispatcher._comm_manager.dispatched_probs
	fp8_combine_grad_handle = {} if self.fp8_dispatch_bwd else None

[adam-xiaoyao]

这里请记得改下。

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-10 20:57:32

📋 Review 摘要

PR 概述：为 SonicMoE 接入 BF16 z/recompute_z、BF16 wgrad 和部分 FP8/metadata 性能优化。
变更范围：src/paddlefleet/transformer/moe/、SonicMoE third-party 子模块和相关单/多卡测试。
影响面 Tag：MoE FP8 Third-Party

问题

未发现新的阻塞性问题。PR 规范问题在下面章节报，不在这里重复。

历史 Findings 修复情况

Finding	问题	状态
F1	FP8 权重缓存仅在缺失时量化	⚠️ 仍存在
F2	dispatcher 公共结果仍被 _comm_manager 直接读取	⚠️ 仍存在
F3	weight2.main_grad 布局转换仍被跳过	⚠️ 仍存在

📝 PR 规范检查

标题缺少官方 Tag；描述已包含 checklist §D2 要求的 PR Category、PR Types 和 Description 章节。

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

• [Performance] Support SonicMoE BF16 z/recompute_z and BF16 wgrad

总体评价

本轮按风险优先审查了 SonicMoE 权重布局/FP8 缓存、router-score metadata、combine backward FP8 handle 和新增测试；没有新增行间评论。历史问题仍需在合入前继续处理，尤其是 FP8 权重缓存和 weight2.main_grad 布局。

[codecov-commenter]

Codecov Report

❌ Patch coverage is 20.00000% with 68 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@01036b8). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
.../paddlefleet/transformer/moe/fusion_layer_utils.py	14.63%	35 Missing ⚠️
src/paddlefleet/transformer/moe/moe_expert.py	21.42%	33 Missing ⚠️

❌ Your patch status has failed because the patch coverage (20.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    #1348   +/-   ##
==========================================
  Coverage           ?   20.00%           
==========================================
  Files              ?        3           
  Lines              ?       85           
  Branches           ?       16           
==========================================
  Hits               ?       17           
  Misses             ?       68           
  Partials           ?        0           

Flag	Coverage Δ
coverage_combine	20.00% <20.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/paddlefleet/transformer/moe/moe_layer.py	100.00% <100.00%> (ø)
src/paddlefleet/transformer/moe/moe_expert.py	21.42% <21.42%> (ø)
.../paddlefleet/transformer/moe/fusion_layer_utils.py	14.63% <14.63%> (ø)

🚀 New features to boost your workflow:

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

[swgu98]

单卡单测已拉PR验证
https://github.com/PaddlePaddle/PaddleFleet/actions/runs/29107273463/job/86413989421?pr=1451

[swgu98]

修改内容只有b卡才能跑