"""Deadline-bounded public-data training, exact artifact reconstruction and eval. Selection uses a recorded validation-only criterion. Calibration and test/whole-family holdout evaluation happen once in the separate finalize command. All checkpoint loads require trusted project artifacts; no weights or optimizer state go in Git. """ import argparse from collections import Counter, defaultdict from datetime import datetime, timezone import hashlib import importlib.metadata import json import math import os from pathlib import Path import platform import random import signal import statistics import subprocess import time import torch from training_model import ADAPTER_VERSION, PROMPT_VERSION, TrainableScorer from selection import SELECTION_METRIC, validation_selection from data_transition import verify_train_data_transition STOP = False class TrainingValidationInterrupted(RuntimeError): """A bounded validation pass yielded to a training stop or deadline.""" def request_stop(*args): global STOP STOP = True def sha256(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest() def write_json(path, value): path = Path(path) tmp = path.with_suffix(".tmp") tmp.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") tmp.replace(path) def save_torch(path, value): path = Path(path) tmp = path.with_suffix(".tmp") torch.save(value, tmp) tmp.replace(path) def mem_available(): return int(next(line.split()[1] for line in Path("/proc/meminfo").read_text().splitlines() if line.startswith("MemAvailable:"))) * 1024 def guard_memory(device="cuda"): Path("/proc/self/oom_score_adj").write_text("0") host_snapshot = subprocess.check_output(['free','-b'],text=True,timeout=10) gpu_snapshot = subprocess.check_output( ['nvidia-smi','--query-compute-apps=pid,process_name,used_memory','--format=csv'], text=True,timeout=10) if device == 'cuda' else None print(json.dumps({'event':'memory_preflight','pid':os.getpid(),'free_bytes':host_snapshot, 'gpu_processes':gpu_snapshot,'mem_available_bytes':mem_available(), 'oom_score_adj':Path('/proc/self/oom_score_adj').read_text().strip()}),flush=True) if mem_available() < 24 * 2**30: raise RuntimeError("Requires at least 24 GiB currently available unified RAM") torch.set_num_threads(8) if device == "cuda": if not torch.cuda.is_available(): raise RuntimeError("CUDA unavailable") torch.cuda.set_per_process_memory_fraction(16 * 2**30 / torch.cuda.get_device_properties(0).total_memory) torch.backends.cuda.matmul.allow_tf32 = False def sync(scorer): if scorer.device.type == "cuda": torch.cuda.synchronize() def load_artifact(path, device="cuda"): artifact = torch.load(path, map_location="cpu", weights_only=False) if artifact.get("format") != "opensysone-adapter-v1": raise ValueError("Unrecognized artifact format") config = artifact["config"] scorer = TrainableScorer(config["model"], rank=config["rank"], alpha=config["alpha"], adapters=config["adapters"], device=device, max_tokens=config["max_tokens"], branch_batch_size=config["branch_batch_size"]) if scorer.provenance != artifact["model_provenance"]: raise ValueError("Artifact/base provenance mismatch") if artifact["prompt_version"] != PROMPT_VERSION or artifact["adapter_version"] != ADAPTER_VERSION: raise ValueError("Artifact implementation version mismatch") scorer.restore_trainable(artifact["trainable_state"]) return scorer, artifact def data_for(scorer, dataset, output): dataset = Path(dataset) source = json.loads((dataset / "manifest.json").read_text()) for split, checksum in source["split_sha256"].items(): if sha256(dataset / f"{split}.jsonl") != checksum: raise ValueError(f"Frozen dataset hash mismatch: {split}") signature = hashlib.sha256(json.dumps({"data": source["split_sha256"], "model": scorer.provenance, "implementation": sha256("training_model.py"), "max_tokens": scorer.max_tokens}, sort_keys=True).encode()).hexdigest() cache = dataset / f"tokens-{signature[:16]}.pt" if cache.exists(): saved = torch.load(cache, map_location="cpu", weights_only=False) if saved["signature"] != signature: raise ValueError("Token cache signature mismatch") data, audit = saved["data"], saved["audit"] else: data, audit = {}, {} for split in source["split_sha256"]: retained, dropped = [], [] for line in (dataset / f"{split}.jsonl").read_text().splitlines(): row = json.loads(line) try: row["_sequences"] = scorer.sequences(row) except ValueError as error: if "no truncation" not in str(error): raise dropped.append(row["id"]) continue retained.append(row) data[split] = retained audit[split] = {"retained": len(retained), "dropped_ids": dropped, "family_counts": dict(Counter(r["family"] for r in retained)), "retained_id_sha256": hashlib.sha256("\n".join(r["id"] for r in retained).encode()).hexdigest(), "max_branch_tokens": max((len(s) for r in retained for s in r["_sequences"]), default=0)} print(json.dumps({"event": "tokenized", "split": split, "audit": {k:v for k,v in audit[split].items() if k != "dropped_ids"}}), flush=True) save_torch(cache, {"signature": signature, "data": data, "audit": audit}) write_json(Path(output) / "data_filter.json", audit) return data, signature def validation_cut(rows, per_family=32): groups = defaultdict(list) for row in rows: groups[row["family"]].append(row) return [row for family in sorted(groups) for row in groups[family][:per_family]] @torch.inference_mode() def predict(scorer, rows, temperature=1.0, token=False, deadline=None): scorer.eval() result = [] start = time.perf_counter() for index,row in enumerate(rows): if deadline is not None and (STOP or time.time() >= deadline): raise TrainingValidationInterrupted("Training validation interrupted before completion") scores = scorer.scores_token_baseline([row]) if token else scorer.score_examples([row]) logits = scores[0].float() / temperature result.append({"id": row["id"], "group": row["group"], "family": row["family"], "target": row["target"], "choices": row["choices"], "logits": scores[0].float().tolist(), "probabilities": logits.softmax(0).tolist(), "log_probabilities": logits.log_softmax(0).tolist()}) if (index+1) % 128 == 0: print(json.dumps({'event':'prediction_progress','decisions':index+1,'total':len(rows), 'elapsed_seconds':time.perf_counter()-start}),flush=True) return result def metrics(predictions): if not predictions: raise ValueError("Cannot report an empty evaluation") bins = [{"count": 0, "confidence_sum": 0.0, "correct_sum": 0.0} for _ in range(10)] values = [] for row in predictions: probabilities = row["probabilities"] predicted = max(range(len(probabilities)), key=probabilities.__getitem__) confidence = probabilities[predicted] correct = float(predicted == row["target"]) nll = -row["log_probabilities"][row["target"]] brier = sum((p - float(i == row["target"])) ** 2 for i,p in enumerate(probabilities)) values.append((correct, nll, brier, confidence)) bucket = bins[min(9, int(confidence * 10))] bucket["count"] += 1 bucket["confidence_sum"] += confidence bucket["correct_sum"] += correct n = len(values) sorted_values = sorted(values, key=lambda v:v[3], reverse=True) coverage = {} for fraction in (0.25, 0.5, 0.75, 1.0): selected = sorted_values[:max(1, math.ceil(n * fraction))] coverage[str(fraction)] = {"n": len(selected), "accuracy": statistics.mean(v[0] for v in selected), "min_confidence": selected[-1][3]} return {"n": n, "accuracy": statistics.mean(v[0] for v in values), "nll": statistics.mean(v[1] for v in values), "brier_multiclass_sum": statistics.mean(v[2] for v in values), "ece_top_label_10_equal_width_bins": sum(abs(b["confidence_sum"] - b["correct_sum"]) for b in bins) / n, "reliability_bins": bins, "accuracy_vs_coverage": coverage} def report(predictions): grouped = defaultdict(list) for row in predictions: grouped[row["family"]].append(row) return {"overall": metrics(predictions), "per_family": {f:metrics(rows) for f,rows in grouped.items()}} def objective(predictions): return statistics.mean(value["nll"] for value in report(predictions)["per_family"].values()) def decision_backward(scorer, row, divisor, two_pass=False): """Exact categorical gradient with one candidate graph alive in two-pass mode. Requires deterministic, dropout-free forwards. The first pass supplies the softmax derivative; recomputation applies that derivative to each scalar score. """ if not two_pass: score = scorer.score_examples([row])[0].float() loss = -score.log_softmax(0)[row["target"]] if not torch.isfinite(loss): raise RuntimeError("Non-finite training loss") (loss / divisor).backward() return loss.item() with torch.no_grad(): score = scorer.score_examples([row])[0].float() loss = -score.log_softmax(0)[row["target"]] derivative = score.softmax(0) derivative[row["target"]] -= 1 if not torch.isfinite(loss): raise RuntimeError("Non-finite training loss") for index, sequence in enumerate(row["_sequences"]): value = scorer.score_examples([{"_sequences": [sequence]}])[0][0] value.backward(gradient=derivative[index] / divisor) return loss.item() @torch.inference_mode() def correctness(scorer, rows, token_initial=False, allow_stop=False): scorer.eval() chosen = validation_cut(rows, 2) original_size = scorer.branch_batch_size def scores(examples, token=False): if allow_stop and STOP: raise TrainingValidationInterrupted("Final training correctness stopped before completion") return scorer.scores_token_baseline(examples) if token else scorer.score_examples(examples) base = [s.detach().clone() for s in scores(chosen)] def difference(a,b): return max((x.float().softmax(0) - y.float().softmax(0)).abs().max().item() for x,y in zip(a,b)) checks = {} try: for size in (1, 2, 4): scorer.branch_batch_size = size checks[f"branch_chunks_{size}_probability_max_abs"] = difference(base, scores(chosen)) scorer.branch_batch_size = original_size single = [scores([row])[0] for row in chosen] checks["question_isolation_probability_max_abs"] = difference(base, single) reversed_rows = [] for row in chosen: reversed_rows.append({**row, "choices": list(reversed(row["choices"])), "_sequences": list(reversed(row["_sequences"]))}) flipped = [s.flip(0) for s in scores(reversed_rows)] checks["candidate_permutation_probability_max_abs"] = difference(base, flipped) checks["repeat_probability_max_abs"] = difference(base, scores(chosen)) if token_initial: checks["pretrained_readout_probability_max_abs"] = difference(base, scores(chosen, token=True)) finally: scorer.branch_batch_size = original_size checks["tolerance_probability_abs"] = 1e-4 if any(value > 1e-4 for key,value in checks.items() if key.endswith("max_abs")): raise RuntimeError(f"FP32 correctness gate failed: {checks}") return checks def restore_warm_start(scorer, saved, config, signature): """Restore compatible trained weights while retaining a fresh optimizer/RNG. Optimization settings may change. An explicitly requested data expansion verifies both datasets and preserves every reserved evaluation split. """ if saved.get("format") != "opensysone-adapter-v1": raise ValueError("Unrecognized warm-start artifact format") if (saved.get("prompt_version") != PROMPT_VERSION or saved.get("adapter_version") != ADAPTER_VERSION): raise ValueError("Warm-start implementation version mismatch") for key in ("rank", "alpha", "adapters"): if saved["config"].get(key) != config[key]: raise ValueError(f"Warm-start must preserve {key}") if saved["model_provenance"] != scorer.provenance: raise ValueError("Warm-start model/data mismatch") transition = None if config.get("allow_train_data_change", False): transition = verify_train_data_transition(saved, config, signature) elif saved["data_signature"] != signature: raise ValueError("Warm-start model/data mismatch") scorer.restore_trainable(saved["trainable_state"]) return transition def validation_fits(deadline, measured_seconds, now=None): """Do not start a full validation pass that would consume the stop margin.""" now = time.time() if now is None else now return now + max(60, measured_seconds * 1.25 + 30) < deadline def selection_result(predictions, metric): if metric == "raw_nll": return {"metric": metric, "score": objective(predictions), "raw_macro_nll": objective(predictions)} if metric != SELECTION_METRIC: raise ValueError("Unknown checkpoint selection metric") return validation_selection(predictions) def reselect_inherited(artifact, predictions, metric): """Change selection metadata only; retained weights/provenance stay intact.""" result = selection_result(predictions, metric) updated = dict(artifact) updated.update(selection_metric=metric, best_validation_macro_nll=result["raw_macro_nll"], best_validation_selection_score=result["score"], validation_selection=result) return updated, result def train(args): if getattr(args, "allow_train_data_change", False) and not args.warm_start: raise ValueError("--allow-train-data-change requires --warm-start") out = Path(args.output).resolve() out.mkdir(parents=True, exist_ok=False) guard_memory() signal.signal(signal.SIGTERM, request_stop) signal.signal(signal.SIGINT, request_stop) random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) config = vars(args).copy() config["model"] = str(Path(args.model).resolve()) config["dataset"] = str(Path(args.dataset).resolve()) config["adapters"] = not args.head_only config["selection_metric"] = getattr(args, "selection_metric", "raw_nll") parent_path = args.resume or args.warm_start initialization = {"kind": "resume" if args.resume else "warm_start" if args.warm_start else "pretrained", "parent_checkpoint": str(Path(parent_path).resolve()) if parent_path else None, "parent_checkpoint_sha256": sha256(parent_path) if parent_path else None, "restores_optimizer": bool(args.resume), "restores_rng": bool(args.resume)} manifest = {"config": config, "pid": os.getpid(), "hostname": platform.node(), "started_utc": datetime.now(timezone.utc).isoformat(), "git_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), "git_status": subprocess.check_output(["git", "status", "--porcelain"], text=True), "source_sha256": {str(p):sha256(p) for p in [*Path('.').glob('*.py'), *Path('scripts').glob('*')] if p.is_file()}, "packages": {p:importlib.metadata.version(p) for p in ("torch", "transformers", "pyarrow", "numpy")}, "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(), "capability": torch.cuda.get_device_capability(), "cuda_cap_bytes": 16 * 2**30, "initial_mem_available_bytes": mem_available(), "oom_score_adj": Path('/proc/self/oom_score_adj').read_text().strip(), "parent_checkpoint_sha256": initialization["parent_checkpoint_sha256"], "initialization": initialization} write_json(out / "manifest.json", manifest) start = time.perf_counter() scorer = TrainableScorer(config["model"], rank=args.rank, alpha=args.alpha, adapters=config["adapters"], max_tokens=args.max_tokens, branch_batch_size=args.branch_batch_size) if args.two_pass: if args.branch_batch_size != 1: raise ValueError("Two-pass training requires branch_batch_size=1 to preserve recomputation shapes") if any(isinstance(m, torch.nn.Dropout) and m.p for m in scorer.modules()) or getattr(scorer.lm.config, "attention_dropout", 0): raise ValueError("Two-pass gradients require dropout-free forwards") data, signature = data_for(scorer, args.dataset, out) manifest.update(model_provenance=scorer.provenance, data_signature=signature, total_parameters=sum(p.numel() for p in scorer.parameters()), trainable_parameters=sum(p.numel() for p in scorer.parameters() if p.requires_grad), adapter_modules=scorer.adapter_names, load_and_tokenize_seconds=time.perf_counter() - start) write_json(out / "manifest.json", manifest) optimizer = torch.optim.AdamW([ {"params": [p for p in scorer.lm.parameters() if p.requires_grad], "lr": args.lr}, {"params": scorer.head.parameters(), "lr": args.head_lr}], weight_decay=0.01) parameters = [p for p in scorer.parameters() if p.requires_grad] maximum_steps = math.ceil(len(data["train"]) * args.epochs / args.effective_batch) target_steps = min(args.steps, maximum_steps) if args.steps else maximum_steps completed, best, best_selection, stale = 0, math.inf, math.inf, 0 parent_best = None if args.warm_start: saved = torch.load(args.warm_start, map_location="cpu", weights_only=False) transition = restore_warm_start(scorer, saved, config, signature) initialization.update(parent_step=saved["step"], parent_source_commit=saved["source_commit"]) if transition is not None: initialization["data_transition"] = transition del saved if args.resume: saved = torch.load(args.resume, map_location="cpu", weights_only=False) initialization.update(parent_step=saved["step"], parent_source_commit=saved["source_commit"], parent_initialization=saved.get("initialization")) for key in ("rank", "alpha", "adapters", "max_tokens", "branch_batch_size", "seed", "effective_batch", "lr", "head_lr", "schedule_steps", "epochs", "validation_per_family", "two_pass"): if saved["config"].get(key,False) != config[key]: raise ValueError(f"Resume must preserve {key}") if saved["model_provenance"] != scorer.provenance or saved["data_signature"] != signature: raise ValueError("Resume model/data mismatch") scorer.restore_trainable(saved["trainable_state"]) optimizer.load_state_dict(saved["optimizer"]) random.setstate(saved["random_state"]) torch.set_rng_state(saved["torch_rng"]) torch.cuda.set_rng_state_all(saved["cuda_rng"]) completed, best, stale = saved["step"], saved["best_validation_macro_nll"], saved["stale_evaluations"] previous_metric = saved.get("selection_metric", "raw_nll") initialization["selection_policy_change"] = { "from": previous_metric, "to": config["selection_metric"], "optimizer_and_rng_unchanged": True} parent_best = Path(args.resume).resolve().parent / "best.pt" if not parent_best.exists(): raise ValueError("Resume requires the parent's validation-selected best.pt") inherited = torch.load(parent_best, map_location="cpu", weights_only=False) if inherited["data_signature"] != signature: raise ValueError("Parent best artifact data mismatch") # Legacy campaigns copied an inherited best artifact without its raw # prediction file. Its recorded output locates that original evidence. evidence_directories = [parent_best.parent] if inherited["config"].get("output"): evidence_directories.append(Path(inherited["config"]["output"])) candidates = [] for directory in evidence_directories: candidates.extend([directory / f"validation_step_{inherited['step']:06d}_predictions.json", directory / "best_validation_predictions.json"]) if inherited["step"] == 0: candidates.append(directory / "initial_validation_predictions.json") best_predictions = None for candidate in candidates: if candidate.exists(): candidate_predictions = json.loads(candidate.read_text()) # A canonical file can belong to a different step after interruption. if abs(objective(candidate_predictions) - inherited["best_validation_macro_nll"]) > 2e-5: continue best_predictions = candidate_predictions write_json(out / "best_validation_predictions.json", best_predictions) write_json(out / f"validation_step_{inherited['step']:06d}_predictions.json", best_predictions) break if best_predictions is None: raise ValueError("Inherited best requires matching raw validation evidence") inherited, inherited_selection = reselect_inherited(inherited, best_predictions, config["selection_metric"]) best, best_selection = inherited_selection["raw_macro_nll"], inherited_selection["score"] if previous_metric != config["selection_metric"]: stale = 0 write_json(out / "inherited_validation_selection.json", inherited_selection) write_json(out / "best_validation_selection.json", inherited_selection) save_torch(out / "best.pt", inherited) del saved, inherited write_json(out / "manifest.json", manifest) def checkpoint(path, step, resumable=True): artifact = {"format": "opensysone-adapter-v1", "step": step, "config": config, "model_provenance": scorer.provenance, "data_signature": signature, "prompt_version": PROMPT_VERSION, "adapter_version": ADAPTER_VERSION, "trainable_state": scorer.trainable_state(), "best_validation_macro_nll": best, "selection_metric": config["selection_metric"], "best_validation_selection_score": best_selection, "stale_evaluations": stale, "source_commit": manifest["git_commit"], "initialization": initialization} if resumable: artifact.update(optimizer=optimizer.state_dict(), random_state=random.getstate(), torch_rng=torch.get_rng_state(), cuda_rng=torch.cuda.get_rng_state_all()) save_torch(path, artifact) # A durable reconstruction exists before any correctness/evaluation work. checkpoint(out / "checkpoint.pt", completed) checks = correctness(scorer, data["validation"], token_initial=not parent_path) write_json(out / "correctness_initial.json", checks) val = validation_cut(data["validation"], args.validation_per_family) validation_started = time.monotonic() predictions = predict(scorer, val) validation_seconds = time.monotonic() - validation_started write_json(out / "resumed_initial_predictions.json" if args.resume else out / "initial_validation_predictions.json", predictions) initial_report = report(predictions) initial_objective = objective(predictions) initial_selection = selection_result(predictions, config["selection_metric"]) write_json(out / "initial_validation_selection.json", initial_selection) if args.resume: write_json(out / f"validation_step_{completed:06d}_predictions.json", predictions) initial_improved = initial_selection["score"] < best_selection - 0.001 if not args.resume or initial_improved: best, best_selection, stale = initial_objective, initial_selection["score"], 0 write_json(out / "best_validation_predictions.json", predictions) write_json(out / "best_validation_selection.json", initial_selection) checkpoint(out / "best.pt", completed, False) if args.resume: print(json.dumps({"event": "resumed_initial_selection", "step": completed, "macro_nll": initial_objective, "best_macro_nll": best, "selection_metric": config["selection_metric"], "selection_score": initial_selection["score"], "best_selection_score": best_selection, "improved": initial_improved}), flush=True) # Persist an initial/resumed selection before a stop can interrupt new updates. checkpoint(out / "checkpoint.pt", completed) print(json.dumps({"event": "ready", "output": str(out), "trainable_parameters": manifest["trainable_parameters"], "validation_macro_nll": initial_objective, "target_steps": target_steps, "correctness": checks}), flush=True) deadline = datetime.fromisoformat(args.deadline.replace("Z", "+00:00")).timestamp() if args.deadline else math.inf last_save = time.monotonic() history = [] orders = {} train_start = time.monotonic() status = "completed_step_target" if args.steps and args.steps < maximum_steps else "completed_epochs" for step in range(completed + 1, target_steps + 1): if STOP or time.time() >= deadline: status = "interrupted" if STOP else "training_deadline" break if mem_available() < 16 * 2**30: checkpoint(out / "checkpoint.pt", completed) raise RuntimeError("Host availability below 16 GiB; checkpoint saved") scorer.train() optimizer.zero_grad(set_to_none=True) batch = [] for absolute in range((step - 1) * args.effective_batch, min(step * args.effective_batch, len(data["train"]) * args.epochs)): epoch, index = divmod(absolute, len(data["train"])) if epoch not in orders: order = list(range(len(data["train"]))) random.Random(args.seed + epoch).shuffle(order) orders = {epoch:order} batch.append(data["train"][orders[epoch][index]]) sync(scorer) tick = time.perf_counter() loss_sum = actual_tokens = padded_tokens = branches = 0 for row in batch: loss_sum += decision_backward(scorer,row,len(batch),args.two_pass) sequences = row["_sequences"] actual_tokens += sum(map(len, sequences)) branches += len(sequences) for start_index in range(0, len(sequences), args.branch_batch_size): chunk = sequences[start_index:start_index + args.branch_batch_size] padded_tokens += max(map(len, chunk)) * len(chunk) norm = torch.nn.utils.clip_grad_norm_(parameters, 1.0, error_if_nonfinite=True) # Fixed schedule length is persisted separately for resumable pilots/campaigns. schedule_steps = args.schedule_steps or maximum_steps warmup = max(1, min(100, schedule_steps // 20)) factor = min(1.0, step / warmup) if step <= warmup else max(0.1, 0.5 * (1 + math.cos(math.pi * min(1, (step - warmup) / max(1, schedule_steps - warmup))))) for group, lr in zip(optimizer.param_groups, (args.lr, args.head_lr)): group["lr"] = lr * factor optimizer.step() sync(scorer) completed = step item = {"step": step, "loss": loss_sum / len(batch), "gradient_norm": norm.item(), "seconds": time.perf_counter() - tick, "decisions": len(batch), "branches": branches, "actual_branch_tokens": actual_tokens, "padded_branch_tokens": padded_tokens, "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), "lr_factor": factor} history.append(item) with (out / "training.jsonl").open("a") as handle: handle.write(json.dumps(item) + "\n") if step % 10 == 0 or step == target_steps: print(json.dumps({"event": "step", **item}), flush=True) if step % args.save_steps == 0 or time.monotonic() - last_save >= args.save_seconds: checkpoint(out / "checkpoint.pt", completed) last_save = time.monotonic() if step % args.eval_steps == 0 or step == target_steps: checkpoint(out / "checkpoint.pt", completed) if not validation_fits(deadline, validation_seconds): print(json.dumps({"event": "validation_skipped_for_deadline", "step": completed, "measured_validation_seconds": validation_seconds}), flush=True) continue validation_started = time.monotonic() try: predictions = predict(scorer, val, deadline=deadline) except TrainingValidationInterrupted: status = "interrupted" if STOP else "training_deadline" print(json.dumps({"event": "validation_interrupted", "step": completed, "status": status}), flush=True) break validation_seconds = max(validation_seconds, time.monotonic() - validation_started) score = objective(predictions) selected_score = selection_result(predictions, config["selection_metric"]) improved = selected_score["score"] < best_selection - 0.001 # Save evidence before publishing a newly selected checkpoint. Fleet # selection can always retrieve predictions for the durable best step. write_json(out / f"validation_step_{completed:06d}_predictions.json", predictions) if improved: best, best_selection, stale = score, selected_score["score"], 0 write_json(out / "best_validation_predictions.json", predictions) write_json(out / "best_validation_selection.json", selected_score) checkpoint(out / "best.pt", completed, False) else: stale += 1 validation = {"step": completed, "macro_nll": score, "best_macro_nll": best, "selection": selected_score, "best_selection_score": best_selection, "improved": improved, "metrics": report(predictions)} with (out / "validation.jsonl").open("a") as handle: handle.write(json.dumps(validation) + "\n") print(json.dumps({"event": "validation", "step": completed, "macro_nll": score, "best": best}), flush=True) checkpoint(out / "checkpoint.pt", completed) last_save = time.monotonic() if stale >= args.patience: status = "validation_early_stop" break checkpoint(out / "checkpoint.pt", completed) try: checks = correctness(scorer, data["validation"], allow_stop=True) except TrainingValidationInterrupted: checks = None final_correctness_status = "skipped_on_stop" status = "interrupted" else: final_correctness_status = "passed" write_json(out / "correctness_final.json", checks) summary = {"status": status, "completed_steps": completed, "target_steps": target_steps, "initial_validation": initial_report, "best_validation_macro_nll": best, "selection_metric": config["selection_metric"], "best_validation_selection_score": best_selection, "training_seconds": time.monotonic() - train_start, "median_step_seconds": statistics.median(item["seconds"] for item in history) if history else None, "processed_decisions": sum(item["decisions"] for item in history), "actual_branch_tokens": sum(item["actual_branch_tokens"] for item in history), "padded_branch_tokens": sum(item["padded_branch_tokens"] for item in history), "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_cuda_reserved_bytes": torch.cuda.max_memory_reserved(), "final_mem_available_bytes": mem_available(), "correctness": checks, "final_correctness_status": final_correctness_status, "checkpoint_sha256": sha256(out / "checkpoint.pt"), "best_sha256": sha256(out / "best.pt")} write_json(out / "summary.json", summary) print(json.dumps({"event": "training_complete", **summary}), flush=True) def fit_temperature(predictions): candidates = torch.logspace(-1, 1.3, 101).tolist() return min(candidates, key=lambda t: statistics.mean( -torch.tensor(row["logits"]).div(t).log_softmax(0)[row["target"]].item() for row in predictions)) def with_temperature(predictions, temperature): return [{**row, "probabilities": (torch.tensor(row["logits"]) / temperature).softmax(0).tolist(), "log_probabilities": (torch.tensor(row["logits"]) / temperature).log_softmax(0).tolist()} for row in predictions] def bootstrap_difference(base, tuned, repetitions=400): import numpy as np if [row["id"] for row in base] != [row["id"] for row in tuned]: raise ValueError("Bootstrap needs matched predictions") groups = defaultdict(list) def losses(row): p = row["probabilities"] return [float(max(range(len(p)), key=p.__getitem__) == row["target"]), -row["log_probabilities"][row["target"]], sum((v - float(i == row["target"])) ** 2 for i,v in enumerate(p))] for a,b in zip(base,tuned): groups[(a["family"],a["group"])].append(np.asarray(losses(b)) - np.asarray(losses(a))) by_family = defaultdict(list) for (family,_),rows in groups.items(): by_family[family].append((np.sum(rows,axis=0),len(rows))) arrays = [(np.asarray([r[0] for r in rows]),np.asarray([r[1] for r in rows])) for rows in by_family.values()] rng = np.random.default_rng(907) samples = [] for _ in range(repetitions): total, count = np.zeros(3), 0 for sums,counts in arrays: selected = rng.integers(0,len(sums),len(sums)) total += sums[selected].sum(axis=0) count += counts[selected].sum() samples.append(total/count) lo, hi = np.quantile(samples,[0.025,0.975],axis=0) point = sum((sums.sum(axis=0) for sums,_ in arrays),np.zeros(3))/len(base) return {"method": f"{repetitions} stratified source-group bootstrap resamples; decision-weighted tuned minus base", "point_delta":dict(zip(('accuracy','nll','brier'),map(float,point))), "accuracy": [float(lo[0]),float(hi[0])], "nll": [float(lo[1]),float(hi[1])], "brier": [float(lo[2]),float(hi[2])]} def finalize(args): out = Path(args.output).resolve() out.mkdir(parents=True, exist_ok=False) guard_memory() scorer, selected = load_artifact(args.checkpoint) data, signature = data_for(scorer, args.dataset, out) if signature != selected["data_signature"]: raise ValueError("Evaluation dataset differs from training") write_json(out / "manifest.json", {"pid": os.getpid(), "source_commit": subprocess.check_output( ['git','rev-parse','HEAD'],text=True).strip(), "checkpoint_sha256": sha256(args.checkpoint), "selected_step": selected["step"], "data_signature": signature, "model_provenance": scorer.provenance, "training_config": selected["config"], "training_source_commit": selected["source_commit"], "source_sha256": {str(p):sha256(p) for p in [*Path('.').glob('*.py'), *Path('scripts').glob('*')] if p.is_file()}, "packages": {p:importlib.metadata.version(p) for p in ("torch", "transformers", "pyarrow", "numpy")}, "hostname": platform.node(), "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(), "cuda_cap_bytes": 16 * 2**30, "initial_mem_available_bytes": mem_available(), "oom_score_adj": Path('/proc/self/oom_score_adj').read_text().strip(), "selection": {"metric": selected.get("selection_metric", "raw_nll"), "score": selected.get("best_validation_selection_score", selected["best_validation_macro_nll"]), "raw_macro_nll": selected["best_validation_macro_nll"], "scope": "validation only; reserved calibration/test/holdout not used"}, "started_utc": datetime.now(timezone.utc).isoformat()}) write_json(out / "correctness.json", correctness(scorer,data["validation"])) calibration = predict(scorer,data["calibration"]) write_json(out / "calibration_predictions.json",calibration) temperature = fit_temperature(calibration) selected["temperature"] = temperature selected["calibration_status"] = "one global temperature fitted on separate known-family calibration data; whole-family calibration unproven" selected.pop("optimizer",None) # The deployable artifact is durable before untouched evaluation starts. save_torch(out / "model.pt",selected) result = {"temperature": temperature, "selected_step":selected["step"], "status":"pending_test", "claim_scope":"Public decision benchmark; no claim of Jev-level intelligence or general calibration"} all_tuned = {} for split in ("test","holdout"): predictions = predict(scorer,data[split]) calibrated = with_temperature(predictions,temperature) all_tuned[split] = calibrated write_json(out / f"{split}_trained_predictions.json",predictions) write_json(out / f"{split}_calibrated_predictions.json",calibrated) result[split] = {"trained":report(predictions), "calibrated":report(calibrated)} write_json(out / "metrics.json",result) print(json.dumps({"event":"evaluation","split":split,"metrics":result[split]}),flush=True) # Reconstruct the unchanged pretrained readout in the same model allocation. # Reset adapters without creating a second backbone on the GPU. with torch.no_grad(): for name,p in scorer.named_parameters(): if name.endswith("adapter_b"): p.zero_() scorer.head.weight.copy_((scorer.lm.lm_head.weight[scorer.yes_no_ids[0]] - scorer.lm.lm_head.weight[scorer.yes_no_ids[1]]).unsqueeze(0)) bias = scorer.lm.lm_head.bias scorer.head.bias.fill_(0 if bias is None else bias[scorer.yes_no_ids[0]]-bias[scorer.yes_no_ids[1]]) base_calibration = predict(scorer,data["calibration"]) base_temperature = fit_temperature(base_calibration) write_json(out / "base_calibration_predictions.json",base_calibration) result["base_temperature"] = base_temperature for split in ("test","holdout"): predictions = predict(scorer,data[split]) calibrated = with_temperature(predictions,base_temperature) write_json(out / f"{split}_base_predictions.json",predictions) write_json(out / f"{split}_base_calibrated_predictions.json",calibrated) result[split].update(base=report(predictions),base_calibrated=report(calibrated), calibrated_difference_95pct=bootstrap_difference(calibrated,all_tuned[split])) result.update(status="complete",model_sha256=sha256(out / "model.pt"), peak_cuda_allocated_bytes=torch.cuda.max_memory_allocated(), peak_cuda_reserved_bytes=torch.cuda.max_memory_reserved(),final_mem_available_bytes=mem_available()) write_json(out / "metrics.json",result) print(json.dumps({"event":"finalized","output":str(out),"model_sha256":result["model_sha256"]}),flush=True) def main(): parser = argparse.ArgumentParser(description=__doc__) sub = parser.add_subparsers(dest="command",required=True) training = sub.add_parser("train") training.add_argument("--model",required=True) training.add_argument("--dataset",required=True) training.add_argument("--output",required=True) initialization = training.add_mutually_exclusive_group() initialization.add_argument("--resume", help="Restore weights, optimizer and RNG without changing training configuration") initialization.add_argument("--warm-start", help="Initialize compatible trained weights with a fresh optimizer and training configuration") training.add_argument("--allow-train-data-change", action="store_true", help="Explicit warm-start expansion; verify parent lineage and preserve all reserved data") training.add_argument("--steps",type=int) training.add_argument("--epochs",type=int,default=3) training.add_argument("--rank",type=int,default=16) training.add_argument("--alpha",type=float,default=32) training.add_argument("--head-only",action="store_true") training.add_argument("--two-pass",action="store_true") training.add_argument("--lr",type=float,default=1e-4) training.add_argument("--head-lr",type=float,default=1e-4) training.add_argument("--seed",type=int,default=431) training.add_argument("--effective-batch",type=int,default=4) training.add_argument("--branch-batch-size",type=int,default=2) training.add_argument("--max-tokens",type=int,default=768) training.add_argument("--save-steps",type=int,default=250) training.add_argument("--save-seconds",type=int,default=900) training.add_argument("--eval-steps",type=int,default=500) training.add_argument("--validation-per-family",type=int,default=32) training.add_argument("--patience",type=int,default=8) training.add_argument("--deadline") training.add_argument("--schedule-steps",type=int) training.add_argument("--selection-metric", choices=("raw_nll", SELECTION_METRIC), default="raw_nll", help="Validation-only checkpoint selection; changing it preserves optimizer/RNG but re-scores inherited evidence") evaluation = sub.add_parser("finalize") evaluation.add_argument("--checkpoint",required=True) evaluation.add_argument("--dataset",required=True) evaluation.add_argument("--output",required=True) args = parser.parse_args() if args.command == "train": if any(getattr(args,key) <= 0 for key in ('epochs','rank','alpha','lr','head_lr','effective_batch','branch_batch_size','max_tokens','save_steps','save_seconds','eval_steps','validation_per_family','patience')): parser.error("training sizes, rates and cadences must be positive") if args.steps is not None and args.steps <= 0: parser.error("--steps must be positive") if args.schedule_steps is not None and args.schedule_steps <= 0: parser.error("--schedule-steps must be positive") train(args) else: finalize(args) if __name__ == "__main__": main()