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

align / code /tests /test_dedup_arm_b.py

Add files using upload-large-folder tool

201dc77 verified 16 days ago

7.51 kB

"""Arm B shard dedup: the permanent invariant, and the dedup script's logic.
The ref-logprob cache build refused the original Arm B shards with
`ValueError: duplicate pair_id`: 87 train and 2 held_out pair_ids were
present exactly twice (curation's depth-round top-up re-added already
selected pairs; the per-problem merge in `prm800k_pairs.curate` never
deduped by pair_id). `scripts/dedup_arm_b_shards.py` removed the exact
re-adds. From now on, unique pair_ids per split -- and zero train/held
overlap -- is a data invariant, not a hope.
"""

import importlib.util
import os
import sys
from pathlib import Path

import numpy as np
import pytest

REPO = Path("/apps/development/rlhf/nano-mistral-o1")
ARM_B = REPO / "data" / "shards" / "arm_b"


def _load_dedup_module():
    """Import the script by path; scripts/ is not a package."""
    path = REPO / "scripts" / "dedup_arm_b_shards.py"
    spec = importlib.util.spec_from_file_location("dedup_arm_b_shards", path)
    mod = importlib.util.module_from_spec(spec)
    # The script pins CUDA_VISIBLE_DEVICES="" for its own runs; do not let
    # importing it here clobber the test session's environment.
    saved = os.environ.get("CUDA_VISIBLE_DEVICES")
    try:
        spec.loader.exec_module(mod)
    finally:
        if saved is None:
            os.environ.pop("CUDA_VISIBLE_DEVICES", None)
        else:
            os.environ["CUDA_VISIBLE_DEVICES"] = saved
    sys.modules.pop("dedup_arm_b_shards", None)
    return mod


# ---------------------------------------------------------------------------
# (a) Post-condition on the LIVE shards: the permanent invariant
# ---------------------------------------------------------------------------


def _split_pair_ids(split: str) -> list[str]:
    from align.data.store import ShardDataset

    ds = ShardDataset(ARM_B / split)
    return [ds[i].extras["pair_id"] for i in range(len(ds))]


@pytest.mark.skipif(not ARM_B.is_dir(), reason=f"arm_b shards not present at {ARM_B}")
def test_live_arm_b_shards_have_unique_pair_ids_and_disjoint_splits():
    """STRICT: every pair_id unique within each split, zero overlap across.
    A duplicate here means either the dedup regressed or a rebuild
    reintroduced the curation top-up defect; both must fail at CPU-test time,
    before the ref cache refuses them after GPUs have been claimed.
    """
    by_split = {}
    for split in ("train", "held_out"):
        ids = _split_pair_ids(split)
        dupes = sorted({p for p in ids if ids.count(p) > 1}) if len(set(ids)) < len(ids) else []
        assert len(set(ids)) == len(ids), (
            f"arm_b/{split}: {len(ids) - len(set(ids))} duplicate pair_id "
            f"occurrences (e.g. {dupes[:3]}). The dedup invariant is broken; "
            "see scripts/dedup_arm_b_shards.py."
        )
        by_split[split] = set(ids)

    overlap = by_split["train"] & by_split["held_out"]
    assert not overlap, (
        f"{len(overlap)} pair_ids present in BOTH arm_b splits (e.g. "
        f"{sorted(overlap)[:3]}); held-out preference accuracy would measure "
        "memorization while looking like generalization."
    )


@pytest.mark.skipif(
    not (REPO / "data" / "manifest.json").is_file(), reason="build manifest not present"
)
def test_manifest_arm_b_count_matches_live_shards():
    """counts.arm_b_curated must equal what is actually on disk: the manifest
    gate derives the optimizer-step horizon from it."""
    import json

    from align.data.store import ShardDataset

    if not ARM_B.is_dir():
        pytest.skip(f"arm_b shards not present at {ARM_B}")
    manifest = json.loads((REPO / "data" / "manifest.json").read_text())
    n = sum(len(ShardDataset(ARM_B / s)) for s in ("train", "held_out"))
    assert manifest["counts"]["arm_b_curated"] == n


