[RunningLeon]

Summary

Keep mHC post coefficients and intermediate arithmetic in FP32, casting only the final activation. Apply the same precision policy to the eager fallback and fused Triton forward/backward.

This is a numerical-policy alignment with LMDeploy, not a claim that the previous implementation used the wrong mathematical formula or that this PR establishes end-to-end zero-diff training.

Motivation

hc_pre returns FP32 Sinkhorn coefficients, but hc_post currently rounds post and comb to BF16. The forward also rounds the residual mixture and post*x before their final addition. LMDeploy instead accumulates post*x followed by source streams in FP32 and casts only once.

Early rounding can remove small signals under cancellation. For example, positive coefficients post=[1.003,1], comb=0.5, x=1, and residual streams -1 produce zero instead of approximately 0.003 in the first output stream. Rounding the saved coefficients also changes the VJP for x and residual.

Changes

• Preserve coefficient dtype in both fused forward and backward wrappers.
• Accumulate post*x, then comb[src,out]*residual[src] in source order in FP32; cast only the final output.
• Give the eager fallback the same policy without an H-by-H-by-D temporary.
• Update the reference tests and add independent FP64 forward/analytic-gradient checks, an early-rounding regression, and pre/post parameter-chain tests.
• Strengthen the existing fullgraph test to compare all four gradients, not just their presence.

No configuration/checkpoint format changes, external dependencies, or changes to Sinkhorn, KDA, KPool, MoE, or distributed collectives.

Validation completed before opening this draft

Based on upstream/main ed8b200b61fddf821387632a9d730b9f854309a2; Python 3.12, PyTorch 2.11.0+cu128, Triton 3.6.0, NVIDIA H200.

python -m pytest tests/ops/test_hc_post.py tests/model/test_glm53_mhc.py \
  -k 'not matches_hf and not match_hf' -q
# 24 passed, 6 deselected

