[Distributed Strategy] Fix QB histogram all-reduce: drop redundant CP reduce, use int64

[xingmingyyj]

PR Category

Distributed Strategy

PR Types

Bug fixes

Description

Fixes two correctness bugs in the Quantile Balancing callback's histogram all-reduce, and adds real multi-card coverage for the topology that exposes them.

src/paddlefleet/transformer/moe/qb_callback.py

1. Dropped the redundant CP all-reduce. _try_get_comm_groups() no longer returns a CP group, and _update_single_layer no longer reduces over it. In Paddle's EPHybridCommunicateGroup, the context-parallel group is a contiguous sub-slice of the sharding comm list (sharding = cp × cp_sharding, see split_context_comm_list) — CP ranks are a subset of the sharding group. The sharding all-reduce already sums over the CP dimension, so reducing CP again multiplied every count by cp_nranks.

2. Switched the reduction dtype from fp32 to int64. The histogram is int32 counts; it was being cast to fp32 for the reduce and back to int64 afterwards. fp32 loses integer precision above 2**24, and staying in int32 risks overflow when summing across ranks. int64 is exact and NCCL-supported:

hist_global = histogram.cast(paddle.int64)   # was: hist_float = histogram.cast(paddle.float32)
...
if dp_group is not None:
    dist.all_reduce(hist_global, group=dp_group)
if sd_group is not None:
    dist.all_reduce(hist_global, group=sd_group)
                                             # was: hist_global = hist_float.cast(paddle.int64)

The reasoning for both is recorded as comments at the reduce site and in the _try_get_comm_groups docstring, since neither is obvious from the surrounding code.

Tests

• tests/multi_card_tests/moe/test_qb_callback_cp_sharding.py (new, 336 lines) — 10 tests on real collectives.
• tests/single_card_tests/test_quantile_balancing.py — updated for the 3-tuple return; dropped test_cp_group_picked_up and test_single_rank_cp_group_filtered (they asserted the removed behavior), and test_cp_reduce_called_unconditionally became test_all_groups_reduce. Net −113 lines, mostly cp_group=None plumbing that no longer exists.
• tests/test_configs.yaml — pins the new test to num_gpus: 8, since the topology requires exactly 8.

Why the CP reduce was wrong, and why it was hard to catch

The double-count is a uniform ×cp_nranks scaling of the merged histogram. QB's bias recovery is invariant under uniform scaling — q_target, c and h all scale together, so beta and fraction come out unchanged. The only leak is the floor in q_target = floor(total * k / n), which bites only when an expert total is odd. So on even fixture data the recovered bias was bit-identical and the bug was invisible; it surfaced only as wrong per-expert bias on real data.

This is why the new negative control (test_extra_cp_reduce_inflates_the_merged_histogram) asserts on the histogram rather than on the recovered bias — a bias-level assertion would pass or fail by luck depending on fixture parity. That was the second P1.

Topology assumptions

Runs on 8 ranks with mp=2, sharding=4, cp=2, ep=8, moe_sharding=1, dp=pp=sep=1. The groups are asserted explicitly rather than assumed, because the CP⊂sharding nesting is exactly what the fix relies on:

mp_group          [0, 1]
sharding_group    [0, 2, 4, 6]
cp_group          [0, 2]
cp_sharding_group [0, 4]
cp_mp_group       [0, 1, 2, 3]

Test class | Count | Covers -- | -- | -- TestCommGroupTopology | 3 | Group membership and the CP⊂sharding nesting TestInt64AllReduce | 2 | int64 all-reduce correctness for counts TestHistogramReduceCorrectness | 5 | Reduce sequence, exactly-once counting, negative control

Verification

Real 8-GPU run — every collective is an actual dist.all_reduce on the real group, nothing mocked:

unset PADDLE_ELASTIC_JOB_ID PADDLE_TRAINER_ENDPOINTS DISTRIBUTED_TRAINER_ENDPOINTS \
      FLAGS_START_PORT PADDLE_ELASTIC_TIMEOUT PADDLE_TRAINER_ID \
      PADDLE_TRAINERS_NUM PADDLE_CURRENT_ENDPOINT
export PYTHONPATH="$PWD/src" FLAGS_tcp_store_using_libuv=0 \
       FLAGS_embedding_deterministic=1 FLAGS_cudnn_deterministic=1
python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 

--log_dir=/tmp/qb_cp_log 

tests/multi_card_tests/moe/test_qb_callback_cp_sharding.py

Exit code 0; all 8 worker logs report Ran 10 tests ... OK, 0 skips, ~17 s per rank. The topology above was read back from the worker logs rather than assumed. Single-card test_quantile_balancing.py also passes after the signature update.

