File size: 7,508 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 | """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"
|