"""Audit 2026-10-14 (G4-rl-batch) regressions for `align/rl/batch_source.py`. U04-1 fewer live rows than world x accum slots killed the run (RuntimeError after the rollout was paid for); U04-2 row quantisation threw away up to ~50% of live rows and toggled the step size 2x; U04-3 peers waited for rank 0's rollout on a collective bounded by the process-group timeout, and a rank-0 failure left them blocked. """ from __future__ import annotations import os import tempfile import time import pytest import torch from test_grpo_module import _batch, _module from test_grpo_driver import FakeClient, _source from align.rl.batch_source import GrpoBatchSource, RolloutMicroHandle def _payload(B: int, T: int = 5) -> dict[str, torch.Tensor]: """B group-contiguous rows; row i is identifiable by input_ids[i, 0] == i + 1.""" ids = torch.zeros((B, T), dtype=torch.int64) ids[:, 0] = torch.arange(1, B + 1) rm = torch.zeros((B, T), dtype=torch.int64) rm[:, 2:] = 1 return { "input_ids": ids, "attention_mask": torch.ones((B, T), dtype=torch.int64), "response_mask": rm, "advantages": torch.ones(B), "vllm_logprobs": torch.zeros((B, T - 1)), "group_ids": torch.arange(B) // 16, "reward::reward/total": torch.ones(B), } def _bare(world: int, accum: int, rank: int = 0, prompts: int = 32, g: int = 16) -> GrpoBatchSource: src = GrpoBatchSource.__new__(GrpoBatchSource) src.world_size, src.grad_accum, src.rank = world, accum, rank src.prompts_per_step, src.n_samples_per_prompt = prompts, g src.dropped_remainder_rows = 0 return src def _real_rows(slices) -> list[int]: """Row ids carried by real (non-padding) rows of one rank's slices.""" out = [] for s in slices: for i in range(s["input_ids"].shape[0]): if int(s["response_mask"][i].sum()) > 0: out.append(int(s["input_ids"][i, 0])) return out def _all_ranks(world, accum, B): per_rank = [] for r in range(world): src = _bare(world, accum, rank=r) per_rank.append(src._slice_rank(_payload(B))) return per_rank def test_u04_1_fewer_live_rows_than_slots_still_trains_every_row(): """World 3 x accum 85 = 255 slots; 240 live rows (17 of 32 groups dead) used to raise.""" per_rank = _all_ranks(3, 85, 240) for slices in per_rank: assert len(slices) == 85 assert all(s["input_ids"].shape[0] >= 1 for s in slices), "a rank would run an empty micro" used = sorted(i for slices in per_rank for i in _real_rows(slices)) assert used == list(range(1, 241)), "every live row must reach exactly one micro" def test_u04_1_padding_rows_contribute_nothing_and_are_flagged(): per_rank = _all_ranks(3, 85, 240) pad_rows = [ s for slices in per_rank for s in slices if int(s["response_mask"].sum()) == 0 ] assert len(pad_rows) == 255 - 240 for s in pad_rows: assert float(s["advantages"].abs().sum()) == 0.0 assert s["input_ids"].shape[0] == 1 def test_u04_2_one_dead_group_no_longer_discards_half_the_rows(): """B=496 (one dead group of 32x16, 3 ranks, accum 85) used 255 rows and dropped 241.""" per_rank = _all_ranks(3, 85, 496) used = sorted(i for slices in per_rank for i in _real_rows(slices)) assert used == list(range(1, 497)) src = _bare(3, 85) src._slice_rank(_payload(496)) assert src.dropped_remainder_rows == 0 def test_u04_2_full_step_shape_is_unchanged(): """The nominal full-live shape keeps its per-micro row cap (memory budget): 512 rows over 255 slots is 2 rows per micro, 2 rows dropped and counted -- exactly the previous behaviour.""" src = _bare(3, 85) slices = src._slice_rank(_payload(512)) assert all(s["input_ids"].shape[0] == 2 for s in slices) assert src.dropped_remainder_rows == 2 def test_u04_2_micro_sizes_never_exceed_the_nominal_cap_and_ranks_stay_balanced(): for B in (255, 300, 400, 496, 510, 512): per_rank = _all_ranks(3, 85, B) for m in range(85): sizes = [per_rank[r][m]["input_ids"].shape[0] for r in range(3)] assert max(sizes) <= 2 and max(sizes) - min(sizes) <= 1, (B, m, sizes) def test_u04_1_pad_micro_runs_through_the_real_module_with_zero_loss_and_live_graph(): mod, model = _module(grad_accum=1) src = _bare(1, 4, prompts=1, g=8) src.device = "cpu" src.old_logprob_fn = lambda ids, attn: mod.compute_logprobs_nograd(ids, attn) b = _batch(model, B=2, T=10) payload = { "input_ids": b.input_ids, "attention_mask": b.attention_mask, "response_mask": b.response_mask, "advantages": b.advantages, "vllm_logprobs": b.vllm_logprobs, "group_ids": b.group_ids, } slices = src._slice_rank(payload) # 2 rows over 4 slots: two micros are padding pads = [s for s in slices if int(s["response_mask"].sum()) == 0] assert len(pads) == 2 batch = src._finalize(model, pads[0]) out = mod._grpo_loss_adapter(model, batch) assert float(out.loss) == 0.0 out.loss.backward() assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) def test_u04_1_tripwire_is_not_fed_from_a_padding_micro(): """A padding micro reports entropy 0 / length 0; observing it would fake a collapse.""" seen = [] class _Wire: fired = False fire_reason = None def observe(self, step, entropy, length): seen.append((step, entropy, length)) return False src = _source(g=8, accum=2, steps=1) src.collapse_tripwire = _Wire() src._on_step_boundary(0, {"entropy": 0.0, "lengths/response_mean": 0.0}, padded=True) assert seen == [] src._on_step_boundary(0, {"entropy": 1.0, "lengths/response_mean": 3.0}) assert seen == [(0, 1.0, 3.0)] # ---------------------------------------------------------------------------- U04-3 (2 real ranks) def _wait_worker(rank, world, init_file, mode, q): import datetime import torch.distributed as dist os.environ["ALIGN_PG_TIMEOUT_S"] = "3" os.environ["WORLD_SIZE"] = str(world) os.environ["RANK"] = str(rank) os.environ["ALIGN_ROLLOUT_WAIT_TIMEOUT_S"] = "60" dist.init_process_group( "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world, timeout=datetime.timedelta(seconds=3), ) class _Src(GrpoBatchSource): def _produce_step(self, optimizer_step): if mode == "slow": time.sleep(7) # longer than the 3 s PG timeout, far shorter than the rollout budget return _payload(4) raise RuntimeError("rollout engine died (test)") src = _Src( [{"problem": "p", "answer": "1"}], prompts_per_step=1, n_samples_per_prompt=4, max_response_len=4, temperature=1.0, optimizer_steps=1, grad_accum=1, seed=0, prompt_builder=lambda p: p, ) src.bind(client=None, reward=None, verifier=None, old_logprob_fn=lambda i, a: None, weights_dir_fn=lambda: "w", rank=rank, world_size=world) src.old_logprob_fn = lambda ids, attn: torch.zeros(ids.shape[0], ids.shape[1] - 1) t0 = time.time() try: batch = src._materialize_micro(None, RolloutMicroHandle(0, 0)) q.put((rank, "ok", int(batch.input_ids.shape[0]), time.time() - t0)) except Exception as e: # noqa: BLE001 q.put((rank, "err", str(e), time.time() - t0)) def _run_two_ranks(mode): import torch.multiprocessing as mp ctx = mp.get_context("spawn") q = ctx.Queue() with tempfile.TemporaryDirectory(dir="/apps/development/rlhf/align/.tmp/audit_2026-10-14/scratch_g4_rl_batch") as d: f = os.path.join(d, "store") ps = [ctx.Process(target=_wait_worker, args=(r, 2, f, mode, q)) for r in range(2)] for p in ps: p.start() res = {} for _ in range(2): r = q.get(timeout=90) res[r[0]] = r for p in ps: p.join(30) return res def test_u04_3_peers_survive_a_rank0_rollout_longer_than_the_pg_timeout(): res = _run_two_ranks("slow") assert res[0][1] == "ok" and res[1][1] == "ok", res assert res[1][2] == 2 # 4 rows over 2 ranks def test_u04_3_a_rank0_failure_reaches_the_peers_as_a_named_error_not_a_timeout(): res = _run_two_ranks("fail") assert res[0][1] == "err" and "rollout engine died (test)" in res[0][2], res assert res[1][1] == "err" and "rollout engine died (test)" in res[1][2] and "rank 0" in res[1][2], res assert res[1][3] < 30, "the peer must be told promptly, not wait out a timeout"