Add never_give_up (NGU) GRPO retry mechanism

[mnoukhov]

What

Adds the never_give_up (NGU) GRPO retry mechanism, ported clean off the long-lived ngu branch (only the core mechanism + its advantage baseline — no ngu_seq_multiplier continuations, no reinforce_ada_est, no batch_by=completions, no utilization-metric changes).

With probability --never_give_up, an unsolved zero-std prompt group (one that would otherwise be filtered by filter_zero_std_samples) is requeued under a retry-suffixed prompt id ({epoch}_{index}_1, _2, ...) instead of being dropped. Its completions are buffered across accumulation calls until a later attempt in the chain finally has learning signal, at which point:

• the buffered attempts are merged back into one logical group, and
• the group's advantage baseline is computed over every attempt in the chain (not just the accepted batch), then re-anchored so the positive (max-reward) samples keep their advantage and each group re-centers to zero.

Two design choices from the ngu branch are hardcoded here (with explaining comments) rather than exposed as knobs:

• accept rule = "better": a retry is accepted only when its max reward strictly beats the chain's previous best;
• baseline rescaling = "anchor positives": applied whenever the filtered chain-wide baseline is active.

How

• New open_instruct/data_loader_utils.py: chain-id / retry-suffix helpers, NeverGiveUpAccumulationState + pending-state pop/store with model-step age filtering, the accept rule, GenerationResult merge/select, and compute_grouped_advantages (variable group sizes, filtered NGU baseline, always-on anchor-positives rescaling).
• data_loader.py: the monolithic process_group is split into a pure process_group + maybe_dont_resample_prompt + maybe_filter_group + maybe_replenish_prompt pipeline. Group / BatchStatistics carry per-group sample_count / baseline_sample_count / baseline_reward_sum / attempt_count; make_batch_from_groups no longer assumes a uniform generation_config.n per group. DataPreparationActor owns one NGU state and routes NGU runs (never_give_up > 0) through compute_grouped_advantages. The non-NGU path is unchanged (never_give_up == 0 keeps the original scores.reshape(-1, n) math, byte-for-byte).
• Config: --never_give_up (float prob), --maintain_pending_ngu_age, --maintain_pending_ngu_completions. never_give_up > 0 is rejected together with mask_truncated_completions (would desync merged-group counts).
• Minor: with no_resampling_pass_rate set, the "exclude this index" call now happens just before replenishing rather than just after — no training-behavior change in practice.

Tests

• open_instruct/test_data_loader_utils.py (new): chain ids, accept rule, pending-state round-trip / age filtering, and compute_grouped_advantages — including a regression test that it reproduces the old inline reshape advantages for uniform groups.
• open_instruct/test_data_loader.py: maybe_filter_group — zero-std filter with NGU off, requeue + buffer with NGU on, and accept + merge (sample_count/baseline_sample_count == 2n, attempt_count == 2).
• uv run pytest open_instruct/test_data_loader.py open_instruct/test_data_loader_utils.py — green.
• make style && make quality — ruff clean; ty check unchanged from main (51 pre-existing).

Smoke test

scripts/train/debug/never_give_up_single_gpu.sh (= single_gpu_on_beaker.sh + --never_give_up 1.0).

Runs:

1. Single-GPU NGU smoke test: Beaker

GPU_TESTS=bypass