code/tests/test_data_loader.py · viai957/align at main

align / code /tests /test_data_loader.py

Add files using upload-large-folder tool

201dc77 verified 16 days ago

9.15 kB

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

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