#!/usr/bin/env python3 """Audit the complete failure-retaining 27-cell ResNet R2 panel.""" import argparse import hashlib import json import math import os import sys ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) from experiments.resnet_crossover_r2 import ( DEPTHS, METHODS, DEFAULT_SELECTOR, r2_jobs, registry_sha256, selector_report, ) from experiments.crossover_hardware import assert_hardware_report def sha256(path): digest = hashlib.sha256() with open(path, "rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def read_json(path): with open(path, encoding="utf-8") as handle: return json.load(handle) def finite(value): return ( isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(float(value)) ) def expected_mode(method): return { "bp": "bp", "fa": "hfa", "dfa": "dfa", "clean_kp": "kp", "sdil": "kp_traffic", }[method] def history_metrics(record, job): method = job["method"] expected_epochs = job["epochs"] losses = [] validation_accuracies = [] history_examples = 0 if method == "ff": layers = record.get("layers", []) layer_indices = [row.get("layer") for row in layers] assert layer_indices == list(range(len(layers))) for layer in layers: epochs = layer.get("epochs", []) assert [row.get("epoch") for row in epochs] == list( range(1, len(epochs) + 1) ) assert len(epochs) <= expected_epochs losses.extend(row.get("loss") for row in epochs) history_examples += sum( int(row.get("examples", 0)) for row in epochs ) complete = ( len(layers) == job["depth"] and all( len(layer.get("epochs", [])) == expected_epochs for layer in layers ) ) completed_epochs = sum( len(layer.get("epochs", [])) for layer in layers ) else: epochs = record.get("epochs", []) assert [row.get("epoch") for row in epochs] == list( range(1, len(epochs) + 1) ) assert len(epochs) <= expected_epochs complete = len(epochs) == expected_epochs completed_epochs = len(epochs) if method in ("pepita", "ep", "dualprop"): losses = [row.get("train_loss") for row in epochs] history_examples = sum( int(row.get("train_examples", 0)) for row in epochs ) validation_accuracies = [ row["validation"]["accuracy"] for row in epochs if row.get("validation") is not None ] else: losses = [row.get("train_loss") for row in epochs] history_examples = sum( int(row.get("train_examples", 0)) for row in epochs ) validation_accuracies = [ row["eval_accuracy"] for row in epochs if row.get("eval_accuracy") is not None ] final = record.get("final", {}) final_accuracy = final.get("accuracy") final_loss = final.get("loss") finite_accuracies = [ float(value) for value in validation_accuracies + [final_accuracy] if finite(value) ] best_accuracy = max(finite_accuracies) if finite_accuracies else None losses_finite = bool(losses) and all(finite(value) for value in losses) final_finite = ( final.get("finite") is True and finite(final_accuracy) and (final_loss is None or finite(final_loss)) ) first_nonfinite = record.get("first_nonfinite_step") return { "complete_trajectory": complete, "epochs_completed": completed_epochs, "history_training_examples": history_examples, "all_training_losses_finite": losses_finite, "final_validation_accuracy": float(final_accuracy) if finite(final_accuracy) else None, "best_validation_accuracy": float(best_accuracy) if best_accuracy is not None else None, "final_validation_loss": float(final_loss) if finite(final_loss) else None, "first_nonfinite_step": first_nonfinite, "finite": ( complete and losses_finite and final_finite and first_nonfinite is None ), } def work_metrics(record, job, history): method = job["method"] if method in ("pepita", "ff", "ep", "dualprop"): work = record["work"] ordinary = int(work["ordinary_training_examples"]) validation = int(work["ordinary_validation_examples"]) presentations = int(work["training_example_presentations"]) relaxation = int(work["relaxation_example_passes"]) candidate = int(work["candidate_label_evaluation_presentations"]) queries = int(work["logical_task_loss_queries"]) local_vjp = int(work["local_vjp_example_evaluations"]) local_backward = int( work["local_target_backward_example_evaluations"] ) assert ordinary == history["history_training_examples"] if method == "ff": assert presentations == 2 * ordinary assert candidate == 10 * validation assert relaxation == 0 assert local_backward == 2 * ordinary assert local_vjp == 0 elif method == "pepita": assert presentations == 2 * ordinary assert relaxation == candidate == local_vjp == local_backward == 0 elif method == "ep": assert presentations == ordinary assert relaxation == 24 * ordinary + 20 * validation assert local_vjp == relaxation assert candidate == local_backward == 0 else: assert presentations == ordinary assert relaxation == 16 * ordinary assert local_vjp == relaxation assert candidate == local_backward == 0 total_work = { "forward_macs_per_example": work["forward_macs_per_example"], "feedforward_example_passes": work["feedforward_example_passes"], "relaxation_example_passes": relaxation, "local_vjp_example_evaluations": local_vjp, "candidate_label_evaluation_presentations": candidate, } else: work = record["work"] counters = record["counters"] ordinary = int(counters["ordinary_examples"]) validation = ( int(record["split"]["validation_examples"]) * int(record["evaluation_protocol"]["validation_evaluations"]) ) presentations = ordinary queries = int(work["logical_batch_loss_queries"]) total_work = { "forward_macs_per_example": work["forward_macs_per_example"], "total_macs_estimate": work["total_macs_estimate"], "elementwise_operations_estimate": work["elementwise_operations_estimate"], "total_forward_equivalent_examples": work["total_forward_equivalent_examples"], } assert ordinary == history["history_training_examples"] assert queries == 0 return { "ordinary_training_examples": ordinary, "ordinary_validation_examples": validation, "training_example_presentations": presentations, "logical_task_loss_queries": queries, "work": total_work, } def audit_job( job, expected_source, expected_selector, hardware_policy ): manifest_path = job["output"] + ".manifest.json" if not os.path.isfile(manifest_path): raise AssertionError(f"missing R2 manifest: {job['experiment_name']}") manifest = read_json(manifest_path) for key in ( "stage", "method", "architecture", "depth", "rate", "epochs", "lr_schedule", "experiment_name", "output", "timeout_seconds", "command", ): assert manifest[key] == job[key], ( f"{job['experiment_name']}: manifest drift in {key}" ) assert manifest["source"] == expected_source assert manifest["selector"] == expected_selector hardware = manifest["hardware_lock"] assert_hardware_report(hardware, hardware_policy) common = { "cell_id": f"resnet{job['depth']}::{job['method']}", "method": job["method"], "depth": job["depth"], "rate": job["rate"], "status": manifest["status"], "manifest": os.path.relpath(manifest_path, ROOT), "driver_wall_seconds": float(manifest["driver_wall_seconds"]), "physical_gpu_index": hardware["physical_gpu_index"], "physical_gpu_uuid": hardware["physical_gpu_uuid"], "output_sha256": manifest["output_sha256"], } if manifest["status"] != "completed": assert manifest["status"] in { "timeout", "nonzero_exit", "missing_output" } assert ( manifest["output_sha256"] is None if not manifest["output_exists"] else manifest["output_sha256"] == sha256(job["output"]) ) return { **common, "finite": False, "complete_trajectory": False, "epochs_completed": 0, "best_validation_accuracy": None, "final_validation_accuracy": None, "final_validation_loss": None, "first_nonfinite_step": None, "ordinary_training_examples": None, "ordinary_validation_examples": None, "training_example_presentations": None, "logical_task_loss_queries": None, "forward_parameter_count": None, "peak_memory_allocated_bytes": None, "total_wall_seconds": None, "work": None, } assert manifest["output_exists"] is True assert manifest["output_sha256"] == sha256(job["output"]) record = read_json(job["output"]) provenance = record["provenance"] assert provenance["git_commit"] == expected_source["git_commit"] assert provenance["git_tracked_dirty"] is False args = record["args"] for key, expected in { "depth": job["depth"], "width": 16, "seed": 0, "loader_seed": 0, "split_seed": 2027, "batch_size": 128, "epochs": job["epochs"], "train_limit": 0, "val_examples": 5000, "eval_split": "validation", "eval_every": 1, "augment_train": 1, "lr": job["rate"], "lr_schedule": job["lr_schedule"], "momentum": 0.9, "weight_decay": 1e-4, }.items(): assert args[key] == expected, ( f"{job['experiment_name']}: argument drift in {key}" ) if job["method"] in ("pepita", "ff", "ep", "dualprop"): assert args["method"] == job["method"] else: assert args["mode"] == expected_mode(job["method"]) split = record["split"] assert split["train_examples"] == 45_000 assert split["validation_examples"] == 5_000 assert split["split_seed"] == 2027 evaluation = record["evaluation_protocol"] assert evaluation["test_evaluations"] == 0 assert evaluation["test_used_for_selection"] is False architecture = record["architecture"] assert architecture["depth"] == job["depth"] assert architecture["base_width"] == 16 history = history_metrics(record, job) work = work_metrics(record, job, history) if history["complete_trajectory"]: expected_ordinary = 45_000 * job["epochs"] if job["method"] == "ff": expected_ordinary *= job["depth"] assert work["ordinary_training_examples"] == expected_ordinary hardware_record = record["hardware"] assert ( hardware_record["cuda_visible_devices"] == hardware["cuda_visible_devices"] ) assert ( hardware_record["cuda_device_name"] == hardware["physical_gpu_name"] ) parameter_count = architecture.get("forward_parameter_count") if parameter_count is None: parameter_count = architecture["forward_parameters"] return { **common, **history, **work, "forward_parameter_count": int(parameter_count), "peak_memory_allocated_bytes": int( hardware_record["peak_memory_allocated_bytes"] ), "total_wall_seconds": float( record.get("total_wall_seconds") if "total_wall_seconds" in record else record["timing"]["total_timed_wall_s"] ), } def main(): parser = argparse.ArgumentParser() parser.add_argument("--selector", default=DEFAULT_SELECTOR) parser.add_argument( "--out", default="results/resnet_crossover/r2_audit.json" ) args = parser.parse_args() selector = selector_report(args.selector) jobs = r2_jobs(selector["selected_rates"]) launch_path = os.path.join( ROOT, "results", "resnet_crossover", "r2_launch.json" ) assert os.path.isfile(launch_path), "missing ResNet R2 launch lock" launch = read_json(launch_path) assert launch["stage"] == "r2" assert launch["selector"] == selector assert launch["registry_sha256"] == registry_sha256(jobs) assert launch["num_jobs"] == 27 hardware_policy = launch["hardware_policy"] assert ( launch["allowed_physical_gpus"] == hardware_policy["allowed_physical_gpu_indices"] ) source = launch["source"] records = [ audit_job(job, source, selector, hardware_policy) for job in jobs ] assert len(records) == 27 assert len({record["cell_id"] for record in records}) == 27 assert { (record["method"], record["depth"]) for record in records } == {(method, depth) for method in METHODS for depth in DEPTHS} for depth in DEPTHS: counts = { record["forward_parameter_count"] for record in records if record["depth"] == depth and record["forward_parameter_count"] is not None } assert len(counts) <= 1, f"forward parameter mismatch at depth {depth}" failures = [ record["cell_id"] for record in records if not record["finite"] ] report = { "audit_status": "passed", "stage": "resnet_crossover_r2", "complete_grid": True, "failure_retaining": True, "num_expected_cells": 27, "num_audited_cells": len(records), "num_finite_cells": len(records) - len(failures), "failed_or_incomplete_cells": failures, "test_policy": "none", "source": source, "selector": selector, "hardware_policy": hardware_policy, "launch_lock": { "path": os.path.relpath(launch_path, ROOT), "sha256": sha256(launch_path), "registry_sha256": launch["registry_sha256"], }, "records": records, } os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) with open(args.out, "w", encoding="utf-8") as handle: json.dump(report, handle, indent=2, sort_keys=True) handle.write("\n") print( json.dumps( { "audit_status": report["audit_status"], "num_audited_cells": len(records), "num_finite_cells": report["num_finite_cells"], "failed_or_incomplete_cells": failures, }, indent=2, sort_keys=True, ) ) if __name__ == "__main__": main()