# ---------------------------------------------------------------------------
# (b) The dedup script's kept-first / content-identity logic, synthetically
# ---------------------------------------------------------------------------


def _write_tiny_shard(tmp_path: Path, rows: list[dict]) -> Path:
    from align.data.store import ShardWriter

    fields = {
        "input_ids": "int32",
        "rejected_input_ids": "int32",
        "rejected_boundary_positions": "int32",
    }
    with ShardWriter(tmp_path, "tiny-0000", fields) as w:
        for r in rows:
            w.add_row(
                r["input_ids"],
                problem_id=r["problem_id"],
                subject="algebra",
                boundary_positions=r["boundary_positions"],
                ragged={
                    "rejected_input_ids": r["rejected_input_ids"],
                    "rejected_boundary_positions": [0],
                },
                chosen_prompt_len=2,
                rejected_prompt_len=2,
                shared_prefix_len=3,
                step_index=r.get("step_index", 1),
                pair_id=r["pair_id"],
                source="tiny",
            )
    return tmp_path


def _row(pair_id, input_ids=(5, 6, 7, 8), rejected=(5, 6, 9), problem_id="p1", **kw):
    n = len(input_ids)
    return {
        "pair_id": pair_id,
        "input_ids": list(input_ids),
        "rejected_input_ids": list(rejected),
        "boundary_positions": [1, n - 1] if n > 2 else [n - 1],
        "problem_id": problem_id,
        **kw,
    }


def test_dedup_keeps_first_and_drops_identical_re_adds(tmp_path):
    from align.data.store import ShardDataset

    dedup = _load_dedup_module()
    # a, b, a-again (identical), c, b-again (identical)
    rows = [
        _row("pair_a"),
        _row("pair_b", input_ids=(1, 2, 3, 4)),
        _row("pair_a"),
        _row("pair_c", input_ids=(9, 9, 9)),
        _row("pair_b", input_ids=(1, 2, 3, 4)),
    ]
    ds = ShardDataset(_write_tiny_shard(tmp_path, rows))
    kept, dropped = dedup.dedup_row_indices(ds)
    assert kept == [0, 1, 3], "first occurrence wins, in original order"
    assert dropped == [2, 4]
    assert [ds[i].extras["pair_id"] for i in kept] == ["pair_a", "pair_b", "pair_c"]


def test_dedup_refuses_non_identical_duplicate(tmp_path):
    """A shared pair_id with DIFFERENT content is a collision, not a re-add;
    the script must refuse rather than resolve it by row order."""
    from align.data.store import ShardDataset

    dedup = _load_dedup_module()
    rows = [
        _row("pair_a"),
        _row("pair_a", input_ids=(5, 6, 7, 8, 8)),  # same id, different tokens
    ]
    ds = ShardDataset(_write_tiny_shard(tmp_path, rows))
    with pytest.raises(ValueError, match="DIFFERENT content"):
        dedup.dedup_row_indices(ds)


