Align Kimi BF16 column linear under accuracy gate - #1411
Conversation
There was a problem hiding this comment.
Pull request overview
该 PR 旨在为 Kimi 的跨框架精度对齐引入一个显式的 use_accuracy_compatible 开关：在不影响默认高性能路径（如 fused / fusion 相关路径）的前提下，为特定 BF16/TP=1/MLA/MoE 形状选择更贴近 Megatron-Core（no-TE）数值行为的实现，以便进行 bitwise/严格精度对齐验证。
Changes:
- 在
TransformerConfig中新增use_accuracy_compatible开关，并在开启时统一调整一批 fused/rope/softmax 等相关开关与 env flag。 - 为 RMSNorm、attention(QK/softmax/mask)、MoE router/单卡 MoE combine、MLP SwiGLU、Embedding dtype 等关键算子引入 accuracy gate 下的数值对齐路径。
- 为 TP 线性层相关实现补充 accuracy gate 下的特定 BF16 路径（包含 TP=1 ColumnParallelLinear 的
F.linear路由与 Kimi routed expert fc1 的 split GEMM PyLayer）。
PR 标题/描述检查：
- 标题未符合约定格式（需要形如
[CLASS]Title）。建议改为例如：[Feature] Align Kimi BF16 column linear under accuracy gate（或按团队分类用[BugFix]/[Refactor]等）。 - 描述信息较完整（包含 Summary/Validation/Scope），满足最低要求。
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/paddlefleet/transformer/transformer_config.py | 新增 use_accuracy_compatible 并在开启时统一设置对齐用的配置与 env flag。 |
| src/paddlefleet/transformer/paddle_norm.py | accuracy gate 下引入显式 FP32 岛与 BF16 边界行为的 RMSNorm PyLayer。 |
| src/paddlefleet/transformer/multi_latent_attention.py | 在实验路径下对 accuracy gate 做 RoPE 分流，确保 Kimi 对齐走标准 MLA YaRN 路径。 |
| src/paddlefleet/transformer/moe/moe_utils.py | 将 env flag 读取改为动态读取，支持 config 构造阶段设置。 |
| src/paddlefleet/transformer/moe/moe_router.py | accuracy gate 下修正 unfused router matmul 权重布局、topk gather 以匹配 Torch scatter backward 等。 |
| src/paddlefleet/transformer/moe/moe_layer.py | accuracy gate 下调整单卡 MoE 的 per-token scale 位置与 combine 的 BF16 累加行为。 |
| src/paddlefleet/transformer/mlp.py | accuracy gate 下对 SwiGLU 使用 FP32 岛并支持 per_token_scale 的对齐位置。 |
| src/paddlefleet/transformer/hyper_connection.py | 动态读取 env flag，避免 import-time 固化。 |
| src/paddlefleet/transformer/dot_product_attention.py | accuracy/eager 下引入显式 QK PyLayer backward、softmax/mask dtype 处理以贴近对齐目标。 |
| src/paddlefleet/tensor_parallel/layers.py | accuracy gate 下为线性层引入 F.linear 分流与 Kimi expert fc1 split GEMM PyLayer 等。 |
| src/paddlefleet/models/gpt/gpt_layer_specs.py | accuracy gate 下调整 MLA 子层 spec（q_a/kv_a 投影选择）。 |
| src/paddlefleet/models/gpt/gpt_embedding.py | accuracy gate 下将 embedding 输出 cast 回表 dtype（对齐 Megatron 输入 RMSNorm 的 BF16 行为）。 |
| src/paddlefleet/fusions/fused_softmax.py | accuracy gate 下提供 FP32 softmax + BF16 输出且显式 backward 的 PyLayer。 |
| packages/paddlefleet_ops/src/paddlefleet_ops/init.py | 改善 ecosystem ops 动态加载失败时的封装与可用性 flag 处理。 |
| if bias is not None: | ||
| output = output + bias | ||
| output = paddle.nn.functional.linear(input, weight, bias) | ||
| elif use_accuracy_compatible: | ||
| output = paddle.nn.functional.linear(input, weight) | ||
| else: |
| mla_q_a_proj = ( | ||
| backend.linear() | ||
| if use_accuracy_compatible | ||
| else backend.column_parallel_linear() | ||
| ) | ||
| mla_kv_a_proj = ( | ||
| backend.linear() | ||
| if use_accuracy_compatible | ||
| else backend.column_parallel_linear() | ||
| ) |
| def forward( | ||
| ctx, | ||
| input, | ||
| weight, | ||
| bias, | ||
| allreduce_dgrad, | ||
| tp_group, | ||
| use_accuracy_compatible=False, |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已完成首轮审查。发现需要先修复的 P1 问题，细节已放在行内评论，主要涉及 accuracy gate 的进程级状态泄漏，以及 frozen-weight 线性层默认路径和反向 dtype 处理的回归风险。CI 当前通过。
- 非行级：GitHub mergeability 当前仍显示存在 merge conflicts；请同步 develop 并解决冲突后再请求复查。
| self.moe_topk_fusion = False | ||
| self.routing_map_fusion = False | ||
| self.fused_linear_ce_loss_chunk = 0 | ||
| os.environ["FLAGS_use_accuracy_compatible_kernel"] = "1" |
There was a problem hiding this comment.
这里在构造一个 use_accuracy_compatible=True 的 config 时会永久写入进程级环境变量。这个 flag 现在被多个运行时路径动态读取（例如 fused softmax、language loss、MoE dispatcher/hyper_connection），因此同一 Python 进程里只要先创建过 accuracy config，之后创建的普通 config 即使 use_accuracy_compatible=False 也会被这些 env-gated 路径污染，默认训练/推理路径不再可恢复。请不要在 TransformerConfig.__post_init__ 里修改 os.environ，而是把 config gate 显式传到需要的模块；如果还要兼容外部手动设置 env，可以把 env 只作为用户显式开关读取。
可以按这个方向收敛：
def accuracy_compatible_enabled(config=None):
return bool(getattr(config, "use_accuracy_compatible", False)) or (
os.environ.get("FLAGS_use_accuracy_compatible_kernel", "0") == "1"
)然后由持有 config 的调用点使用该 helper，而不是从 config 构造阶段反向写全局环境变量。
| if bias is not None: | ||
| output = output + bias | ||
| output = paddle.nn.functional.linear(input, weight, bias) | ||
| elif use_accuracy_compatible: | ||
| output = paddle.nn.functional.linear(input, weight) | ||
| else: | ||
| output = paddle.matmul(input, weight) |
There was a problem hiding this comment.
这里把 frozen-weight 的 bias 分支无条件切到 F.linear，即使 use_accuracy_compatible=False 也会改变默认路径的 GEMM/融合选择；这和 PR 描述的“default fused_linear path unchanged unless use_accuracy_compatible is enabled”不一致，也会让普通 BF16/FP16 冻结权重层出现数值或性能回归。请只在 accuracy gate 下使用 F.linear，默认路径保持原来的 matmul + bias。
| if bias is not None: | |
| output = output + bias | |
| output = paddle.nn.functional.linear(input, weight, bias) | |
| elif use_accuracy_compatible: | |
| output = paddle.nn.functional.linear(input, weight) | |
| else: | |
| output = paddle.matmul(input, weight) | |
| if use_accuracy_compatible: | |
| output = paddle.nn.functional.linear(input, weight, bias) | |
| else: | |
| output = paddle.matmul(input, weight) | |
| if bias is not None: | |
| output = output + bias |
| bias, | ||
| allreduce_dgrad, | ||
| tp_group, | ||
| use_accuracy_compatible=False, |
There was a problem hiding this comment.
这个新参数让 frozen-weight 前向也能走 accuracy F.linear，但同一个类的 backward 仍直接使用保存的 weight.t() 做 dgrad。PR 在 trainable linear backward 已经处理了同类问题：BF16 前向接 CE/上游 FP32 梯度时，grad_output 可能是 float32 而冻结权重仍是 bfloat16，Paddle matmul 需要同 dtype，会在 LoRA/冻结 base weight 路径反向时报 dtype mismatch。请同步修正 frozen-weight backward：
dgrad_weight = (
weight.cast(grad_output.dtype)
if weight.dtype != grad_output.dtype
else weight
)
grad_input = grad_output.matmul(dgrad_weight.t())
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
PaddleFleet Log Analysis
日志分析报告
失败的测试 case: 根本原因分析: 本 PR（
修复建议:
🔄 每次 Re-run 后自动更新 |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。合并冲突已经解除；上一轮 3 个 P1 代码问题在当前代码中仍未修复，继续以原行内线程为准。另发现当前 codestyle 失败中有两个可直接定位的格式/ruff 问题，已补充为新的行内评论。
- 非行级：PR 标题/描述不在 diff 行中，无法行内附着。当前
Check PR Description仍因模板字段缺失失败；建议把标题改成类似[Bug fixes] Align Kimi BF16 column linear under accuracy gate，并在描述里补齐PR Category: Distributed Strategy、PR Types: Bug fixes（或更准确的类型）以及现有 Summary/Validation 内容。
| if per_token_scale is not None: | ||
| scale = per_token_scale.unsqueeze(-1).cast("float32") | ||
| intermediate_parallel = intermediate_parallel * scale | ||
| intermediate_parallel = intermediate_parallel.cast(original_dtype) |
There was a problem hiding this comment.
Codestyle Check 里的 ruff format 仍在失败，本地 ruff format --diff src/paddlefleet/transformer/mlp.py 定位到这一行需要拆行。当前格式会让必需的样式检查继续失败；请按 formatter 输出调整。
| intermediate_parallel = intermediate_parallel.cast(original_dtype) | |
| intermediate_parallel = intermediate_parallel.cast( | |
| original_dtype | |
| ) |
| def _use_accuracy_compatible_config(config: "TransformerConfig") -> bool: | ||
| return bool(getattr(config, "use_accuracy_compatible", False)) |
There was a problem hiding this comment.
Codestyle Check 里 ruff check 失败，本地 ruff check --fix --diff src/paddlefleet/transformer/dot_product_attention.py 定位到这里的 UP037。该文件已经启用 from __future__ import annotations，带引号的类型注解会让样式检查持续失败；请移除注解里的字符串引号。
| def _use_accuracy_compatible_config(config: "TransformerConfig") -> bool: | |
| return bool(getattr(config, "use_accuracy_compatible", False)) | |
| def _use_accuracy_compatible_config(config: TransformerConfig) -> bool: | |
| return bool(getattr(config, "use_accuracy_compatible", False)) |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
a07ffb4 to
20a36d5
Compare
|
Update pushed in
|
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。当前 PR 已重置为最小范围，只保留 ColumnParallelLinear experimental 路径下的 BF16 TP=1 accuracy gate 和对应单卡回归测试；上一轮 Nyanpasu 提出的 P1/P3 问题都已随旧的大范围 diff 移除或由当前模板检查通过而关闭。当前代码未发现需要阻塞合入的问题。
本地已验证 py_compile 和 ruff check 通过；新增测试在本地因环境缺少 paddle 无法收集运行，仍以仓库 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-08 15:24:44
📋 Review 摘要
PR 概述：将 Kimi BF16 column-linear 的 accuracy-compatible 分支收敛到 experimental ColumnParallelLinear 的 TP=1 路径，并补充单卡门禁测试。
变更范围：src/paddlefleet/tensor_parallel/layers.py、tests/single_card_tests/tensor_parallel/test_column_sequence_parallel_accuracy.py
影响面 Tag：Tensor Parallel Tests
问题
未发现阻塞性问题。PR 规范问题在下面章节报，不在这里重复。
历史 Findings 修复情况
| Finding | 问题 | 状态 |
|---|---|---|
| F1 | use_accuracy_compatible 会把进程级 FLAGS_use_accuracy_compatible_kernel 永久置为 1。 |
✅ 已修复 |
| F2 | use_accuracy_compatible 只把 bf16 置为 True，但不会同步 params_dtype/autocast_dtype。 |
✅ 已修复 |
| F3 | use_accuracy_compatible=True 的默认配置会在这里提前走 EC flash attention，绕过后面的 accuracy-compatible QK/softmax 路径。 |
|
| F4 | BF16 accuracy attention backward 会把 FP32 softmax 梯度和 BF16 key/query 直接做 matmul。 |
📝 PR 规范检查
标题仍缺少官方 Tag；描述已按 checklist §D2 补齐 PR Category / PR Types / Description 结构。
标题建议（可直接复制）：
[Bug fixes] Align Kimi BF16 column linear under accuracy gate
总体评价
本次 diff 的默认 fused path 保持 use_accuracy_compatible=False，accuracy gate 也限制在 TP=1 且 input/weight 均为 BF16 的 experimental column-linear 分支，整体范围符合 PR 描述。新增测试覆盖了 gate、默认路径、dtype guard 和 ColumnParallelLinear.forward 的配置透传。
本地验证：python -m py_compile src/paddlefleet/tensor_parallel/layers.py tests/single_card_tests/tensor_parallel/test_column_sequence_parallel_accuracy.py 与 git diff HEAD^ HEAD --check 通过；python -m unittest tests.single_card_tests.tensor_parallel.test_column_sequence_parallel_accuracy 因当前环境缺少 paddle 包未能执行。
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #1411 +/- ##
===========================================
Coverage ? 100.00%
===========================================
Files ? 1
Lines ? 2
Branches ? 1
===========================================
Hits ? 2
Misses ? 0
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Kimi-K2 numerical-alignment scope updateThis PR should remain narrow: it contributes one measured BF16 column-linear mechanism and keeps the default path unchanged. The results below come from the complete gated Kimi stack, not from this PR alone. Validated on the real Kimi-K2 2-layer BF16 SFT checkpoint with native ms-swift/Megatron-Core and PaddleFormers/PaddleFleet runners:
This is not yet a complete-training alignment claim. The following native contracts remain to be closed end to end: global gradient norm/clipping, scheduler and warmup authority, all parameter groups/tied weights, distributed optimizer reduction/update order, and multi-step native training updates. A related authority distinction is important for review. PFCCLab/Megatron-LM PR #2 targets the MiniMax/PyTorch-AdamW stack and cannot be applied wholesale to this Kimi contract:
Therefore this PR's explicit TP1/BF16/accuracy-gated scope is intentional. It should not absorb model-specific optimizer, MoE-backward, or training-normalization changes. |
PR Category
Operator Mechanism
PR Types
Bug fixes
Description
This PR narrows the Kimi accuracy-compatible BF16 column-linear alignment change to the experimental ColumnParallelLinear path only.
use_accuracy_compatibleargument tocolumn_sequence_parallel_linear.paddle.nn.functional.linearto match the Megatron-Core no-TE behavior observed for Kimi shared-expert fc1.fused_linearpath unchanged unlessuse_accuracy_compatible=True.ColumnParallelLinear.forwardconfig propagation.Validation:
python -m py_compile src/paddlefleet/tensor_parallel/layers.py tests/single_card_tests/tensor_parallel/test_column_sequence_parallel_accuracy.pypre-commit run --files tests/single_card_tests/tensor_parallel/test_column_sequence_parallel_accuracy.py src/paddlefleet/tensor_parallel/layers.py --show-diff-on-failureBITWISE_PASSin the repro workspace.