File size: 4,968 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Tests for the rollout assembly seam (fake verifier; no model, no engine)."""
import enum
import pytest
import torch
from align.rl.advantages import GroupFilterStats
from align.rl.assembly import assemble_rollout
from align.rl.rewards import CompositeReward
class _Outcome(enum.Enum):
CORRECT = 1
WRONG = 2
def _verify(golds, preds):
return [
_Outcome.CORRECT if (p is not None and p.strip() == g) else _Outcome.WRONG
for g, p in zip(golds, preds)
]
def _text(answer: str) -> str:
return f"Step 1: reason.\n\n# Answer\n\n{answer}"
def _mk(n_prompts=2, g=2, answers=None):
"""Two prompts x G completions with controllable answers."""
prompt_ids = [[1, 2, 3], [4, 5]]
completions = [[[10, 11], [12, 13, 14]], [[20], [21, 22]]]
logprobs = [[[-0.1, -0.2], [-0.3, -0.4, -0.5]], [[-0.6], [-0.7, -0.8]]]
answers = answers or [["7", "8"], ["9", "9"]]
texts = [[_text(a) for a in row] for a in [None] for row in answers]
return prompt_ids, completions, logprobs, texts
def test_assembles_shapes_masks_and_alignment():
prompt_ids, completions, logprobs, texts = _mk(answers=[["7", "8"], ["9", "1"]])
out = assemble_rollout(
prompt_ids,
completions,
logprobs,
texts,
golds=["7", "9"],
reward=CompositeReward(format_weight=0.1),
verify_batch=_verify,
pad_token_id=0,
)
assert out is not None
B, T = out.input_ids.shape
assert B == 4 and T == 6 # longest row: prompt 3 + completion 3
# row 0: prompt [1,2,3] + completion [10,11]
assert out.input_ids[0].tolist() == [1, 2, 3, 10, 11, 0]
assert out.response_mask[0].tolist() == [0, 0, 0, 1, 1, 0]
# label-aligned engine logprobs: completion token t -> position len(p)+t-1
assert out.vllm_logprobs[0, 2].item() == pytest.approx(-0.1)
assert out.vllm_logprobs[0, 3].item() == pytest.approx(-0.2)
assert out.vllm_logprobs[0, 1].item() == 0.0 # prompt region zero-filled
# group ids: two kept groups
assert out.group_ids.tolist() == [0, 0, 1, 1]
# advantages: group means subtracted -> each group sums to ~0
a = out.advantages
assert abs(float(a[0] + a[1])) < 1e-6 and abs(float(a[2] + a[3])) < 1e-6
# rewards itemised
assert set(out.rewards) >= {"reward/accuracy", "reward/format", "reward/total"}
assert out.rewards["reward/accuracy"].tolist() == [1.0, 0.0, 1.0, 0.0]
def test_dead_groups_dropped_whole_with_telemetry():
# group 0: both wrong (all_zero-ish variance 0 at total level? totals equal)
# -> dead; group 1: mixed -> kept
prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["9", "2"]])
stats = GroupFilterStats(max_possible_score=1.1) # acc 1.0 + fmt 0.1
out = assemble_rollout(
prompt_ids,
completions,
logprobs,
texts,
golds=["7", "9"],
reward=CompositeReward(format_weight=0.1),
verify_batch=_verify,
pad_token_id=0,
filter_stats=stats,
)
assert out is not None
assert out.kept_groups == 1 and out.dropped_groups == 1
assert out.input_ids.shape[0] == 2 # only group 1's two rows
# both-wrong completions still earn the format term (0.1 each), so the dead
# group's identical totals are NONZERO -> nonzero_std0, not all_zero. With a
# welded soft-format reward, all_zero is reserved for unformatted garbage.
assert stats.nonzero_std0 == 1 and stats.kept == 1
def test_all_dead_returns_none_not_error():
prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["2", "2"]])
out = assemble_rollout(
prompt_ids,
completions,
logprobs,
texts,
golds=["7", "9"],
reward=CompositeReward(format_weight=0.1),
verify_batch=_verify,
pad_token_id=0,
)
assert out is None
def test_ragged_logprobs_refused():
prompt_ids, completions, logprobs, texts = _mk()
logprobs[0][0] = [-0.1] # 1 logprob for 2 completion tokens
with pytest.raises(ValueError, match="ragged|logprobs"):
assemble_rollout(
prompt_ids,
completions,
logprobs,
texts,
golds=["7", "9"],
reward=CompositeReward(format_weight=0.1),
verify_batch=_verify,
pad_token_id=0,
)
def test_unanswerable_text_scores_zero_accuracy():
prompt_ids, completions, logprobs, _ = _mk()
texts = [["no answer marker here", _text("8")], [_text("9"), _text("9")]]
out = assemble_rollout(
prompt_ids,
completions,
logprobs,
texts,
golds=["8", "9"],
reward=CompositeReward(format_weight=0.1),
verify_batch=_verify,
pad_token_id=0,
)
assert out is not None
# group 0 kept (one right, one no-answer); the no-answer row scores acc 0
row0_acc = out.rewards["reward/accuracy"][0].item()
assert row0_acc == 0.0
|