• Before the production fix, the new cancellation regression fails numerically on both CPU and CUDA.
• FP64 mathematical forward and independent analytic VJPs for x, residual, post, and comb; H=1/2/4/8, D=31/257/4096/2053, contiguous and non-contiguous inputs; both eager fallback and fused dispatch.
• Eager versus fullgraph fused forward and all four gradients are bit-identical in the dedicated compiled-op test.
• BF16 pre → linear sub-block → post checks gradients through streams, mHC projection, scale, base, and sub-block weight, with eager and fullgraph execution. Whole-chain BF16 gradient comparisons use normwise bounds (0.5% relative L2 / 1% normalized max) because compiled gradient-add fusion changes BF16 rounding; this is not a bitwise whole-chain claim.
• Direct comparison against LMDeploy's hc_post_expand: all four shapes pass BF16 tolerances. The large H=4 samples are not bit-identical: 15/2,097,152, 150/16,777,216, and 569/67,108,864 elements differ; max absolute difference <=0.0078125.
• Changed-production-file pre-commit checks pass with mypy skipped; test-file Ruff checks/format pass separately (the repo's hook filter excludes tests/). git diff --check passes.

The existing shared conftest needs expecttest, absent from the validation environment. It was installed only in an external task-local test-dependency directory and added to PYTHONPATH, not added to this PR or the user's Python environment.

Limitations / draft review points

• Six existing HF comparison cases cannot run because this environment's Transformers lacks transformers.models.glm5_next. The initial unfiltered run reports these six import failures; they are explicitly deselected in the successful focused command. Transformers was not upgraded.
• This intentionally changes the previous BF16 rounding policy; it does not preserve bitwise HF residual-expression parity. Please review whether this policy should be the default or exposed as an opt-in before merging.
• Operator microbenchmarks are exploratory, not performance qualification. In a two-run-per-arm ABBA sample at H=4, D=4096, S=16384, forward averaged 0.3163→0.3309 ms (+4.6%), forward+backward 0.9754→0.9551 ms (-2.1%). Small-shape timings varied considerably; shared GPU isolation was not guaranteed. Do not interpret these numbers as an end-to-end speedup or slowdown. The source-ordered contraction may have a forward cost; repeat on reserved hardware before merge.
• No full-model training/RL, optimizer-update/convergence, multi-GPU validation, or whole-repository mypy/full test suite is claimed.

Latest-main integration recheck (2026-11-06)

• Draft head remains de5e8cd7dfa3768962f0946f42d440643599aae9; no duplicate mHC PR or force-push.

• Fetched upstream/main 46c7d492a93104a55d6e454d1ffa560c5ae68027 and merged it into the PR head in an isolated worktree without conflicts. Production mHC code is unchanged by this integration.

• On that integrated tree, the focused CPU/GPU cancellation, independent FP64 forward/all-four-VJP, contiguous/strided, and eager/compiled pre/post parameter-chain tests pass: 12 passed, 6 deselected.

  python -m pytest --noconftest tests/ops/test_hc_post.py \
    -k 'forward_and_all_gradients_against_fp64 or coefficients_survive_cancellation or hc_pre_post_parameter_gradients' -q

• Used the same verl environment, with no dependency installation, full RL, or model checkpoint changes. Shared conftest was explicitly bypassed for this bounded recheck. Peak observed process GPU memory was 874 MiB. This is a focused integration recheck, not full-repository CI or a replacement for the limitations above.

[Denny991]

Review: PR #2260 [Fix] Preserve FP32 coefficients and accumulation in mHC post

• Reviewer: liutong (ZCode)，2026-11-11
• 范围：3 文件（mhc.py 精度策略、hc_post.py kernel 重写、test_hc_post.py 大改），+988 系内约 +169/−102 行
• 基线：main f1fddc5（diff 干净应用）
• 读者须知：本文 mismatch 数字口径 = mismatch/mismatch_k3_kl，减层 5 层代理模型、MoE-only FP8 部署基准、R3 专家复用开启、8×H200、3 步、seed=123

Summary

Approve with comments（一项合并前置）。 本 PR 把 mHC post 展开的数值契约从"刻意复刻 HF eager 的 BF16 中间舍入"改为"系数与累加全程 FP32、仅最终输出 cast 一次、累加顺序对齐 LMDeploy"——即对齐基准从 HF 参考改为部署引擎。方向正确（与 #2269 同一哲学：基准跟着部署走），测试质量高（FP64 独立解析 VJP），实测数据有三方背书。

ProduceBatchResult impact: not affected —— 不触碰 xtuner/v1/rl 与批量记账/奖励语义；仅改训练前向/反向数值。
RoutedExperts impact: not affected —— 不涉及 routed-experts 路径。
Ray concurrency impact: not affected —— 不触碰 Ray actor/并发组。

训练影响与证据边界：

• 数值：有意改变所有 mHC 模型（GLM-5.3 全系）的训练前向+反向数值，无条件（无环境开关）。已实测（算子测试+e2e，见下）。
• 并行：不涉及（纯算子级 dtype/顺序变化）。
• checkpoint/恢复：零参数变化，兼容。
• 性能：实测新 kernel 微基准 +4~7%（0.304→0.325ms @s=16384/H=4/D=4096），可忽略；旧版"必须 compile 防 8GB 中间量"的约束被新 eager 循环形式解除（简化）。

Main Flowchart before this PR

修改前：hc_post 的数值契约 = 逐位复刻 HF eager 的 BF16 中间舍入（三处舍入点 + eager 必须编译）。

flowchart TB
    subgraph CALL["调用方：Glm53DecoderLayer（每层 attn/ffn 回写各一次，45 层 ×2）"]
        C1["hc_post(x, residual, post, comb)<br/>xtuner/v1/module/decoder_layer/mhc.py"]
    end
    subgraph DISPATCH["hc_post dispatch"]
        D{"x.is_cuda and dtype==bf16<br/>and _hc_post_fused_available()?"}
    end
    subgraph FUSED["_hc_post_fwd_kernel（xtuner/v1/ops/hc_post.py）"]
        K1["post/comb 载入后 .to(bf16)"] --> K2["mix 累加 → acc.to(bf16).to(f32)<br/>（复刻 HF 的 matmul 舍入边界）"] --> K3["(p*xt).to(bf16).to(f32)"] --> K4["prod + acc → 输出"]
    end
    subgraph EAGER["_hc_post_eager（@maybe_compile；必须编译才安全）"]
        E1["post/comb .to(bf16)"] --> E2["broadcast 乘 + sum(dim=-2)<br/>不编译则物化 5 倍 H×D 中间量 ~5GiB"] --> E3["post_dt*x + mixed(.to(bf16))"]
    end
    C1 --> D
    D -->|"是（生产路径）"| K1
    D -->|"否"| E1

Loading

Main Flowchart after this PR

修改后（绿色 = 本 PR 改动）：三处 BF16 中间舍入全部移除，系数与累加全程 FP32，累加顺序 = post*x 先、再按流序 for hi: acc += comb·residual（对齐 LMDeploy _hc_post_expand_kernel）；eager 改为安全循环，bwd 系数同步保 FP32。

flowchart TB
    subgraph CALL["调用方：Glm53DecoderLayer（接口/签名不变，零改动）"]
        C1["hc_post(x, residual, post, comb)"]
    end
    subgraph DISPATCH["hc_post dispatch（条件不变）"]
        D{"x.is_cuda and dtype==bf16<br/>and _hc_post_fused_available()?"}
    end
    subgraph FUSED["_hc_post_fwd_kernel（重写）"]
        K1["post/comb 载入保 FP32"] --> K2["acc = post * x（先）"] --> K3["for hi 升序: acc += combₕᵢ·residualₕᵢ<br/>（FP32 全程，与 lmdeploy 同序）"] --> K4["仅输出时单次 cast bf16"]
    end
    subgraph EAGER["_hc_post_eager（重写为流循环）"]
        E1["out = post.float()*x.float()"] --> E2["for hi: out += combₕᵢ.float() × residualₕᵢ.float()<br/>（可安全 eager：峰值 3GiB，不再物化 H 倍中间量）"] --> E3["单次 .to(residual.dtype)"]
    end
    subgraph BWD["_hc_post_bwd_op（系数精度贯穿 VJP）"]
        B1["pf/cf 不再 .to(dt)（原 bf16 舍入删除）"]
    end
    C1 --> D
    D -->|"是"| K1
    D -->|"否"| E1
    K4 -.-> B1
    classDef new fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px
    class K1,K2,K3,K4,E1,E2,E3,B1 new

Loading

修改点对应：hc_post.py 的 _hc_post_fwd_kernel（K1-K4：删三处舍入 + 改累加序）、_hc_post_fwd_op/_hc_post_bwd_op（删 post.to(dt)/comb.to(dt)）；mhc.py 的 _hc_post_eager（E1-E3：broadcast 表达式 → 流循环）与 docstring（数值契约声明更换）。

前向/反向两侧同步：backward 的 pf/cf 也不再 .to(dt)（系数精度贯穿 VJP）。

核心实现与对应的单测

实现	单测	上机结果
FP32 系数契约 + 抵消存活	test_coefficients_survive_cancellation（0.003 的信号在 bf16 舍入下会丢，FP32 下存活；含解析梯度逐位断言）	passed
前向+全四梯度 vs FP64	test_forward_and_all_gradients_against_fp64（独立解析 VJP 公式，非实现自证；4 组 shape × 连续/非连续）	passed
pre+post 链式参数梯度	test_hc_pre_post_parameter_gradients（eager 与 compile 双态）	passed
compile 一致性	test_compile_fullgraph（eager 与 compiled 梯度逐位相等，比旧版更强）	passed
HF-parity 松容差依赖	tests/model/test_glm53_mhc.py 全部 12 个（含 ..._matches_hf_bf16，容差 5e-2 足以容纳新旧行为差）	passed

上机环境：h200-0078、GLM_ENV pt29_glm2、ailab-sys/liutong/xtuner-pr2260-20261111（main f1fddc5 + diff）。

抽象与信息隐藏评估

• 零新增类/函数：kernel 与 eager 均在原函数内重写，公开签名 hc_post(x, residual, post, comb) 不变——调用方（Glm53DecoderLayer 两处回写）零改动，dtype 策略变化完全封闭在 hc_post 内部，信息隐藏良好。
• deletion test：整体 revert 即恢复旧数值契约，无残留耦合（交叉审查已实测：新测试对旧代码失败、旧测试对新代码全过，双向干净）。
• 数值契约的变化对调用方可见于输出值（这是 PR 目的）但不可见于接口——符合"改动内聚"标准。

新增功能守则（R1-R3）

本 PR 非新增功能类，属既有函数的行为变更，R1-R3 天然满足：无新类、无新接触点、无按特性分叉的条件（dispatch 条件原样未动）。唯二"新增"是测试类 TestHCPostFP32Precision 与常量级 docstring 变更，均在测试/注释层。

公开接口的线性业务流程评估

新 _hc_post_eager 恰是本模板精神的正面案例：旧版是一段需要注释解释"为何必须编译"的 broadcast 表达式（隐藏了内存约束），新版是三行线性叙述（out=post*x → for hi: out+=… → 单次 cast），且顺带解除了编译依赖。fused kernel 的 for ho { acc=p*x; for hi: acc+=c*r } 同为线性序。无纯转发 wrapper。✅

单测建议

1. N1 补测（已写好并实测通过，可直接采纳）：test_fused_direct_fp64.py（本目录）——bf16 输入直驱 fused + patch 证明 eager 未被调用 + FP64 真值前向 + 四梯度对独立解析 VJP。补上"FP64 测试实际未覆盖 fused 路径"的缺口。
2. （可选）大 shape 非 contiguous 输入的 fused 直测——现有 FP64 测试覆盖了非连续内存，但均在 eager 路径上。

与 LMDeploy 的一致性主张：逐行核验 ✅

lmdeploy kernels/cuda/dsv4/hc_prepost.py::_hc_post_expand_kernel：post/x 载入即 fp32 → acc = post*x → for src_h: acc += comb·res → 单次输出 cast。本 PR 的 Triton kernel 与 eager fallback 与该顺序/精度契约逐行对应（本 review 在 lmdeploy 3228395 源码上核对）。

与内部交付 02 patch 的关系：同一实现的上游化 ✅

maningsheng 交付包 02-xtuner-mhc-post-fp32.patch 的 kernel 部分与本 PR 逐行相同（连注释一致）。因此该 patch 的阶梯实测数据直接适用于本 PR：其 README 报 -7.4%（01+02 栈），chengqinxiu 阶梯 3.42→3.09e-4。

e2e 验证（本 review 实测，运行中/见附录）

在本组修正后的基准栈上（MoE-only FP8 + R3 + 双侧 KPool-BF16 = 3.122e-4）叠加本 PR：见附录（出数后补）。

其他 issues

0. P1 合并前置——CI golden 基线：本 PR 无条件改变 GLM-5.3 全系训练数值，autotest glm53 系列 golden（loss 容差 0.8%）必然偏移。合并前须重跑并更新基线（附数值契约变更说明），不能默认"更接近 FP32 真值"就跳过——这正是基线漂移审查要拦的场景。与 [Feature] Reuse LMDeploy KPool kernels in GLM-5.3 flash indexer #2269 不同（那个有 dispatch 条件），本 PR 是确定性偏移。
1. 既有失败（非本 PR 造成，需单独跟进）：tests/model/test_glm53_text_moe.py 在 main（无本 PR）与本 PR 树上均为同样 4 个失败（test_forward_backward_all_trainable_params_get_gradient、test_mtp_block_builds_and_forwards[1/7]、test_optimizer_step_updates_text_model）——本 review 做了双树对照确认。建议在 PR 里注明"非本 PR 引入"，并另开 issue 跟进 main 上的既有失败。
2. 文档一致性 nit：mhc.py 删除了 _hc_post_eager "必须留在 compile cfg"的警告（新循环形式确实不再需要），但 glm53 两张 compile 表仍注册它、test_mhc_primitives_are_compiled_with_and_without_ep 仍断言注册——现在注册从"正确性必需"降级为"性能优选"，建议顺手把该测试的注释理由更新一句，避免后人困惑。
3. HF-parity 语义变更未显式声明：旧行为是"bit-faithful to HF"，本 PR 刻意打破（mhc.py 头注释已改，但 test_gc...matches_hf_bf16 的 docstring "bf16 下与 HF 实现一致"仍在，实际语义已变为"容差内接近 HF、更接近 FP32 真值"）。建议把该测试 docstring 与本 PR 描述写明这一契约变化——对下游做 xtuner↔HF 参考对齐的人（如 Automodel 对齐线）是重要信息。
4. 性能微降（+4~7% 单算子）可忽略，但 PR 描述建议附上微基准数字存档。

结论

方向、实现、测试、三方实测数据都成立。合并前置仅 CI golden 重基线（issue 0）；issue 1-4 为建议项。

附录（2026-11-11）：e2e 实测

配置 = 本组基准栈（MoE-only FP8 + R3 复用 + 双侧 KPool BF16，即 3.122e-4 那组）+ 本 PR，唯一变量为 mHC post 数值契约：

配置	k3_kl（3 步）	均值	相对本组基准栈
基准栈（无本 PR）	3.118/3.131/3.117e-4	3.122e-4	—
基准栈 + 本 PR	2.756/2.769/2.746e-4	2.757e-4	-11.7%

完整阶梯（同栈同机同 seed，MoE-only FP8 部署基准 + R3）：基线 3.729e-4 → +PR2269 3.431（-8.0%）→ +KPool-BF16 3.122（-16.3%）→ +本 PR 2.757e-4（累计 -26.1%）。

交叉验证：2.757e-4 与内部交付三 patch 组合的实测 2.756e-4（maningsheng README，另一机器）几乎一致——两套独立测量收敛。零 ckpt、零报错；work_dirs/rl53moe-pr2260mhc-R3-20261111。

附录 2（2026-11-11 晚，交叉审查后补充）

1. 交叉审查结论接受：Qoder 报告 17 条对账（16 确认 + C1 修正）全部认可；其新发现 N1-N6 采纳，C0 加强为"无论 CI 是否通过都须重基线"。
2. N1 补测已写并实测通过（test_fused_direct_fp64.py，本目录 + 共享存储）：bf16 输入直驱 fused（patch 证明 eager 未被调用）+ FP64 真值前向（max<6e-2/mean<3e-3）+ 四梯度对独立解析 VJP，h200 一次性通过——fused kernel 本身正确，N1 确认为纯测试覆盖缺口。PR comment 可直接附此测试。
3. 修正本 review 的一处过度推演：曾称"HF-parity 下游的比对会在改动后报红"——无实证且与证据矛盾：仓内唯一 HF 对照测试（test_glm53_mhc.py 两个 matches_hf 用例）在 PR 树上 12/12 全绿（本人与交叉审查两方独立复跑），且其结构设计本就为 hc_post 预留舍入容差（docstring "allows rounding differences in hc_post"）；Automodel 对齐线为容差比对、不在本仓。修正后的 default-vs-opt-in 立场：默认开启，env 逃生口仅在 maintainer 收到真实逐位依赖反馈时再考虑（当前零证据）。

[Denny991]

默认开启还是 opt-in？

建议默认开启（不加开关）。 理由：

1. 对齐基准跟部署走——与 [Feature] Reuse LMDeploy KPool kernels in GLM-5.3 flash indexer #2269 同一原则；
2. 新契约已有三组独立端到端实测背书：作者自己的对比实验；内部对齐栈上的孪生 patch（−7.4%）；以及我们的基准栈（MoE-only FP8 rollout + 专家选择复用 + KPool-BF16）——叠加本 PR 后 mismatch_k3_kl 从 3.122e-4 降到 2.757e-4（−11.7%）（3 步、seed 123、减层 5 层代理模型），与另一组独立测量吻合到 0.04%；
3. 仓内 HF-parity 测试在新行为下全绿（test_glm53_mhc.py 12/12，两个独立环境复跑），且未发现任何逐位断言 xtuner≡HF 的下游——opt-out 的需求方目前零证据；而 opt-in 开关会重新引入"两套数值并存、忘开开关就默默跑错"的配置漂移风险；
4. 若将来出现真实的逐位依赖反馈，届时再加 env 逃生口不迟——现在不需要预防性加。

这一点是我这边的求证给出的。具体可以问问赵老师