Download code/scripts/dedup_arm_b_shards.py from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 13 kB
-
https://huggingface.co/viai957/align/resolve/main/code/scripts/dedup_arm_b_shards.py
- Command line
-
hf download hf://viai957/align/code/scripts/dedup_arm_b_shards.py
-
curl -L -o dedup_arm_b_shards.py https://huggingface.co/viai957/align/resolve/main/code/scripts/dedup_arm_b_shards.py
13 kB
| """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()) | |