def test_rows_identical_covers_every_stored_field(tmp_path):
    """The identity check must see extras and the rejected side, not just
    input_ids -- otherwise a real divergence hides behind a shared pair_id."""
    from align.data.store import ShardDataset

    dedup = _load_dedup_module()
    ds = ShardDataset(
        _write_tiny_shard(
            tmp_path,
            [
                _row("pair_a"),
                _row("pair_b"),  # same tokens as pair_a, different pair_id extra
                _row("pair_c", rejected=(5, 6, 1)),  # differs only on the rejected side
            ],
        )
    )
    assert dedup.rows_identical(ds[0], ds[0])
    assert not dedup.rows_identical(ds[0], ds[1]), "extras differ (pair_id)"
    assert not dedup.rows_identical(ds[0], ds[2]), "rejected_input_ids differ"
    assert np.array_equal(ds[0].input_ids, ds[1].input_ids), "sanity: tokens do match"

	"""Arm B shard dedup: the permanent invariant, and the dedup script's logic.

	The ref-logprob cache build refused the original Arm B shards with
	`ValueError: duplicate pair_id`: 87 train and 2 held_out pair_ids were
	present exactly twice (curation's depth-round top-up re-added already
	selected pairs; the per-problem merge in `prm800k_pairs.curate` never
	deduped by pair_id). `scripts/dedup_arm_b_shards.py` removed the exact
	re-adds. From now on, unique pair_ids per split -- and zero train/held
	overlap -- is a data invariant, not a hope.
	"""

	import importlib.util
	import os
	import sys
	from pathlib import Path

	import numpy as np
	import pytest

	REPO = Path("/apps/development/rlhf/nano-mistral-o1")
	ARM_B = REPO / "data" / "shards" / "arm_b"


	def _load_dedup_module():
	"""Import the script by path; scripts/ is not a package."""
	path = REPO / "scripts" / "dedup_arm_b_shards.py"
	spec = importlib.util.spec_from_file_location("dedup_arm_b_shards", path)
	mod = importlib.util.module_from_spec(spec)
	# The script pins CUDA_VISIBLE_DEVICES="" for its own runs; do not let
	# importing it here clobber the test session's environment.
	saved = os.environ.get("CUDA_VISIBLE_DEVICES")
	try:
	spec.loader.exec_module(mod)
	finally:
	if saved is None:
	os.environ.pop("CUDA_VISIBLE_DEVICES", None)
	else:
	os.environ["CUDA_VISIBLE_DEVICES"] = saved
	sys.modules.pop("dedup_arm_b_shards", None)
	return mod


	# ---------------------------------------------------------------------------
	# (a) Post-condition on the LIVE shards: the permanent invariant
	# ---------------------------------------------------------------------------


	def _split_pair_ids(split: str) -> list[str]:
	from align.data.store import ShardDataset

	ds = ShardDataset(ARM_B / split)
	return [ds[i].extras["pair_id"] for i in range(len(ds))]


	@pytest.mark.skipif(not ARM_B.is_dir(), reason=f"arm_b shards not present at {ARM_B}")
	def test_live_arm_b_shards_have_unique_pair_ids_and_disjoint_splits():
	"""STRICT: every pair_id unique within each split, zero overlap across.

	A duplicate here means either the dedup regressed or a rebuild
	reintroduced the curation top-up defect; both must fail at CPU-test time,
	before the ref cache refuses them after GPUs have been claimed.
	"""
	by_split = {}
	for split in ("train", "held_out"):
	ids = _split_pair_ids(split)
	dupes = sorted({p for p in ids if ids.count(p) > 1}) if len(set(ids)) < len(ids) else []
	assert len(set(ids)) == len(ids), (
	f"arm_b/{split}: {len(ids) - len(set(ids))} duplicate pair_id "
	f"occurrences (e.g. {dupes[:3]}). The dedup invariant is broken; "
	"see scripts/dedup_arm_b_shards.py."
	)
	by_split[split] = set(ids)

	overlap = by_split["train"] & by_split["held_out"]
	assert not overlap, (
	f"{len(overlap)} pair_ids present in BOTH arm_b splits (e.g. "
	f"{sorted(overlap)[:3]}); held-out preference accuracy would measure "
	"memorization while looking like generalization."
	)


	@pytest.mark.skipif(
	not (REPO / "data" / "manifest.json").is_file(), reason="build manifest not present"
	)
	def test_manifest_arm_b_count_matches_live_shards():
	"""counts.arm_b_curated must equal what is actually on disk: the manifest
	gate derives the optimizer-step horizon from it."""
	import json

	from align.data.store import ShardDataset

	if not ARM_B.is_dir():
	pytest.skip(f"arm_b shards not present at {ARM_B}")
	manifest = json.loads((REPO / "data" / "manifest.json").read_text())
	n = sum(len(ShardDataset(ARM_B / s)) for s in ("train", "held_out"))
	assert manifest["counts"]["arm_b_curated"] == n


	# ---------------------------------------------------------------------------
	# (b) The dedup script's kept-first / content-identity logic, synthetically
	# ---------------------------------------------------------------------------


	def _write_tiny_shard(tmp_path: Path, rows: list[dict]) -> Path:
	from align.data.store import ShardWriter

	fields = {
	"input_ids": "int32",
	"rejected_input_ids": "int32",
	"rejected_boundary_positions": "int32",
	}
	with ShardWriter(tmp_path, "tiny-0000", fields) as w:
	for r in rows:
	w.add_row(
	r["input_ids"],
	problem_id=r["problem_id"],
	subject="algebra",
	boundary_positions=r["boundary_positions"],
	ragged={
	"rejected_input_ids": r["rejected_input_ids"],
	"rejected_boundary_positions": [0],
	},
	chosen_prompt_len=2,
	rejected_prompt_len=2,
	shared_prefix_len=3,
	step_index=r.get("step_index", 1),
	pair_id=r["pair_id"],
	source="tiny",
	)
	return tmp_path


	def _row(pair_id, input_ids=(5, 6, 7, 8), rejected=(5, 6, 9), problem_id="p1", **kw):
	n = len(input_ids)
	return {
	"pair_id": pair_id,
	"input_ids": list(input_ids),
	"rejected_input_ids": list(rejected),
	"boundary_positions": [1, n - 1] if n > 2 else [n - 1],
	"problem_id": problem_id,
	**kw,
	}


	def test_dedup_keeps_first_and_drops_identical_re_adds(tmp_path):
	from align.data.store import ShardDataset

	dedup = _load_dedup_module()
	# a, b, a-again (identical), c, b-again (identical)
	rows = [
	_row("pair_a"),
	_row("pair_b", input_ids=(1, 2, 3, 4)),
	_row("pair_a"),
	_row("pair_c", input_ids=(9, 9, 9)),
	_row("pair_b", input_ids=(1, 2, 3, 4)),
	]
	ds = ShardDataset(_write_tiny_shard(tmp_path, rows))
	kept, dropped = dedup.dedup_row_indices(ds)
	assert kept == [0, 1, 3], "first occurrence wins, in original order"
	assert dropped == [2, 4]
	assert [ds[i].extras["pair_id"] for i in kept] == ["pair_a", "pair_b", "pair_c"]


	def test_dedup_refuses_non_identical_duplicate(tmp_path):
	"""A shared pair_id with DIFFERENT content is a collision, not a re-add;
	the script must refuse rather than resolve it by row order."""
	from align.data.store import ShardDataset

	dedup = _load_dedup_module()
	rows = [
	_row("pair_a"),
	_row("pair_a", input_ids=(5, 6, 7, 8, 8)), # same id, different tokens
	]
	ds = ShardDataset(_write_tiny_shard(tmp_path, rows))
	with pytest.raises(ValueError, match="DIFFERENT content"):
	dedup.dedup_row_indices(ds)


	def test_rows_identical_covers_every_stored_field(tmp_path):
	"""The identity check must see extras and the rejected side, not just
	input_ids -- otherwise a real divergence hides behind a shared pair_id."""
	from align.data.store import ShardDataset

	dedup = _load_dedup_module()
	ds = ShardDataset(
	_write_tiny_shard(
	tmp_path,
	[
	_row("pair_a"),
	_row("pair_b"), # same tokens as pair_a, different pair_id extra
	_row("pair_c", rejected=(5, 6, 1)), # differs only on the rejected side
	],
	)
	)
	assert dedup.rows_identical(ds[0], ds[0])
	assert not dedup.rows_identical(ds[0], ds[1]), "extras differ (pair_id)"
	assert not dedup.rows_identical(ds[0], ds[2]), "rejected_input_ids differ"
	assert np.array_equal(ds[0].input_ids, ds[1].input_ids), "sanity: tokens do match"