"""Matched continuation study with optional primitive-fact gradients into the backbone.""" import argparse import hashlib import json import math import random import time from collections import Counter from datetime import datetime, timezone from pathlib import Path import torch from safetensors.torch import load_file, save_file from train_foundation_head import write from train_workflow_backbone import forward from workflow_workspace_supervision import WorkspaceSupervisor from veyra.constants import QUESTION_TYPES from veyra.data import read_records from veyra.option_model import OptionModel from veyra.proper_learning import distribution_losses, retention_losses from veyra.workflow_facts import auxiliary_targets from veyra.workflow_learning import SHARES, annotate, selection_key, summarize_logits def read(path): return json.loads(Path(path).read_text(encoding="utf-8")) def digest(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest() def development_risk(logits, tensors, indices, temperatures): types, valid = tensors["types"][indices], tensors["valid"][indices] values = logits / torch.tensor(temperatures)[types, None] confidence, prediction = values.masked_fill(~valid, -1e9).softmax(-1).max(-1) gold = tensors["gold"][indices] errors = 1 - tensors["targets"][indices].gather(1, prediction[:, None]).squeeze(1) errors = torch.where(gold >= 0, (prediction != gold).float(), errors) result = {} for type_id, kind in enumerate(QUESTION_TYPES): selected = torch.where(types == type_id)[0] order = selected[torch.argsort(confidence[selected], descending=True, stable=True)] count = len(order) counts = torch.arange(1, count + 1) risk = errors[order].cumsum(0) / counts boundaries = torch.ones(count, dtype=torch.bool) boundaries[:-1] = confidence[order[:-1]] != confidence[order[1:]] eligible = torch.where((risk <= 0.15) & (counts >= 20) & boundaries)[0] result[kind] = { "questions": count, "risk_at_60pct_coverage": float(risk[max(0, math.floor(count * 0.6) - 1)]), "maximum_coverage_at_15pct_error": (int(eligible[-1]) + 1) / count if len(eligible) else 0.0, } return result @torch.inference_mode() def assess(model, supervisor, enabled, records, root, rows, tensors, indices, targets, eligible): model.eval() logits, fact_prediction, fact_target = [], [], [] for i in indices.tolist(): with supervisor.observe(enabled and bool(eligible[i])): values, _ = forward(model, records[rows[i]["id"]], root) logits.append(values[0].cpu()) if enabled and eligible[i]: fact_prediction.append(supervisor.logits()[0].sigmoid().cpu()) fact_target.append(targets[i]) logits = torch.stack(logits) metrics = summarize_logits(logits, tensors, rows, indices, model.calibration.temperatures) result = { "metrics": metrics, "development_risk_diagnostic": development_risk( logits, tensors, indices, model.calibration.temperatures ), } if fact_prediction: prediction, target = torch.stack(fact_prediction), torch.stack(fact_target) observed = (target == 0) | (target == 1) result["auxiliary_diagnostic"] = { "questions": len(prediction), "condition_probability_mse": float((prediction - target).square().mean()), "observed_condition_accuracy": float( ((prediction[observed] >= 0.5) == (target[observed] == 1)).float().mean() ), "constant_half_mse": float((target - 0.5).square().mean()), } return result def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, required=True) parser.add_argument("--arm", required=True) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() if args.output.exists(): raise FileExistsError("workspace study outputs are immutable") config = read(args.config) if config["status"] != "declared_before_training": raise ValueError("declare the follow-up experiment before optimization") arm = config["arms"][args.arm] parent, data, features = map(Path, (config["parent"], config["records"], config["features"])) if ( digest(parent / "head.safetensors") != config["parent_weights_sha256"] or digest(parent / "manifest.json") != config["parent_manifest_sha256"] or digest(data) != config["records_sha256"] or digest(config["release_protocol"]) != config["release_protocol_sha256"] or digest(config["baseline_development"]) != config["baseline_development_sha256"] ): raise ValueError("declared model, data or acceptance protocol changed") spec, cache = read(features / "specification.json"), read(features / "manifest.json") if spec["records_sha256"] != digest(data) or spec["calibration_or_final_encoded"] is not False: raise ValueError("feature cache is not the declared training/development set") for name, key in ( ("features.safetensors", "features_sha256"), ("records.json", "metadata_sha256"), ("specification.json", "specification_sha256"), ): if digest(features / name) != cache[key]: raise ValueError("feature cache bytes changed") tensors, rows = load_file(features / "features.safetensors"), read(features / "records.json") records = {r.id: r for r in read_records(data) if r.split in {"train", "dev"}} annotate(rows, records.values()) targets, eligible = auxiliary_targets(rows, records.values()) train = [i for i, row in enumerate(rows) if row["split"] == "train"] dev = torch.tensor([i for i, row in enumerate(rows) if row["split"] == "dev"]) if int(eligible[train].sum()) != 3840 or int(eligible[dev].sum()) != 1200: raise ValueError("unexpected primitive supervision population") conditional_audit = Path("data/workflow-v12/conditional-target-audit.json") included_groups = {row["group"] for row in rows} audited_targets = { item["group"]: item["posterior"] for item in read(conditional_audit) if item["group"] in included_groups } for i, row in enumerate(rows): if row["domain"] == "uncertainty" and not torch.allclose( targets[i], torch.tensor(audited_targets[row["group"]]), atol=1e-6 ): raise ValueError("primitive target differs from the original evidence model") counts = Counter(rows[i]["domain"] for i in train) weights = {domain: SHARES[domain] * len(train) / count for domain, count in counts.items()} grouped = {} for i in train: grouped.setdefault(rows[i]["group"], []).append(i) orders = [] for epoch in range(config["epochs"]): groups = list(grouped.values()) random.Random(config["seed"] + epoch).shuffle(groups) orders.append([i for group in groups for i in group]) baseline = read(config["baseline_development"])["dev"] protocol = { "started_at_utc": datetime.now(timezone.utc).isoformat(), "configuration": config, "config_sha256": digest(args.config), "arm": args.arm, "source_sha256": { str(path): digest(path) for path in ( Path(__file__), Path("scripts/workflow_workspace_supervision.py"), Path("scripts/train_workflow_backbone.py"), Path("src/veyra/workflow_facts.py"), Path("src/veyra/workflow_learning.py"), Path("src/veyra/proper_learning.py"), Path("src/veyra/option_model.py"), ) }, "feature_cache": cache, "conditional_audit_sha256": digest(conditional_audit), "conditional_targets_match_audit": True, "domain_shares": SHARES, "calibration_or_final_used": False, "primary_experiment_candidates_changed": False, "inference_modules_added": 0, "generated_intermediate_tokens": 0, } write(args.output / "protocol.json", protocol) torch.manual_seed(config["seed"]) torch.set_num_threads(4) model = OptionModel.load(parent, local_files_only=True, merge=False) model.condition_readout.requires_grad_(False) supervisor = WorkspaceSupervisor(model, config["seed"]) adapters = [p for name, p in model.named_parameters() if ".lora_" in name] groups = [ {"params": adapters, "lr": config["learning_rate"]}, { "params": list(model.readout.parameters()) + list(model.binding_head.parameters()), "lr": config["learning_rate"] * 0.1, }, ] if arm["auxiliary_weight"] > 0: groups.append({"params": supervisor.decoder.parameters(), "lr": config["auxiliary_lr"]}) optimizer = torch.optim.AdamW(groups, weight_decay=0.01) accumulation = config["accumulation"] total_steps = len(orders) * math.ceil(len(train) / accumulation) def rate(step): if step < 30: return (step + 1) / 30 return 0.1 + 0.45 * (1 + math.cos(math.pi * min(1, (step - 30) / max(1, total_steps - 30)))) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, rate) temperatures = model.calibration.temperatures history, steps, processed, total_loss, auxiliary_count, auxiliary_sum = [], 0, 0, 0.0, 0, 0.0 started = time.perf_counter() torch.cuda.reset_peak_memory_stats() try: for epoch, order in enumerate(orders, 1): model.train() model.encoder.model.eval() for start in range(0, len(order), accumulation): batch = order[start : start + accumulation] optimizer.zero_grad(set_to_none=True) for i in batch: use_auxiliary = arm["auxiliary_weight"] > 0 and bool(eligible[i]) with supervisor.observe(use_auxiliary): logits, hidden = forward(model, records[rows[i]["id"]], data.parent) valid = tensors["valid"][i : i + 1].to(logits.device) types = tensors["types"][i : i + 1].to(logits.device) scale = torch.tensor(temperatures, device=logits.device)[types, None] scaled = logits / scale loss = distribution_losses( scaled, tensors["targets"][i : i + 1].to(logits.device), valid, types, config["method"], ).mean() if rows[i]["retention"]: reference = ( tensors["parent_logits"][i : i + 1].to(logits.device) / scale ) anchor = ( 1 - torch.nn.functional.cosine_similarity( hidden, tensors["states"][i : i + 1].to(hidden.device) ).mean() ) loss = ( 0.25 * loss + 4 * retention_losses(scaled, reference, valid).mean() + 0.1 * anchor ) loss = loss * weights[rows[i]["domain"]] if use_auxiliary: if not supervisor.captured.requires_grad: raise ValueError("primitive loss is detached from the backbone") aux_loss = torch.nn.functional.binary_cross_entropy_with_logits( supervisor.logits(), targets[i : i + 1].to(logits.device) ) if auxiliary_count == 0: gradients = torch.autograd.grad( aux_loss, adapters, retain_graph=True, allow_unused=True ) norms = [ float(g.detach().abs().max()) for g in gradients if g is not None ] if ( not norms or not all(math.isfinite(n) for n in norms) or max(norms) == 0 ): raise ValueError( "no finite nonzero primitive-to-backbone gradient" ) write( args.output / "auxiliary-gradient-path.json", { "backbone_parameters_with_gradient": len(norms), "backbone_parameters_with_nonzero_gradient": sum( n > 0 for n in norms ), "maximum_absolute_gradient": max(norms), "same_decision_forward": True, "optimizer_steps_before_check": steps, }, ) del gradients, norms loss = loss + arm["auxiliary_weight"] * aux_loss auxiliary_count += 1 auxiliary_sum += float(aux_loss.detach()) if not torch.isfinite(loss): raise ValueError("nonfinite training objective") (loss / len(batch)).backward() total_loss += float(loss.detach()) processed += 1 torch.nn.utils.clip_grad_norm_( [p for group in optimizer.param_groups for p in group["params"]], 1.0, error_if_nonfinite=True, ) optimizer.step() scheduler.step() steps += 1 if steps % 25 == 0: progress = { "epoch": epoch, "processed": processed, "total": len(train) * len(orders), "steps": steps, "mean_loss": total_loss / processed, "auxiliary_bce": auxiliary_sum / auxiliary_count if auxiliary_count else None, "elapsed_seconds": time.perf_counter() - started, } write(args.output / "progress.json", progress) print(json.dumps(progress), flush=True) result = assess( model, supervisor, arm["auxiliary_weight"] > 0, records, data.parent, rows, tensors, dev, targets, eligible, ) checkpoint = args.output / f"epoch-{epoch}" / "checkpoint" metadata = { "completed": True, "intermediate": True, "method": "workspace-" + args.arm, "seed": config["seed"], "epochs": epoch, "optimizer_steps": steps, "backbone_updated": True, "parent_weights_sha256": config["parent_weights_sha256"], "dataset_sha256": digest(data), "protocol_sha256": digest(args.output / "protocol.json"), "selected_dev": result["metrics"], "selection": selection_key(result["metrics"], baseline), "development_risk_diagnostic": result["development_risk_diagnostic"], "auxiliary_diagnostic": result.get("auxiliary_diagnostic"), "elapsed_seconds": time.perf_counter() - started, "peak_memory_mib": torch.cuda.max_memory_allocated() / 2**20, "final_evaluated_at_selection": False, "smoke": False, } model.save(checkpoint, metadata) if arm["auxiliary_weight"] > 0: save_file( { name: value.detach().cpu().contiguous() for name, value in supervisor.decoder.state_dict().items() }, checkpoint.parent / "auxiliary-training-only.safetensors", ) history.append({"checkpoint": str(checkpoint), **metadata}) write(args.output / "history.json", history) print( json.dumps( { "checkpoint": str(checkpoint), "selection": metadata["selection"], "auxiliary": result.get("auxiliary_diagnostic"), } ), flush=True, ) finally: supervisor.close() write( args.output / "complete.json", { "completed": True, "candidates": [item["checkpoint"] for item in history], "release_allowed": False, }, ) if __name__ == "__main__": main()