source/smoke_train.py

"""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()