support high_precision_norm in swa

[wangyuwen1999]

PR Category

Execute Infrastructure

PR Types

New features

Description

新增swa 中norm高精模式，支持qkv进行norm与rope的计算过程中以fp32进行，以支持训推一致。

新增开关：

• swa_high_precision_norm：在QKV的Norm过程中，计算应用fp32，返回时重新cast回原有精度，如bf16。

RMS_Norm forward函数新增入参

• high_precision_norm：RMS_Norm计算过程在fp32下进行
• return_high_precision_norm：RMS_Norm返回的结果是否为fp32。

备注
当swa_high_precision_norm开启，kv的rms_norm将在fp32下进行，并返回kv为fp32

是否引起精度变化

是

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

root seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Paddle-CI-Bot]

PaddleFleet Log Analysis

Run #28454452326 · Attempt 1

日志分析报告

流水线名称	问题标签	修复建议	日志片段
Unit test (multi-card)	测试超时 hang（test_csa_attention_cp.py）	PR 引入的 dsv4_hybrid_attention.py 修改影响 CSA attention + CP 的 8 卡并行路径，导致进程 hang 超出 60 min 超时；需在本地复现 test_csa_attention_cp.py，确认 CP=8 下 attention 计算是否死锁或无限等待	报错代码
Coverage Upload And Check	增量覆盖率不达标（exit 9）	PR 改动 68 行，diff coverage 仅 86%（阈值 90%），需补充 dsv4_hybrid_attention.py L531/L612 和 paddle_norm.py L86/L98 的单测覆盖	报错代码

失败的测试case:

tests/multi_card_tests/transformer/test_csa_attention_cp.py  — 8 GPUs hang，CI 60min 超时强杀

根本原因分析:
PR fix_attn_precision 修改了 dsv4_hybrid_attention.py（36 行）、paddle_norm.py（19 行）、transformer_config.py（13 行），其中 attention precision 相关改动在 CP=8 的 test_csa_attention_cp 场景下触发进程 hang（nvidia-smi 连续超时、主进程长时间无输出），同时新增代码未覆盖 4 个分支行，导致增量覆盖率跌至 86%。

修复建议:

1. 定位 hang：在本地执行 python -m paddle.distributed.launch --gpus 0,1,2,3,4,5,6,7 tests/multi_card_tests/transformer/test_csa_attention_cp.py，重现后用 py-spy dump --pid <rank0_pid> 或 gdb 查看卡在哪个 NCCL/attention kernel 调用。
2. 排查 attention precision 修改：检查 dsv4_hybrid_attention.py 中 precision 转换（如 fp32 ↔ bf16 cast）是否在 CP ring-allreduce 通信前后引入了数据类型不一致，导致 collective 挂起；重点关注 W0630 ... Got different data type, run type promotion automatically 警告与 hang 的关联。
3. 补齐覆盖率：在现有单测中增加 dsv4_hybrid_attention.py L531、L612 和 paddle_norm.py L86、L98 对应分支的测试用例，确保 diff coverage ≥ 90%。

---

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

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

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-06-30 21:29:14

📋 Review 摘要

PR 概述：新增 DSv4 SWA 高精度 RMSNorm 开关，并扩展 RMSNorm/DSv4 RoPE 路径。
变更范围：src/paddlefleet/transformer/、DSv4 attention、RMSNorm、TransformerConfig、相关单测
影响面 Tag：Transformer TransformerConfig Tests

问题

级别	文件	概述
🔴 Bug	src/paddlefleet/transformer/transformer_config.py:1236	swa_high_precision_norm 未限制 norm 实现，合法 LayerNorm 配置会在 DSv4 KV path 传参时报错
🟡 建议	tests/single_card_tests/ai_edited_test/transformer/test_ai_swa_high_precision_norm.py:29	新测试手动改 sys.path，容易导入非当前 PR 代码
🟡 建议	tests/single_card_tests/ai_edited_test/transformer/test_ai_swa_high_precision_norm.py:65	高精度 norm 测试只断言 dtype/shape，不能验证 fp32 计算语义

历史 Findings 修复情况

Finding	问题	状态
F1	新增的 swa_high_precision_norm 只落到了 dataclass，命令行入口没有同步。	⚠️ 仍存在
F2	swa_high_precision_norm 没有覆盖 Q path 里已有的 q_layernorm。	⚠️ 仍存在
F3	high_precision_norm=True 但 return_high_precision_norm=False 时没有恢复输入 dtype。	⚠️ 仍存在

📝 PR 规范检查

