File size: 9,151 Bytes
201dc77 | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """Length-grouped batching + Preference collate.
Run: uv run pytest tests/test_data_loader.py -m "not gpu" -v
Padding waste is not a minor inefficiency here: the per-micro-batch cost is
dominated by a fixed host-staged all-gather, so wasted padding is wasted
wall-clock in direct proportion. And a DPO batch that reports 0 for an
undefined statistic is indistinguishable from one reporting a real 0 -- every
undefined statistic in this module is NaN.
"""
from __future__ import annotations
import math
import numpy as np
import pytest
from align.data.loader import (
LengthGroupedSampler,
Preference,
collate_preference,
preference_length_key,
)
PAD = 0
def _pair_row(
prompt_len: int, chosen_completion: int, rejected_completion: int, pid: str = "p0"
) -> dict:
"""One preference row in the on-disk shape the pair shards produce."""
chosen = list(range(10, 10 + prompt_len + chosen_completion))
rejected = chosen[:prompt_len] + list(range(500, 500 + rejected_completion))
return {
"chosen_input_ids": np.asarray(chosen, dtype=np.int32),
"rejected_input_ids": np.asarray(rejected, dtype=np.int32),
"chosen_prompt_len": prompt_len,
"rejected_prompt_len": prompt_len,
"problem_id": pid,
"subject": "Algebra",
}
def test_length_grouped_batches_bound_padding_waste():
"""Batches drawn from a synthetic length distribution have intra-batch
length spread below the configured bound (grouping actually groups).
Both directions: the same distribution batched WITHOUT grouping blows the
bound, so the assertion is measuring grouping and not the distribution."""
rng = np.random.default_rng(0)
lengths = rng.integers(64, 1024, size=2000)
grouped = LengthGroupedSampler(lengths, max_tokens=8192, max_rows=8, seed=0, shuffle=True)
spreads = [int(lengths[b].max() - lengths[b].min()) for b in grouped if len(b) > 1]
assert spreads, "sampler produced no multi-row batches"
assert max(spreads) <= 128, f"intra-batch spread too large: {max(spreads)}"
# Firing direction: ungrouped batches of the same size are far worse.
order = rng.permutation(len(lengths))
naive = [order[i : i + 8] for i in range(0, len(order), 8)]
naive_spread = max(int(lengths[b].max() - lengths[b].min()) for b in naive)
assert naive_spread > 128
# A token budget, not a row count: padded cost stays under the budget.
for b in grouped:
assert int(lengths[b].max()) * len(b) <= 8192 or len(b) == 1
def test_length_grouping_is_seeded_and_reproducible():
"""Same seed => identical batch order across two loader constructions
(required by the section-07 seeded-resume contract). Different seed =>
different order, so the test is not passing on a degenerate sampler."""
lengths = np.random.default_rng(1).integers(64, 1024, size=500)
a = [list(map(int, b)) for b in LengthGroupedSampler(lengths, 4096, seed=7)]
b = [list(map(int, x)) for x in LengthGroupedSampler(lengths, 4096, seed=7)]
c = [list(map(int, x)) for x in LengthGroupedSampler(lengths, 4096, seed=8)]
assert a == b
assert a != c
# set_epoch changes the order deterministically, and reproducibly so.
s1 = LengthGroupedSampler(lengths, 4096, seed=7)
s1.set_epoch(1)
s2 = LengthGroupedSampler(lengths, 4096, seed=7)
s2.set_epoch(1)
assert [list(map(int, x)) for x in s1] == [list(map(int, x)) for x in s2]
assert [list(map(int, x)) for x in s1] != a
def test_preference_collate_shapes_and_masks():
"""Preference batch exposes chosen/rejected ids, attention masks, prompt
lengths, and per-row metadata; completion-token masks exclude prompt
positions.
Incident class: a completion mask that includes prompt positions makes the
DPO logprob a function of the (identical) prompt, which cancels in the
difference for shared-prefix pairs and quietly shrinks the signal."""
rows = [_pair_row(5, 3, 4, "pA"), _pair_row(7, 9, 2, "pB")]
batch = collate_preference(rows, pad_token_id=PAD)
assert isinstance(batch, Preference)
assert batch.chosen_input_ids.shape == (2, 16) # max(5+3, 7+9)
assert batch.rejected_input_ids.shape == (2, 9) # max(5+4, 7+2)
assert batch.chosen_attention_mask.sum().item() == (5 + 3) + (7 + 9)
assert batch.rejected_attention_mask.sum().item() == (5 + 4) + (7 + 2)
# Completion mask: excludes prompt positions AND padding.
assert batch.chosen_completion_mask[0].tolist() == [0] * 5 + [1] * 3 + [0] * 8
assert batch.rejected_completion_mask[1].tolist() == [0] * 7 + [1] * 2
assert batch.chosen_prompt_lens.tolist() == [5, 7]
assert batch.rejected_prompt_lens.tolist() == [5, 7]
assert [m["problem_id"] for m in batch.metadata] == ["pA", "pB"]
stats = batch.stats()
assert stats["n_pairs"] == 2
assert stats["mean_chosen_completion_tokens"] == pytest.approx((3 + 9) / 2)
assert stats["mean_rejected_completion_tokens"] == pytest.approx((4 + 2) / 2)
assert stats["chosen_longer_fraction"] == pytest.approx(0.5)
# Grouping key for preference data is max(len(chosen), len(rejected)).
assert preference_length_key(rows[0]) == 9
assert preference_length_key(rows[1]) == 16
def test_preference_collate_empty_stats_are_nan():
"""Any statistic undefined on an empty/degenerate batch is NaN, never 0.
Incident: a 0 that means 'no data' is indistinguishable from a 0 that means
'measured zero' -- the predecessor's reward channel read as an honest
measurement of a weak policy for two full runs."""
empty = collate_preference([], pad_token_id=PAD)
s = empty.stats()
assert s["n_pairs"] == 0
for key in (
"mean_chosen_completion_tokens",
"mean_rejected_completion_tokens",
"chosen_longer_fraction",
"mean_prompt_tokens",
):
assert math.isnan(s[key]), f"{key} must be NaN on an empty batch, got {s[key]}"
# Degenerate: rows exist but no completion tokens at all.
degenerate = collate_preference([_pair_row(4, 0, 0)], pad_token_id=PAD)
ds = degenerate.stats()
assert ds["n_pairs"] == 1
assert ds["mean_prompt_tokens"] == pytest.approx(4.0)
assert math.isnan(ds["mean_chosen_completion_tokens"])
assert math.isnan(ds["chosen_longer_fraction"])
# Silent direction: a normal batch reports real numbers, not NaN.
ok = collate_preference([_pair_row(4, 2, 1)], pad_token_id=PAD).stats()
assert not math.isnan(ok["mean_chosen_completion_tokens"])
def test_pad_token_never_enters_the_completion_mask():
"""Right-padding only: boundary receipts are absolute indices into the
unpadded sequence, so left padding would shift every stored position.
The test pins the padding side by construction."""
rows = [_pair_row(2, 1, 1), _pair_row(2, 6, 6)]
b = collate_preference(rows, pad_token_id=PAD)
# Row 0 is the short one: real tokens first, padding at the tail.
assert b.chosen_input_ids[0, :3].tolist() != [PAD, PAD, PAD]
assert b.chosen_input_ids[0, 3:].tolist() == [PAD] * (b.chosen_input_ids.shape[1] - 3)
assert b.chosen_completion_mask[0].sum().item() == 1
def test_stored_pair_rows_collate_into_a_preference_batch(tmp_workdir):
"""Shards -> Preference batch, through the documented adapter. This is the
seam section-12 loads both DPO arms across, so it is tested rather than
assumed: the pair shards store the CHOSEN side as the primary input_ids
and the rejected side as a ragged extra, and a consumer that guessed that
mapping wrong would train on chosen-vs-chosen and report a healthy loss."""
from align.data.loader import preference_dataset
from align.data.store import ShardWriter
with ShardWriter(
tmp_workdir,
"arm-0000",
{
"input_ids": "int32",
"rejected_input_ids": "int32",
"rejected_boundary_positions": "int32",
},
) as w:
for i in range(3):
chosen = list(range(10, 20 + i))
rejected = chosen[:5] + [900, 901]
w.add_row(
chosen,
problem_id=f"p{i}",
subject="Algebra",
boundary_positions=[3, 6],
ragged={"rejected_input_ids": rejected, "rejected_boundary_positions": [3]},
chosen_prompt_len=5,
rejected_prompt_len=5,
shared_prefix_len=5,
step_index=2,
pair_id=f"pair{i}",
source="arm_b",
)
rows = preference_dataset(tmp_workdir)
assert len(rows) == 3
batch = collate_preference(rows, pad_token_id=PAD)
assert batch.chosen_input_ids.shape[0] == 3
assert batch.rejected_input_ids.shape[1] == 7
# Prompt positions are masked out of both completion masks.
assert batch.chosen_completion_mask[:, :5].sum().item() == 0
assert batch.rejected_completion_mask[0].tolist() == [0] * 5 + [1, 1]
assert [m["pair_id"] for m in batch.metadata] == ["pair0", "pair1", "pair2"]
assert not math.isnan(batch.stats()["mean_chosen_completion_tokens"])
|