"""Bounded decision-model training, correctness checks, and illustrative timing. Run via scripts/run_smoke.sh on GX10. No customer datasets, generation services, or model-server processes are used. Checkpoints include optimizer and RNG state. """ import argparse import hashlib import importlib.metadata import json import os from pathlib import Path import platform import random import signal import statistics import subprocess import time os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") REVISION = "060db6499f32faf8b98477b0a26969ef7d8b9987" STOP = False def request_stop(*_): global STOP STOP = True def write_json(path, value): path = Path(path) tmp = path.with_suffix(path.suffix + ".tmp") tmp.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") tmp.replace(path) def command(args): proc = subprocess.run(args, capture_output=True, text=True) return proc.stdout.strip() if proc.returncode == 0 else None def mem_available(): return int(next(line.split()[1] for line in Path("/proc/meminfo").read_text().splitlines() if line.startswith("MemAvailable:"))) * 1024 def main(): parser = argparse.ArgumentParser() parser.add_argument("--model", default=str(Path.home() / "ai/models/opensysone" / f"Qwen2.5-0.5B-{REVISION[:8]}")) parser.add_argument("--output", required=True) parser.add_argument("--steps", type=int, default=60) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--train-layers", type=int, default=2) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--resume", help="Trusted checkpoint.pt from this script; steps is total target") parser.add_argument("--cuda-cap-gib", type=float, default=16) parser.add_argument("--dtype", choices=["float32", "bfloat16"], default="float32") args = parser.parse_args() out = Path(args.output) out.mkdir(parents=True, exist_ok=True) if (out / "manifest.json").exists(): raise ValueError("Use a fresh output directory; previous run artifacts are immutable") provenance = json.loads((Path(args.model) / "opensysone-provenance.json").read_text()) if provenance != {"model_id": "Qwen/Qwen2.5-0.5B", "revision": REVISION}: raise ValueError("Model directory must contain the pinned download's provenance") if args.steps < 1 or args.batch_size < 1 or args.train_layers < 1: raise ValueError("steps, batch size, train layers must be positive") # SSH children inherit gx10's absolute OOM exemption. Give it back before torch. Path("/proc/self/oom_score_adj").write_text("0") if mem_available() < 24 * 2**30: raise RuntimeError("Smoke requires at least 24 GiB currently available unified RAM") signal.signal(signal.SIGTERM, request_stop) signal.signal(signal.SIGINT, request_stop) import torch from decision_model import DecisionScorer from smoke_data import make_data, VERSION torch.set_num_threads(8) random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) if not torch.cuda.is_available(): raise RuntimeError("This run requires a CUDA GPU") total_memory = torch.cuda.get_device_properties(0).total_memory if not 0 < args.cuda_cap_gib <= 16: raise ValueError("This smoke permits a CUDA allocation cap of at most 16 GiB") torch.cuda.set_per_process_memory_fraction(args.cuda_cap_gib * 2**30 / total_memory) data = make_data() for split, rows in data.items(): (out / f"{split}.jsonl").write_text("".join(json.dumps(r) + "\n" for r in rows)) files = [*Path(".").glob("*.py"), *Path("scripts").glob("*")] manifest = { "config": vars(args), "model_id": "Qwen/Qwen2.5-0.5B", "model_revision": REVISION, "pid": os.getpid(), "dataset_version": VERSION, "hostname": platform.node(), "platform": platform.platform(), "python": platform.python_version(), "started_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "git_commit": command(["git", "rev-parse", "HEAD"]), "git_status": command(["git", "status", "--porcelain"]), "parent_checkpoint": str(Path(args.resume).resolve()) if args.resume else None, "source_sha256": {str(p): hashlib.sha256(p.read_bytes()).hexdigest() for p in files if p.is_file()}, "dataset_sha256": {s: hashlib.sha256((out / f"{s}.jsonl").read_bytes()).hexdigest() for s in data}, "packages": {n: importlib.metadata.version(n) for n in ["torch", "transformers", "accelerate", "huggingface-hub", "safetensors", "numpy"]}, "cuda": torch.version.cuda, "cudnn": torch.backends.cudnn.version(), "nccl": torch.cuda.nccl.version(), "gpu": torch.cuda.get_device_name(), "capability": torch.cuda.get_device_capability(), "unified_total_bytes": total_memory, "initial_mem_available_bytes": mem_available(), "oom_score_adj": Path("/proc/self/oom_score_adj").read_text().strip(), "gpu_snapshot": command(["nvidia-smi", "--query-gpu=name,driver_version,temperature.gpu,utilization.gpu,power.draw", "--format=csv"]), } write_json(out / "manifest.json", manifest) start = time.perf_counter() scorer = DecisionScorer(args.model, train_layers=args.train_layers, dtype=getattr(torch, args.dtype)) trainable = [p for p in scorer.parameters() if p.requires_grad] optimizer = torch.optim.AdamW([ {"params": [p for p in scorer.lm.parameters() if p.requires_grad], "lr": 2e-5}, {"params": list(scorer.head.parameters()), "lr": 1e-3}, ], weight_decay=0.01) manifest["total_parameters"] = sum(p.numel() for p in scorer.parameters()) manifest["trainable_parameters"] = sum(p.numel() for p in trainable) manifest["trainable_names"] = [n for n, p in scorer.named_parameters() if p.requires_grad] manifest["precision"] = f"{args.dtype} backbone and backbone AdamW states; FP32 scalar head and its AdamW states" probe_name = next(n for n, p in scorer.named_parameters() if p.requires_grad and n.startswith("lm.")) initial_probe = dict(scorer.named_parameters())[probe_name].detach().float().cpu().clone() manifest["load_seconds"] = time.perf_counter() - start write_json(out / "manifest.json", manifest) def evaluate(rows, method="scalar", temperature=1.0): scorer.eval() values = [] with torch.inference_mode(): for j in range(0, len(rows), args.batch_size): group = rows[j:j + args.batch_size] scores = (scorer.scores_token_baseline(group) if method == "token" else scorer.score_examples(group)) for row, score in zip(group, scores): probabilities = torch.softmax(score.float() / temperature, dim=-1) values.append({"id": row["id"], "family": row["family"], "target": row["target"], "choices": row["choices"], "logits": score.float().tolist(), "probabilities": probabilities.tolist(), "log_probabilities": torch.log_softmax(score.float() / temperature, -1).tolist()}) return values def metrics(predictions): n = len(predictions) accuracy = nll = brier = 0.0 bins = [{"count": 0, "confidence_sum": 0.0, "correct_sum": 0.0} for _ in range(10)] for row in predictions: p = torch.tensor(row["probabilities"]) target = row["target"] confidence, prediction = p.max(dim=0) correct = float(prediction.item() == target) accuracy += correct nll -= row["log_probabilities"][target] brier += float(((p - torch.nn.functional.one_hot(torch.tensor(target), len(p)))**2).sum()) bucket = bins[min(9, int(confidence.item() * 10))] bucket["count"] += 1 bucket["confidence_sum"] += confidence.item() bucket["correct_sum"] += correct ece = sum(abs(b["confidence_sum"] - b["correct_sum"]) for b in bins) / n return {"n": n, "accuracy": accuracy / n, "nll": nll / n, "brier_multiclass_sum": brier / n, "ece_top_label_10_equal_width_bins": ece, "reliability_bins": bins} def checkpoint(step): path = out / "checkpoint.pt" tmp = out / "checkpoint.tmp" state = {n: p.detach().cpu() for n, p in scorer.named_parameters() if p.requires_grad} torch.save({"step": step, "trainable_state": state, "optimizer": optimizer.state_dict(), "random_state": random.getstate(), "torch_rng": torch.get_rng_state(), "cuda_rng": torch.cuda.get_rng_state_all(), "config": vars(args), "model_revision": REVISION}, tmp) tmp.replace(path) completed = 0 if args.resume: saved = torch.load(args.resume, map_location="cpu", weights_only=False) if saved["model_revision"] != REVISION or saved["config"]["train_layers"] != args.train_layers: raise ValueError("Resume base revision or train_layers mismatch") if saved["config"].get("dtype", "bfloat16") != args.dtype: raise ValueError("Resume must preserve checkpoint backbone precision") if saved["config"]["batch_size"] != args.batch_size or saved["config"]["seed"] != args.seed: raise ValueError("Resume requires original batch_size and seed") if set(saved["trainable_state"]) != set(manifest["trainable_names"]): raise ValueError("Resume checkpoint does not cover exactly the trainable parameters") scorer.load_state_dict(saved["trainable_state"], strict=False) 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 = saved["step"] results = {"claim_scope": "Synthetic same-template engineering smoke only; no zero-shot or calibration claim."} print(json.dumps({"event": "loaded", "trainable_parameters": manifest["trainable_parameters"], "load_seconds": manifest["load_seconds"]}), flush=True) if not args.resume: for label, method in [("base_token_yes_minus_no", "token"), ("initial_scalar", "scalar")]: pred = evaluate(data["test"], method) write_json(out / f"{label}_predictions.json", pred) results[label] = metrics(pred) print(json.dumps({"event": label, "metrics": results[label]}), flush=True) else: pred = evaluate(data["test"]) write_json(out / "resumed_initial_predictions.json", pred) results["resumed_initial"] = metrics(pred) checkpoint(completed) history = [] train_start = time.perf_counter() scorer.train() for step in range(completed + 1, args.steps + 1): if STOP: break if mem_available() < 12 * 2**30: raise RuntimeError("Available host memory fell below 12 GiB; last periodic checkpoint is retained") rows = random.sample(data["train"], args.batch_size) torch.cuda.synchronize() tick = time.perf_counter() optimizer.zero_grad(set_to_none=True) scores = scorer.score_examples(rows) loss = torch.stack([-torch.log_softmax(s.float(), dim=0)[row["target"]] for s, row in zip(scores, rows)]).mean() if not torch.isfinite(loss): raise RuntimeError("Non-finite loss") loss.backward() gradient_norm = torch.nn.utils.clip_grad_norm_(trainable, 1.0, error_if_nonfinite=True) optimizer.step() torch.cuda.synchronize() completed = step item = {"step": step, "loss": loss.item(), "gradient_norm": gradient_norm.item(), "seconds": time.perf_counter() - tick, "cuda_peak_allocated_bytes": torch.cuda.max_memory_allocated()} history.append(item) print(json.dumps(item), flush=True) if step % 10 == 0: checkpoint(step) write_json(out / "training.json", history) checkpoint(completed) write_json(out / "training.json", history) results["training"] = {"completed_steps": completed, "seconds": time.perf_counter() - train_start, "median_step_seconds": statistics.median(x["seconds"] for x in history) if history else None} results["parameter_changes"] = { "probe_name": probe_name, "probe_max_abs_from_base": (dict(scorer.named_parameters())[probe_name].detach().float().cpu() - initial_probe).abs().max().item(), "head_weight_norm": scorer.head.weight.norm().item(), } if history: assert results["parameter_changes"]["probe_max_abs_from_base"] > 0 assert results["parameter_changes"]["head_weight_norm"] > 0 if STOP: results["status"] = "interrupted_checkpoint_saved" write_json(out / "metrics.json", results) return trained = evaluate(data["test"]) write_json(out / "trained_predictions.json", trained) results["trained_scalar"] = metrics(trained) calibration = evaluate(data["calibration"]) write_json(out / "calibration_predictions.json", calibration) # Fit ONE global temperature to the separate calibration split only. grid = torch.logspace(-1, 1, 81).tolist() def temperature_loss(t): return statistics.mean(-torch.log_softmax(torch.tensor(r["logits"]) / t, dim=0)[r["target"]].item() for r in calibration) temperature = min(grid, key=temperature_loss) calibrated = [{**r, "probabilities": torch.softmax(torch.tensor(r["logits"]) / temperature, 0).tolist(), "log_probabilities": torch.log_softmax(torch.tensor(r["logits"]) / temperature, 0).tolist()} for r in trained] write_json(out / "calibrated_predictions.json", calibrated) results["temperature"] = temperature results["calibrated_scalar"] = metrics(calibrated) results["status"] = "trained_pending_correctness" write_json(out / "metrics.json", results) scorer.eval() group = [r for r in data["test"] if r["group"] == "TEST-000"] parity_tolerance = 1e-4 if args.dtype == "float32" else 0.02 with torch.inference_mode(): full = scorer.score_examples(group) shared = scorer.scores_shared(group[0]["state"], group, branch_batch_size=4) again = scorer.scores_shared(group[0]["state"], group, branch_batch_size=4) isolated = [scorer.scores_shared(r["state"], [r], branch_batch_size=4)[0] for r in group] reversed_rows = [{**r, "choices": list(reversed(r["choices"]))} for r in group] reversed_scores = scorer.scores_shared(group[0]["state"], reversed_rows, branch_batch_size=4) def difference(a, b): return max((torch.softmax(x.float(), 0) - torch.softmax(y.float(), 0)).abs().max().item() for x, y in zip(a, b)) checks = {"full_vs_shared_probability_max_abs": difference(full, shared), "repeat_cache_probability_max_abs": difference(shared, again), "question_isolation_probability_max_abs": difference(shared, isolated), "candidate_permutation_probability_max_abs": difference(shared, [s.flip(0) for s in reversed_scores]), "tolerance_probability_abs": parity_tolerance} write_json(out / "correctness.json", checks) assert all(v <= parity_tolerance for k, v in checks.items() if k.endswith("max_abs")), checks # Check that the compact saved artifact reconstructs predictions, without reloading the base. saved = torch.load(out / "checkpoint.pt", map_location="cpu", weights_only=False) assert set(saved["trainable_state"]) == set(manifest["trainable_names"]) original = full with torch.no_grad(): for p in scorer.parameters(): if p.requires_grad: p.zero_() scorer.load_state_dict(saved["trainable_state"], strict=False) checks["checkpoint_reload_probability_max_abs"] = difference(original, scorer.score_examples(group)) assert checks["checkpoint_reload_probability_max_abs"] <= 0.00001 results["correctness"] = checks write_json(out / "correctness.json", checks) print(json.dumps({"event": "correctness", **checks}), flush=True) # Isolate actual shared-prefix timings against the SAME full-forward branch batching. # Small repetition count: raw values are retained and not advertised as production p95. benchmark = [] with torch.inference_mode(): for context_target in [128, 1024]: unit = " Shelf log: the warehouse is open and parcel records follow." state = unit while len(scorer.prefix_ids(state)) < context_target: state += unit state += " The label is blue." for questions_count in [1, 4, 16]: questions = [{"state": state, "question": "What color is the label?", "choices": ["red", "blue"]} for _ in range(questions_count)] for mode in ["full", "shared"]: durations = [] detail = [] for repeat in range(4): torch.cuda.synchronize() tick = time.perf_counter() if mode == "shared": scorer.scores_shared(state, questions, branch_batch_size=8) else: for j in range(0, len(questions), 4): scorer.score_examples(questions[j:j+4]) torch.cuda.synchronize() elapsed = (time.perf_counter() - tick) * 1000 if repeat: durations.append(elapsed) if mode == "shared": detail.append(scorer.last_shared_timings.copy()) benchmark.append({"context_target_tokens": context_target, "actual_prefix_tokens": len(scorer.prefix_ids(state)), "questions": questions_count, "candidates_each": 2, "mode": mode, "warm_ms_raw": durations, "warm_ms_median": statistics.median(durations), "shared_detail_raw": detail}) print(json.dumps({"event": "benchmark", **benchmark[-1]}), flush=True) write_json(out / "benchmark.json", benchmark) results["resources"] = {"peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_cuda_reserved_bytes": torch.cuda.max_memory_reserved(), "final_mem_available_bytes": mem_available(), "gpu_snapshot": command(["nvidia-smi", "--query-gpu=temperature.gpu,utilization.gpu,power.draw", "--format=csv"])} results["status"] = "complete" write_json(out / "metrics.json", results) print(json.dumps({"event": "complete", "output": str(out), "accuracy": results["trained_scalar"]["accuracy"], "nll": results["trained_scalar"]["nll"]}), flush=True) if __name__ == "__main__": main()