标题缺少官方 Tag；描述结构符合 checklist §D2。

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

• [New features] Support high precision norm in SWA

PR 描述建议（点击展开，可直接复制）

### PR Category
<!-- One of [ User Experience | Execute Infrastructure | Operator Mechanism | Custom Device | Performance Optimization | Distributed Strategy | Parameter Server | Communication Library  | Environment Adaptation ] -->
Execute Infrastructure
### PR Types
<!-- One of [ New features | Bug fixes | Improvements | Performance | BC Breaking | Deprecations | Docs | Devs | Not User Facing | Security | Others ] -->
New features
### Description
<!-- Describe what you’ve done -->
新增 `swa_high_precision_norm` 开关：在 DSv4 Hybrid Attention 的 Q/KV RMSNorm 计算中使用 fp32，并在 KV RoPE 前保留 fp32 结果以支持高精度 RoPE；同时为 DSv4 非融合 RoPE 路径传递 `high_precision_rope`。本 PR 还扩展 `RMSNorm.forward`，增加 `high_precision_norm` 和 `return_high_precision_norm` 参数。

总体评价

当前实现补上了部分 DSv4/KV 高精度路径，但新增开关还缺少对 norm 实现的约束，存在可触发的运行时崩溃组合。历史三条 finding 在当前 diff 中仍未完全修复，建议和本轮新增问题一起处理后再合入。

[PaddlePaddle-bot]

🔴 Bug swa_high_precision_norm 当前只校验了 DSv4 variant，但没有限制 norm 实现。DSv4 layer spec 会按 config.normalization 构造 q_layernorm/kv_layernorm；当用户设置 normalization="LayerNorm" 时，kv_layernorm 是 LayerNorm.forward(self, hidden_states)，而 get_query_key_value_tensors() 在该开关下会传 high_precision_norm/return_high_precision_norm，运行时会报 unexpected keyword argument。

建议修复方式：
在这里同时要求 normalization == "RMSNorm"，或给 LayerNorm/所有可作为 kv_layernorm 的实现补齐相同 kwargs 语义；按 PR 描述这是 RMSNorm 专项，优先在 config 校验阶段拒绝非 RMSNorm。

[PaddlePaddle-bot]

🟡 建议 新增测试里手动 sys.path.insert 会绕过仓库的测试入口依赖。这个路径从当前文件只回到 tests/ 目录，若工作目录/PYTHONPATH 不同，可能导入到环境里已安装的 paddlefleet 而不是本 PR 的 src/paddlefleet；文件末尾的 unittest.main() 也属于同一类本地运行入口依赖。

建议修复策略：
删除 sys.path.insert 和脚本入口，按仓库测试配置提供 PYTHONPATH=src/editable install 后直接 import，或将公共测试路径配置放到测试 runner/conftest。

[PaddlePaddle-bot]

🟡 建议 这个用例只断言 dtype/shape，不能证明 high_precision_norm=True 真的按 fp32 算；即使分支里没有 q.float()、或计算顺序退回 bf16，测试仍会通过。对 precision feature 需要数值对齐保护，否则当前 Q/KV 路径的精度回归很容易被漏掉。

建议修复策略：
构造一组 bf16 输入，分别计算 fp32 reference 和低精度 reference，断言 high_precision 分支更接近 fp32 reference；同时覆盖 RMSNorm(..., high_precision_norm=True) 和 DSv4 get_query_key_value_tensors() 的实际调用路径。

[codecov-commenter]

Codecov Report

❌ Patch coverage is 76.66667% with 7 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (release/0.3@3b32e04). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...c/paddlefleet/transformer/dsv4_hybrid_attention.py	73.33%	2 Missing and 2 partials ⚠️
src/paddlefleet/transformer/paddle_norm.py	75.00%	2 Missing and 1 partial ⚠️

❌ Your patch status has failed because the patch coverage (76.66%) 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               @@
##             release/0.3    #1333   +/-   ##
==============================================
  Coverage               ?   76.66%           
==============================================
  Files                  ?        3           
  Lines                  ?       30           
  Branches               ?        9           
==============================================
  Hits                   ?       23           
  Misses                 ?        4           
  Partials               ?        3           

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

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

Files with missing lines	Coverage Δ
src/paddlefleet/transformer/transformer_config.py	100.00% <100.00%> (ø)
src/paddlefleet/transformer/paddle_norm.py	75.00% <75.00%> (ø)
...c/paddlefleet/transformer/dsv4_hybrid_attention.py	73.33% <73.33%> (ø)

🚀 New features to boost your workflow:

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