| """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) |
| 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)] |
|
|
|
|
| |
|
|
|
|
| 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) |
| 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: |
| 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 |
|
|
|
|
| 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" |
|
|
