| """Amendment-3 matched pilot: immutable plan, three arms and exact block boundaries. |
| |
| Captured rounds are evaluation-only. Only an explicitly synthetic manifest can |
| use rehearsal sizes; it can never qualify the production fast path. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from dataclasses import asdict |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import signal |
| import threading |
| from typing import Literal, NotRequired, TypedDict, cast |
|
|
| import torch |
| from torch import Tensor |
|
|
| from ..live_reference import Reference |
| from .data import Block, D2Blocks, sha256 |
| from .loop import BatchPlan, Run, train |
| from .model import Drafter, load |
| from .replay import read_checkpoint |
| from .pilot_units import Units, TokenCounts, work_units, require_matched |
| from .pilot_sample import Sample, sample_plan, validate_sample, DEV_POPULATION |
|
|
| ARMS = ("REF", "FAST", "FAST-B") |
| CHECKPOINTS = (2500, 5000, 10000, 20000) |
| SEED = 23 |
| Recipe = Literal["kl", "tvd", "ace", "lr"] |
| RECIPES: tuple[Recipe, ...] = ("kl", "tvd", "ace", "lr") |
|
|
|
|
| class Spec(TypedDict): |
| schema: str |
| synthetic: bool |
| recipe: NotRequired[Recipe] |
| manifest_sha256: str |
| parity_sha256: str |
| g1_sha256: str |
| source_sha256: dict[str, str] |
| frozen_identity: str |
| initialization_sha256: str |
| run: dict[str, object] |
| checkpoints: list[int] |
| order_seed: int |
| order_seed_b: int |
| population_sha256: str |
| arm_orders: dict[str, str] |
| sample: NotRequired[Sample] |
| matched_units: dict[str, dict[str, Units]] |
|
|
|
|
| def digest(value: object) -> str: |
| return hashlib.sha256(json.dumps(value, sort_keys=True, allow_nan=False, separators=(",", ":")).encode()).hexdigest() |
|
|
|
|
| def sources() -> dict[str, str]: |
| root = Path(__file__).parent |
| files = list(root.glob("*.py")) + [root/"run_pilot_locked.sh", root.parent.parent/"native/formats.py", |
| root.parent.parent/"native/official.py", root.parent/"g1.py", root.parent/"live_reference.py"] |
| return {str(p.relative_to(root.parent.parent)): sha256(p) for p in sorted(files)} |
|
|
|
|
| def write_json(path: Path, value: object) -> None: |
| path.parent.mkdir(parents=True, mode=0o700, exist_ok=True) |
| with path.open("x") as f: |
| os.fchmod(f.fileno(), 0o600) |
| json.dump(value, f, indent=2, allow_nan=False) |
| f.write("\n") |
|
|
|
|
| def read_spec(path: Path, *, sampled: bool = False) -> Spec: |
| spec = cast(Spec, json.loads(path.read_text())) |
| validate_spec(spec, sampled=sampled) |
| return spec |
|
|
|
|
| def validate_spec(spec: Spec, *, sampled: bool = False) -> None: |
| if spec["schema"] != "dsv41.d2.pilot-plan.v1" or spec["source_sha256"] != sources(): |
| raise ValueError("pilot source identity mismatch") |
| if spec["run"] != dict(asdict(pilot_run(spec.get("recipe", "kl")))) or spec["order_seed"] != SEED or spec["order_seed_b"] != SEED+1: |
| raise ValueError("pilot matched settings mismatch") |
| points = spec["checkpoints"] |
| if (len(points) != 4 or points != sorted(set(points)) or points[0] <= 0 |
| or (not spec["synthetic"] and tuple(points) != CHECKPOINTS)): |
| raise ValueError("pilot checkpoint schedule mismatch") |
| if set(spec["arm_orders"]) != set(ARMS) or spec["arm_orders"]["REF"] != spec["arm_orders"]["FAST"]: |
| raise ValueError("pilot arm/order mismatch") |
| require_matched(spec["matched_units"]) |
| if "sample" in spec: |
| validate_sample(spec["sample"], spec["manifest_sha256"]) |
| if not spec["synthetic"] and set(spec["sample"]["keys"]) != set(DEV_POPULATION): |
| raise ValueError("real pilot requires amendment-7 dev population") |
| elif sampled: |
| raise ValueError("pilot requires its once-only pre-training sample") |
|
|
|
|
| def pilot_run(recipe: Recipe = "kl") -> Run: |
| if recipe not in RECIPES: |
| raise ValueError("unknown pilot recipe: "+recipe) |
| return Run(epochs=1, seed=SEED, chunk=8192, rank=32, batch_size=128, |
| divergence=recipe if recipe in ("tvd", "ace") else "kl", |
| lr=5e-5 if recipe == "lr" else 2e-4, warmup_fraction=.1 if recipe == "lr" else 0.) |
|
|
|
|
| def schedule(length: int, checkpoints: list[int], arm: str, recipe: Recipe = "kl") -> tuple[BatchPlan, list[int], str]: |
| if arm not in ARMS or length < checkpoints[-1]: |
| raise ValueError("pilot requires all arms and enough unique training blocks") |
| |
| |
| population = torch.randperm(length, generator=torch.Generator().manual_seed(SEED)).tolist()[:checkpoints[-1]] |
| order = population |
| if arm == "FAST-B": |
| order = [] |
| generator = torch.Generator().manual_seed(SEED+1) |
| start = 0 |
| for boundary in checkpoints: |
| segment = population[start:boundary] |
| permutation = torch.randperm(len(segment), generator=generator).tolist() |
| order.extend(segment[i] for i in permutation) |
| start = boundary |
| batches, stops = [], [] |
| start = 0 |
| for boundary in checkpoints: |
| for lo in range(start, boundary, 128): |
| batches.append(tuple(order[lo:min(lo+128, boundary)])) |
| stops.append(len(batches)) |
| start = boundary |
| identity = digest({"order": order, "boundaries": checkpoints, "run": asdict(pilot_run(recipe))}) |
| return BatchPlan(tuple(batches), tuple(stops), identity), population, digest(order) |
|
|
|
|
| def matched_units(data: TokenCounts, points: list[int], recipe: Recipe = "kl") -> dict[str, dict[str, Units]]: |
| profiles = {} |
| for arm in ARMS: |
| batch, _, _ = schedule(len(data), points, arm, recipe) |
| profiles[arm] = {str(n): work_units(data, batch, pilot_run(recipe), step) |
| for n, step in zip(points, batch.checkpoints)} |
| require_matched(profiles) |
| return profiles |
|
|
|
|
| def verify_units(spec: Spec, data: TokenCounts) -> None: |
| if spec["matched_units"] != matched_units(data, spec["checkpoints"], spec.get("recipe", "kl")): |
| raise ValueError("pilot matched-unit contract differs from actual dataset") |
|
|
|
|
| def initialization(model: Drafter) -> str: |
| h = hashlib.sha256() |
| for name, p in model.named_parameters(): |
| h.update(name.encode()) |
| h.update(p.detach().cpu().contiguous().numpy().tobytes()) |
| if name.endswith(".b") and bool(p.detach().count_nonzero()): |
| raise ValueError("pilot requires zero-init LoRA B") |
| return h.hexdigest() |
|
|
|
|
| def manifest_is_synthetic(manifest: Path) -> bool: |
| """Recognize all dataset writers' synthetic markers; any positive wins.""" |
| metadata = cast(dict[str, object], json.loads(manifest.read_text())) |
| checks = metadata.get("checks", {}) |
| generation = metadata.get("generation", {}) |
| return (metadata.get("synthetic") is True or metadata.get("taps_source") == "synthetic" |
| or (isinstance(checks, dict) and (checks.get("synthetic") is True or checks.get("taps_source") == "synthetic")) |
| or (isinstance(generation, dict) and generation.get("synthetic") is True)) |
|
|
|
|
| def validate_inputs(spec: Spec, manifest: Path, parity: Path, g1: Path) -> None: |
| for path, expected in ((manifest, spec["manifest_sha256"]), (parity, spec["parity_sha256"]), (g1, spec["g1_sha256"])): |
| if sha256(path) != expected: |
| raise ValueError("pilot manifest/parity/G1 receipt mismatch") |
| is_synthetic = manifest_is_synthetic(manifest) |
| if is_synthetic != spec["synthetic"]: |
| raise ValueError("pilot synthetic label mismatch") |
| p = json.loads(parity.read_text()) |
| if p.get("pass") is not True or p.get("samples", 0) < 50 or p.get("reference_math") is not True or p.get("frozen_identity") != spec["frozen_identity"]: |
| raise ValueError("pilot unadapted reference parity missing or mismatched") |
| gate = json.loads(g1.read_text()) |
| if not spec["synthetic"]: |
| validate_live_g1(g1, manifest) |
|
|
|
|
| def validate_live_g1(path: Path, manifest: Path) -> Reference: |
| from ..live_reference import verify_reference |
| gate = json.loads(path.read_text()) |
| binding = gate.get("g1", {}) |
| if (manifest_is_synthetic(path) or gate.get("split") != "dev" or binding.get("split") != "dev" |
| or binding.get("mode") != "live-reference" or binding.get("amendment") != "7b" |
| or binding.get("statistic") != "position0_acceptance" or binding.get("minimum_live_fraction") != .5 |
| or binding.get("tripwire_passed") is not True |
| or binding.get("pass") is not True or binding.get("binding") is not True |
| or binding.get("status") != "PASS" or gate.get("manifest_sha256") != sha256(manifest)): |
| raise ValueError("genuine amendment-7b binding dev G1(b) PASS required; legacy acap refused") |
| if not gate.get("live_reference_path") or not binding.get("live_reference_sha256"): |
| raise ValueError("missing live reference receipt/hash") |
| live = verify_reference(Path(gate["live_reference_path"])) |
| if live["status"] != "VALID" or binding["live_reference_sha256"] != live["sha256"]: |
| raise ValueError("live reference receipt hash/status mismatch") |
| from ..g1 import g1_verdict |
| if binding != g1_verdict(gate["position0_match_rate"], "dev", live_reference=live): |
| raise ValueError("G1 verdict differs from live position-0 score") |
| return live |
|
|
|
|
| def prepare(a: argparse.Namespace) -> None: |
| from .safety import assert_gpu_lease |
| assert_gpu_lease() |
| data = D2Blocks(a.manifest, "train") |
| synthetic = manifest_is_synthetic(a.manifest) |
| if a.rehearsal != synthetic: |
| raise ValueError("rehearsal flag and manifest synthetic label must match") |
| points = [25, 50, 100, 200] if synthetic else list(CHECKPOINTS) |
| torch.manual_seed(SEED) |
| model = load(a.ckpt, "cuda", rank=32, reference_math=True) |
| recipe: Recipe = a.recipe |
| plans = {arm: schedule(len(data), points, arm, recipe) for arm in ARMS} |
| spec: Spec = {"schema": "dsv41.d2.pilot-plan.v1", "synthetic": synthetic, "recipe": recipe, "matched_units": matched_units(data, points, recipe), "manifest_sha256": sha256(a.manifest), |
| "parity_sha256": sha256(a.parity_receipt), "g1_sha256": sha256(a.g1_receipt), "source_sha256": sources(), |
| "frozen_identity": model.frozen_identity, "initialization_sha256": initialization(model), |
| "run": dict(asdict(pilot_run(recipe))), "checkpoints": points, "order_seed": SEED, "order_seed_b": SEED+1, |
| "population_sha256": digest(plans["REF"][1]), "arm_orders": {arm: plans[arm][2] for arm in ARMS}} |
| validate_inputs(spec, a.manifest, a.parity_receipt, a.g1_receipt) |
| write_json(a.plan, spec) |
|
|
|
|
| class Diagnostics: |
| """Private per-step drafts for matched-arm agreement; no gating thresholds. |
| |
| REF/FAST comparisons at equal steps measure the evolving matched training |
| trajectories, not fixed-weight arithmetic parity. FAST-B has a different |
| order, so its per-step agreement is to the committed training targets. |
| """ |
| def __init__(self, root: Path): |
| self.root = root |
| root.mkdir(parents=True, mode=0o700, exist_ok=True) |
|
|
| @torch.no_grad() |
| def __call__(self, model: Drafter, block: Block, h: Tensor, hx: Tensor, row: dict[str, float]) -> None: |
| inp = {k: v.to(model.activation_device) for k, v in block["inputs"].items()} |
| picks, _ = model.greedy(h, hx, inp["anchor"], chunk=8192) |
| truth = block["targets"]["labels"].cpu() |
| mask = block["targets"]["position_mask"].cpu() |
| picks = picks.cpu() |
| matches = (picks == truth) & mask |
| |
| accepted = matches.long().cumprod(-1).sum(-1).float().mean() |
| payload = {**row, "drafts": picks.tolist(), "target_position1_agreement": float(matches[:, 0].float().mean()), |
| "target_exact_blocks": float((matches | ~mask).all(-1).float().mean()), |
| "mean_accepted_prefix": float(accepted), "informational_only": True} |
| path = self.root/f"step-{int(row['step']):09d}.json" |
| if path.exists(): |
| if json.loads(path.read_text()) != payload: |
| raise ValueError("resumed per-step diagnostics changed") |
| else: |
| write_json(path, payload) |
|
|
|
|
| def arm(a: argparse.Namespace) -> None: |
| from .safety import assert_gpu_lease |
| assert_gpu_lease() |
| spec = read_spec(a.plan, sampled=True) |
| validate_inputs(spec, a.manifest, a.parity_receipt, a.g1_receipt) |
| data = D2Blocks(a.manifest, "train") |
| verify_units(spec, data) |
| plan, population, order = schedule(len(data), spec["checkpoints"], a.arm, spec.get("recipe", "kl")) |
| if digest(population) != spec["population_sha256"] or order != spec["arm_orders"][a.arm]: |
| raise ValueError("pilot data order or population changed") |
| torch.manual_seed(SEED) |
| model = load(a.ckpt, "cuda", rank=32, reference_math=a.arm == "REF") |
| if model.frozen_identity != spec["frozen_identity"] or initialization(model) != spec["initialization_sha256"]: |
| raise ValueError("pilot arm frozen weights or initialization mismatch") |
| if a.arm != "REF": |
| model.enable_fast_training() |
| |
| plan = BatchPlan(plan.batches, plan.checkpoints, digest({"plan": spec, "arm": a.arm, "schedule": plan.identity})) |
| root = a.output/a.arm |
| stop = threading.Event() |
| signal.signal(signal.SIGTERM, lambda *_: stop.set()) |
| signal.signal(signal.SIGINT, lambda *_: stop.set()) |
| result = train(model, data, pilot_run(spec.get("recipe", "kl")), root/"latest.pt", stop, a.resume, plan=plan, observer=Diagnostics(root/"steps")) |
| if not result["complete"]: |
| raise SystemExit("pilot checkpointed before completion; resume this same plan") |
| checkpoints = {} |
| for blocks, step in zip(spec["checkpoints"], plan.checkpoints): |
| p = root/f"latest.checkpoints/step-{step:09d}.pt" |
| state, checkpoint_sha = read_checkpoint(p) |
| if state["blocks"] != blocks or state["batch_plan_identity"] != plan.identity: |
| raise ValueError("pilot archived boundary mismatch") |
| checkpoints[str(blocks)] = {"path": str(p.resolve()), "sha256": checkpoint_sha, "step": step} |
| write_json(root/"arm.json", {"schema": "dsv41.d2.pilot-arm.v1", "arm": a.arm, "recipe": spec.get("recipe", "kl"), "plan_sha256": sha256(a.plan), |
| "batch_plan_identity": plan.identity, "checkpoints": checkpoints, |
| "units": spec["matched_units"][a.arm], "step_count": result["step"], "peak_reserved_bytes": torch.cuda.max_memory_reserved(), |
| "complete": True, "seconds": result["seconds"], "synthetic": spec["synthetic"]}) |
|
|
|
|
| def assert_same_sample(plan: Path, previous_plan: Path) -> None: |
| """Compare the complete pure draw to pilot 1 without changing either plan. |
| |
| The first plan's old source map is deliberately not validated against this |
| new snapshot. Its file hash is admitted by pilot_stage.check_inputs instead. |
| """ |
| previous = cast(Spec, json.loads(previous_plan.read_text())) |
| current = read_spec(plan, sampled=True) |
| validate_sample(previous["sample"], previous["manifest_sha256"]) |
| if current["sample"] != previous["sample"]: |
| raise ValueError("second pilot sample differs from first pilot keys/hash/inputs") |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(__doc__) |
| sub = ap.add_subparsers(dest="command", required=True) |
| sample = sub.add_parser("sample", help="freeze the Amendment-4 held-out sample on CPU, before training") |
| for name in ("plan", "manifest", "k2-dev", "k2-legal"): |
| sample.add_argument("--"+name, required=True, type=Path) |
| sample.add_argument("--previous-plan", type=Path) |
| for cmd in ("prepare", "arm"): |
| p = sub.add_parser(cmd) |
| for n in ("manifest", "parity-receipt", "g1-receipt", "plan"): |
| p.add_argument("--"+n, type=Path, required=True) |
| p.add_argument("--ckpt", required=True) |
| if cmd == "prepare": |
| p.add_argument("--recipe", choices=RECIPES, default="kl") |
| p.add_argument("--rehearsal", action="store_true") |
| else: |
| p.add_argument("--arm", choices=ARMS, required=True) |
| p.add_argument("--output", type=Path, required=True) |
| p.add_argument("--resume", action="store_true") |
| a = ap.parse_args() |
| torch.set_num_threads(8) |
| if a.command == "sample": |
| sample_plan(a.plan, a.manifest, {"dev": a.k2_dev, "heldout_legal": a.k2_legal}) |
| if a.previous_plan is not None: |
| assert_same_sample(a.plan, a.previous_plan) |
| elif a.command == "prepare": |
| prepare(a) |
| else: |
| arm(a) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
