| """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 |
|
|
| |
| 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 = {"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." |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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." |
| ) |
|
|
| |
| live.rename(predup) |
| staged.rename(live) |
| logger.info("swapped: %s -> %s, %s -> %s", "arm_b", "arm_b_predup", "arm_b_dedup", "arm_b") |
|
|
| |
| 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" |
| ], |
| } |
|
|
| |
| |
| |
| 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__": |
| |
| |
| raise SystemExit(main()) |
|
|
