fix: whiten advantages over the DP group that includes context parallel

[keepkeen]

What

--normalize-advantages is computed with the wrong process group under context parallelism, so advantages are whitened with per-CP-rank statistics instead of global ones.

Root cause

compute_advantages_and_returns builds all_advs / all_masks from this CP rank's zigzag slice of every sequence (the cp_size != 1 branch right above constructs the per-rank mask chunks explicitly), and then reduces over the CP-excluding group:

dp_group = mpu.get_data_parallel_group()   # with_context_parallel defaults to False
whitened_advs_flat = distributed_masked_whiten(all_advs, all_masks, process_group=dp_group, ...)

distributed_masked_whiten all-reduces (sum, sum_sq, mask_sum) over that group, so the mean/variance only cover ranks that share the same cp_rank. Two consequences:

• each CP rank applies a different affine transform to a different part of the same sequence;
• no rank's statistics are the global ones, so even a single-CP-rank view is wrong.

Every other data-parallel group lookup in the repo is explicit about the flag — data.py:186,188, model.py:693, loss.py:1295 use with_context_parallel=True, actor.py:106,242 use False. This call site was the only implicit one.

Impact

Hits any run combining context parallelism with --normalize-advantages. Note slime_validate_args requires --normalize-advantages for reinforce_plus_plus and reinforce_plus_plus_baseline, and CI already runs --context-parallel-size 2 together with --normalize-advantages in test_qwen3_4B_ppo.py (and the disaggregate / critic-only variants) — nothing there checks the whitened values, which is why it went unnoticed.

Driving the real compute_advantages_and_returns at cp_size=2, dp_size=1 with three sequences (total, resp) = (64,48) (100,90) (40,12) and group-centered rewards +1.5 / -0.5 / -1.0:

cp_rank 0:  65 local tokens  mean=-0.076923
cp_rank 1:  85 local tokens  mean=+0.235294
GLOBAL   : 150 tokens        mean=+0.100000

 seq    raw |  CURRENT cp0  CURRENT cp1 |  FIXED cp0  FIXED cp1
   0   +1.5 |    +1.707037    +1.273883 |  +1.439168  +1.439168
   1   -0.5 |    -0.457986    -0.740629 |  -0.616786  -0.616786
   2   -1.0 |    -0.999241    -1.244257 |  -1.130775  -1.130775

The two halves of sequence 0 disagree by ~34%, and both are off the correct value.

Why the numel() > 0 guard has to go too

A CP rank can legitimately own zero response tokens: for a prompt-heavy sequence both of its chunks land inside the prompt.

cp=2 prompt=8000 resp=1000 -> response tokens per CP rank: [1000, 0]
cp=4 prompt=3000 resp= 500 -> response tokens per CP rank: [433, 67, 0, 0]

Today if all_masks.numel() > 0: makes participation in the all-reduce data-dependent. That is already latently unsafe (two DP ranks at the same cp_rank can disagree under dynamic batching), and it becomes an unconditional hang the moment the group spans CP. So the call is now unconditional — distributed_masked_whiten handles an empty local tensor, contributing 0 to all three reduced sums, and still raises if the mask is globally empty.

Test

tests/test_advantage_whiten_cp.py (new, registered in the cpu-unittest job) spawns dp_size * cp_size gloo workers, runs the real compute_advantages_and_returns, and asserts the whitened advantage of each sample is the same for every (dp, cp) factorization of the world and equal to the single-rank baseline. It covers (1,1) (2,1) (1,2) (2,2) (1,4) (4,1) and includes a prompt-heavy sequence so some ranks contribute an empty local mask.

On main it fails:

AssertionError: dp=1 cp=2: CP ranks disagree on sample 0: 1.504895 vs 1.762075

The spawn join is bounded at 180s so that a future rank-dependent branch around the collective surfaces as a test failure rather than a stuck CI job.

[keepkeen]

Hi @zhuzilin, could you please take a look when you have time? This fixes a reproducible training-correctness issue in advantage whitening under context parallelism, and adds a multi-process CPU regression test covering different DP/CP factorizations and empty local masks. GitHub currently reports the PR as mergeable, and all CI checks pass. Thanks!