"""One-shot, surgical dedup of the Arm B preference shards. WHY THIS EXISTS. The ref-logprob cache build refused the Arm B shards with `ValueError: duplicate pair_id` -- correctly. The train split held 11,404 rows with 87 pair_ids present exactly twice, and held_out held 596 rows with 2 doubled ids; every second occurrence is byte-identical to its first (same input_ids, boundaries, rejected side, and extras). Root cause: PRM800K carries multiple annotation records per problem, so two records can yield the IDENTICAL pair site; `PairCapper.offer` caps per (problem_id, step_index) but never checks pair_id, and `curate()` merges everything per problem (nmo/data/prm800k_pairs.py, `by_problem.setdefault(...).append(p)`) without a seen-pair_id check -- so the depth-round top-up re-selects a pair the first round already took. WHAT IT DOES. For each split (train, held_out): 1. read the live shards via the store's own reader; 2. keep the FIRST occurrence of every pair_id, drop second-and-later occurrences -- REFUSING unless the dropped row is content-identical to the kept one (a non-identical duplicate is a different bug and must not be silently resolved by row order); 3. write fresh shards through the same ShardWriter/receipt API the build used, preserving shard naming (`arm_b-train-NNNN` etc.). Then it asserts zero train/held pair_id overlap, atomically swaps `arm_b -> arm_b_predup` and `arm_b_dedup -> arm_b` (the predup directory is KEPT as evidence), and updates data/manifest.json in place: `counts.arm_b_curated`, the two `arm_b-*` shard hashes (same content_hash-of-receipt-hashes chain `_ShardSet.finalize` uses), and the `length_bias_baseline` Arm B entries re-measured over the deduped TRAIN split with the exact function scripts/dpo.py uses (align.train.dpo_stage.measure_chosen_longer_fraction). A provenance block `arm_b_dedup` records old -> new for everything it touched. Deterministic (pure function of the input shards; no RNG), CPU-only. The ref-logprob caches keyed on the old manifest hashes are stale after this and must be re-precomputed -- that is the point. Run: cd /apps/development/rlhf/nano-mistral-o1 CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py --dry-run CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py Refuses to run twice: an existing `arm_b_predup` means the swap already happened. """ from __future__ import annotations import os # CPU-only by construction: never let a lazily-imported torch see a GPU. os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") import argparse import json import logging from pathlib import Path import numpy as np logger = logging.getLogger("dedup_arm_b") REPO = Path("/apps/development/rlhf/nano-mistral-o1") SPLITS = (("train", "arm_b-train"), ("held_out", "arm_b-held")) #: manifest hash key per split prefix, matching build_all's `hashes` dict. MANIFEST_HASH_KEY = {"arm_b-train": "arm_b-train", "arm_b-held": "arm_b-held"} def rows_identical(a, b) -> bool: """Full content identity between two store Rows -- every stored field. Anything short of this (say, comparing input_ids only) would let a real divergence hide behind a shared pair_id and get resolved by row order. """ return ( np.array_equal(a.input_ids, b.input_ids) and np.array_equal(a.boundary_positions, b.boundary_positions) and a.problem_id == b.problem_id and a.subject == b.subject and sorted(a.ragged) == sorted(b.ragged) and all(np.array_equal(a.ragged[f], b.ragged[f]) for f in a.ragged) and a.extras == b.extras ) def dedup_row_indices(ds) -> tuple[list[int], list[int]]: """(kept, dropped) row indices: first occurrence per pair_id wins. Raises if any dropped row is NOT content-identical to its kept first occurrence -- that would be a semantic collision, not a re-added pair, and this script is only licensed to remove exact re-adds. """ first_by_pair: dict[str, int] = {} kept: list[int] = [] dropped: list[int] = [] for i in range(len(ds)): row = ds[i] pair_id = row.extras["pair_id"] if pair_id not in first_by_pair: first_by_pair[pair_id] = i kept.append(i) continue j = first_by_pair[pair_id] if not rows_identical(ds[j], row): raise ValueError( f"pair_id {pair_id!r} occurs at rows {j} and {i} with " "DIFFERENT content. That is a hash/semantics collision, not a " "duplicate re-add; refusing to pick a winner by row order." ) dropped.append(i) return kept, dropped def write_split(ds, kept: list[int], out_dir: Path, prefix: str) -> str: """Write the kept rows through the build's own shard-set API. Returns the manifest-level content hash (content_hash over the per-shard receipt hashes -- the same chain `_ShardSet.finalize` produces, so the updated manifest entry is computed by the code that defined it). """ from align.data.build import _PREF_FIELDS, _ShardSet shard_set = _ShardSet(out_dir, prefix, _PREF_FIELDS) for i in kept: row = ds[i] shard_set.add( input_ids=row.input_ids, problem_id=row.problem_id, subject=row.subject, boundary_positions=row.boundary_positions, ragged={ "rejected_input_ids": row.ragged["rejected_input_ids"], "rejected_boundary_positions": row.ragged["rejected_boundary_positions"], }, **row.extras, ) return shard_set.finalize() def step_basis_fraction(rows) -> float: """Arm B's step-basis chosen-longer fraction over collator-shaped rows. Mirrors build.py's `chosen_longer_fraction_step_tokens`: length measured from the SHARED PREFIX (the one step that differs), counted over every row. Kept next to the completion-basis measurement so the manifest's two Arm B baselines stay siblings of the same dataset state. """ longer = sum( int( len(r["chosen_input_ids"]) - int(r["shared_prefix_len"]) > len(r["rejected_input_ids"]) - int(r["shared_prefix_len"]) ) for r in rows ) return longer / len(rows) if rows else float("nan") def main(argv=None) -> int: logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) ap.add_argument("--root", type=Path, default=REPO, help="repo root") ap.add_argument( "--dry-run", action="store_true", help="measure and report; write and swap nothing" ) args = ap.parse_args(argv) from align.data.loader import preference_dataset from align.data.store import ShardDataset shards_root = args.root / "data" / "shards" live = shards_root / "arm_b" staged = shards_root / "arm_b_dedup" predup = shards_root / "arm_b_predup" manifest_path = args.root / "data" / "manifest.json" if predup.exists(): raise SystemExit( f"{predup} already exists: the swap already happened. This script " "is one-shot; a second run would dedup the already-deduped shards " "and overwrite the evidence." ) if staged.exists(): raise SystemExit( f"{staged} already exists (a previous run died mid-write?). " "Inspect and delete it before re-running; a half-written shard " "set masquerades as resumable." ) # -- pass 1: dedup into the staging directory --------------------------- pair_ids: dict[str, set[str]] = {} before_after: dict[str, tuple[int, int]] = {} new_hashes: dict[str, str] = {} for split, prefix in SPLITS: ds = ShardDataset(live / split) kept, dropped = dedup_row_indices(ds) before_after[split] = (len(ds), len(kept)) pair_ids[split] = {ds[i].extras["pair_id"] for i in kept} logger.info( "%s: %d rows -> %d kept (%d exact-duplicate re-adds dropped)", split, len(ds), len(kept), len(dropped), ) if not args.dry_run: new_hashes[prefix] = write_split(ds, kept, staged / split, prefix) overlap = pair_ids["train"] & pair_ids["held_out"] if overlap: raise SystemExit( f"{len(overlap)} pair_ids appear in BOTH deduped splits (e.g. " f"{sorted(overlap)[:3]}); the train/held split is broken and a " "dedup must not paper over it." ) if args.dry_run: logger.info("dry run: nothing written, nothing swapped") return 0 # -- pass 2: verify the staged shards through a fresh reader ------------ for split, _prefix in SPLITS: staged_ds = ShardDataset(staged / split) staged_ids = [staged_ds[i].extras["pair_id"] for i in range(len(staged_ds))] if len(staged_ids) != len(set(staged_ids)) or len(staged_ids) != before_after[split][1]: raise SystemExit( f"staged {split} re-read: {len(staged_ids)} rows, " f"{len(set(staged_ids))} unique -- the write did not produce " "what the dedup computed; refusing to swap." ) # -- pass 3: atomic swap; the pre-dedup shards are KEPT as evidence ----- live.rename(predup) staged.rename(live) logger.info("swapped: %s -> %s, %s -> %s", "arm_b", "arm_b_predup", "arm_b_dedup", "arm_b") # -- pass 4: manifest update, through the numbers' own definitions ------ manifest = json.loads(manifest_path.read_text()) old = { "counts.arm_b_curated": manifest["counts"]["arm_b_curated"], "hashes.arm_b-train": manifest["hashes"]["arm_b-train"], "hashes.arm_b-held": manifest["hashes"]["arm_b-held"], "arm_b_chosen_longer_fraction": manifest["length_bias_baseline"][ "arm_b_chosen_longer_fraction" ], "arm_b_chosen_longer_fraction_step_basis": manifest["length_bias_baseline"][ "arm_b_chosen_longer_fraction_step_basis" ], } # Re-measure the length-bias baseline over the deduped TRAIN split with # the exact function scripts/dpo.py falls back to, on the exact rows it # would build (read-only imports; nothing here re-implements the metric). from align.train.dpo_stage import measure_chosen_longer_fraction rows = preference_dataset(live / "train") new_fraction = measure_chosen_longer_fraction(rows) new_fraction_step = step_basis_fraction(rows) n_train, n_held = before_after["train"][1], before_after["held_out"][1] manifest["counts"]["arm_b_curated"] = n_train + n_held manifest["hashes"]["arm_b-train"] = new_hashes["arm_b-train"] manifest["hashes"]["arm_b-held"] = new_hashes["arm_b-held"] manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction"] = new_fraction manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction_step_basis"] = new_fraction_step manifest["arm_b_dedup"] = { "reason": ( "87 train + 2 held_out pair_ids were present exactly twice, every " "second occurrence content-identical to its first (curation " "depth-round top-up re-added pairs already selected; the merge in " "prm800k_pairs.curate never deduped by pair_id). The ref-logprob " "cache build correctly refused the duplicates." ), "script": "scripts/dedup_arm_b_shards.py", "train_rows": {"before": before_after["train"][0], "after": n_train}, "held_out_rows": {"before": before_after["held_out"][0], "after": n_held}, "old": old, "predup_shards": str(predup), } manifest_path.write_text(json.dumps(manifest, indent=1)) logger.info("train: %d -> %d rows", before_after["train"][0], n_train) logger.info("held_out: %d -> %d rows", before_after["held_out"][0], n_held) logger.info("counts.arm_b_curated: %s -> %d", old["counts.arm_b_curated"], n_train + n_held) logger.info( "arm_b_chosen_longer_fraction: %s -> %.6f (step basis %s -> %.6f)", old["arm_b_chosen_longer_fraction"], new_fraction, old["arm_b_chosen_longer_fraction_step_basis"], new_fraction_step, ) logger.info( "hashes: arm_b-train %s -> %s, arm_b-held %s -> %s", old["hashes.arm_b-train"], new_hashes["arm_b-train"], old["hashes.arm_b-held"], new_hashes["arm_b-held"], ) logger.info( "NOTE: ref caches keyed on the old manifest hashes are now stale; " "re-run scripts/precompute_ref_logprobs.py before scripts/dpo.py." ) return 0 if __name__ == "__main__": # House rule: everything above is import-safe because spawn-based worker # pools re-import __main__; the entry point stays behind this guard. raise SystemExit(main())