From 1c2324587a7372241fbedbec12cf208eb11ed0c3 Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 27 May 2026 21:32:55 +0000 Subject: [PATCH] fix: pad agg_logp with NaN on zero-variance skip to keep is_s alignment The zero-variance bail at train.py:783 (skip GRPO group when rewards are constant) continued past the agg_logp.extend at line 821. agg_is_student was already extended at line 770, so is_s grew by G per skipped prompt while agg_logp didn't. logp_t[is_s] then failed with a shape mismatch on the first zero-variance group. Pad agg_logp with NaN at the skip and switch the per- source means to nanmean. Caught by #52 vanilla matched-control crashing at step 0. --- src/projected_grpo/train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 3957992..c1f835c 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -781,6 +781,10 @@ def main(cfg: Config) -> int: # dominant pathology with our binary-ish reward shape on a weak 2B # substrate (every group can clip to 0.25 = format_only). if (rewards.max() - rewards.min()).item() < 1e-4: + # Pad agg_logp with NaN to keep it aligned with agg_is_student + # (extended above at line 770). Skipping the gen_logp forward + # here is the whole point of the zero-variance bail. + agg_logp.extend([float("nan")] * len(rs)) continue centered = rewards - rewards.mean() adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4) @@ -943,9 +947,11 @@ def main(cfg: Config) -> int: gt_s_n = int((g_t & is_s).sum()) gt_t_n = int((g_t & ~is_s).sum()) rew_s_mean = rewards_t[is_s].mean().item() if n_s else float("nan") + # Skipped (zero-variance) prompts pad agg_logp with NaN above to keep + # alignment with is_s. nanmean drops them from the per-source means. logp_t = torch.tensor(agg_logp, dtype=torch.float32) if agg_logp else torch.zeros(0) - lp_s_mean = logp_t[is_s].mean().item() if n_s else float("nan") - lp_t_mean = logp_t[~is_s].mean().item() if n_t else float("nan") + lp_s_mean = logp_t[is_s].nanmean().item() if n_s else float("nan") + lp_t_mean = logp_t[~is_s].nanmean().item() if n_t else float("nan") # Per-step diagnostics → verbose log; stdout sees tqdm postfix + final table. n_fin = sum(agg_finished)