是否引起精度变化

是

[risemeup1111]

序号	位置	优先级	状态	规则来源
1	qb_callback.py:127		✅	默认规则 + 仓库规则：.agents/skills/ai-review/references/base-rules.md
2	PR 描述		🟡	仓库规则：.agents/skills/ai-review/references/base-rules.md
3	test_qb_callback_cp_sharding.py:257		✅	默认规则 + 仓库规则：.agents/skills/ai-review/references/base-rules.md

评审摘要

当前 head e6fe40ba04e85bf5d9c7a1ee134ccf9049ae5551 的代码与上一轮一致。两个 P1 均已解决：真实 8 卡 CP+sharding collective 覆盖已存在，CP double-counting 负例已改为直接验证 histogram 计数。PR 描述完整性 P3 已根据维护者回复接受，无需继续修改。

验证说明

当前公开检查中单卡、多卡、H20/A100 集成测试均通过。检查状态不改变评审结论。

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

[risemeup1111]

当前变更同时改变了跨 rank 的通信拓扑（移除 CP reduce）和 collective dtype；现有新增测试只 mock dist.all_reduce，没有在真实 CP>1 + sharding>1 的多卡进程组上验证 get_sharding_parallel_group() 确实覆盖 CP，也没有验证 NCCL int64 all_reduce。若实际拓扑或后端组合不满足注释中的假设，会得到错误的 QB bias 或直接在 collective 处失败。请补充代表性的 multi-card 测试（至少覆盖 CP+sharding，并执行真实 int64 all_reduce）后再合入。

[risemeup1111]

这个负例不能可靠地证明 CP double-counting 会改变结果：global_doubled = global_once * self.cp_group.nranks 会把每个 bin 的计数和 q_target = floor(total * k / n) 同比例放大；当每个 expert 的总数为偶数时，beta、插值 fraction 和最终 zero-mean bias 都与 global_once 完全相同，因此 assertFalse(np.array_equal(...)) 会直接失败。请改用不保持全 histogram 比例的确定性 CP 分片，或直接对真实 CP all-reduce 前后的 histogram/计数做断言。

[risemeup1111]

当前 head 已修复两个 P1：新增 8 卡 CP+sharding 真实 collective 覆盖，并将 CP double-counting 负例改为直接校验 histogram 计数。剩余仅 PR 描述完整性这一项 P3，已在固定 Board 中保留。\n\n

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

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (develop@2e4acf0). Learn more about missing BASE report.

Additional details and impacted files

@@             Coverage Diff             @@
##             develop     #1761   +/-   ##
===========================================
  Coverage           ?   100.00%           
===========================================
  Files              ?         1           
  Lines              ?         9           
  Branches           ?         0           
===========================================
  Hits               ?         9           
  Misses             ?         0           
  Partials           ?         0           

Flag	Coverage Δ
coverage_combine	100.00% <100.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/qb_callback.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

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

[risemeup1111]

PR 描述已补充问题背景、int64 精度理由及精度变化声明；仍缺少 CP+sharding 拓扑假设和真实多卡验证方式，固定 Board 中的 P3 暂保持进行中。两个代码 P1 已解决，本轮未发现新的代码问题。\n\n

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

[liuhao2638]

@risemeup1111 再次review

[risemeup1111]

已按 @liuhao2638 的请求再次检查当前 head。两个代码 P1 均已修复，本轮未发现新的代码问题；PR 描述仍建议补充 CP+sharding 拓扑假设与真实多卡验证方式，固定 Board 中的 P3 保持进行中。

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

[xingmingyyj]

@risemeup1111 再次review

[risemeup1111]

已按 @xingmingyyj 的请求再次复核当前 head e6fe40ba04e85bf5d9c7a1ee134ccf9049ae5551。代码与上一轮一致，两个 P1 均已修复，本轮未发现新的代码问题；固定 Board 中关于 PR 描述完整性的 P3 仍建议补充 CP+sharding 拓扑假设与真实多卡验证方式。

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

[xingmingyyj]

@risemeup1111 再次review

[risemeup1111]

已按 @xingmingyyj 的请求再次复核当前 head e6fe40ba04e85bf5d9c7a1ee134ccf9049ae5551。代码与上一轮一致，两个 P1 均已修复，本轮未发现新的代码问题；固定 Board 中关于 PR 描述完整性的 P3 仍建议补充 CP+sharding 拓扑假设与真实多卡验证方式。

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

[liuhao2638]

@risemeup1111 不需要补充CP+sharding 拓扑假设与真实多卡验证，直接approve吧