#!/usr/bin/env python3 """Training-only localization of the frozen MT-1 nonfinite failure. This reproduces MT-1 initialization, traffic calibration, neutral predictor warmup, and task minibatch order, but performs no validation or test evaluation. It stops after the first update that makes any model, optimizer, feedback, predictor, or BatchNorm tensor nonfinite. """ import argparse import json import math import os import subprocess import sys import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import (CIFARKPMixedTrafficResNet, ConvSDILConfig, conv_kp_mixed_traffic_step) from sdil.data import DATA_DIR, get_cifar_image_splits ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) TRAINING_INACTIVE_GROUPS = {"bn_running_mean", "bn_running_var"} def provenance(): def run(command): return subprocess.run( command, cwd=ROOT, check=True, capture_output=True, text=True).stdout.strip() return { "git_commit": run(["git", "rev-parse", "HEAD"]), "git_tracked_dirty": bool(run( ["git", "status", "--porcelain", "--untracked-files=no"])), } def tensor_groups(net): return { "forward_weight": net.W, "forward_bn_scale": net.gamma, "forward_bn_bias": net.beta, "forward_readout": [net.W_out, net.b_out], "feedback_weight": net.Q, "feedback_readout": [net.R_out], "predictor_slope": net.P_traffic, "predictor_bias": net.P_traffic_bias, "forward_momentum": net.mW, "forward_bn_scale_momentum": net.mgamma, "forward_bn_bias_momentum": net.mbeta, "forward_readout_momentum": [net.mW_out, net.mb_out], "feedback_momentum": net.mQ, "feedback_readout_momentum": [net.mR_out], "bn_running_mean": net.running_mean, "bn_running_var": net.running_var, } def network_state(net): output = {} for name, tensors in tensor_groups(net).items(): bad = [] max_abs = 0.0 for index, value in enumerate(tensors): finite = bool(torch.isfinite(value).all()) if finite: max_abs = max(max_abs, float(value.abs().max())) else: bad.append(index) output[name] = { "all_finite": not bad, "nonfinite_indices": bad, "max_abs_over_finite_tensors": max_abs, } return output def nonfinite_groups(state): return {name: value["nonfinite_indices"] for name, value in state.items() if not value["all_finite"]} def main(): parser = argparse.ArgumentParser() parser.add_argument("--rule", choices=("raw", "matched", "innovation"), default="innovation") parser.add_argument("--predictor_mode", choices=("nlms20", "closed_form"), default="nlms20") parser.add_argument("--predictor_every", type=int, default=16) parser.add_argument("--stability_margin", type=float, default=0.0) parser.add_argument("--neutral_projection", action="store_true") parser.add_argument("--device", default="cuda") parser.add_argument("--max_steps", type=int, default=352) parser.add_argument("--out", required=True) args = parser.parse_args() if args.max_steps < 1: raise ValueError("max_steps must be positive") if args.predictor_every < 0: raise ValueError("predictor_every must be nonnegative") if args.stability_margin < 0: raise ValueError("stability_margin must be nonnegative") torch.manual_seed(0) if str(args.device).startswith("cuda"): if not torch.cuda.is_available(): raise RuntimeError("CUDA device requested but unavailable") torch.cuda.manual_seed_all(0) train, _, _, input_shape, n_out, split = get_cifar_image_splits( batch_size=128, data_dir=DATA_DIR, device=args.device, val_examples=5000, split_seed=2027, loader_seed=0, augment_train=True) if input_shape != (3, 32, 32) or n_out != 10: raise AssertionError("unexpected CIFAR dimensions") net = CIFARKPMixedTrafficResNet( depth=20, base_width=16, n_classes=10, device=args.device, seed=0, weight_scale=1.0, residual_scale=1.0, normalization="batchnorm", feedback_seed=None, feedback_scale=1.0, traffic_seed=4000) net.traffic_rule = args.rule config = ConvSDILConfig( eta=0.1, eta_output=0.1, eta_A=0.0, eta_P=0.1, momentum=0.9, weight_decay=1e-4, learn_A=False, learn_P=True) config.validate() loader_state = train.g.get_state().clone() calibration_x = train.x[:64] calibration_y = train.y[:64] calibration_forward = net.forward( calibration_x, return_cache=True, training=True, update_stats=False) calibration_error = ( torch.softmax(calibration_forward["logits"], dim=1) - torch.nn.functional.one_hot( calibration_y, 10).to(calibration_forward["logits"].dtype)) calibration_instruction = net.hierarchical_teaching( calibration_error, calibration_forward) calibration = net.calibrate_traffic_gain( calibration_instruction, calibration_forward["hiddens"], 4.0) closed_form_fit = None warmup_mse = [] if args.predictor_mode == "closed_form": closed_form_fit = net.predictor_closed_form_fit( calibration_forward["hiddens"], stability_margin=args.stability_margin) warmup_mse.append(closed_form_fit["mse"]) else: iterator = iter(train) for _ in range(20): x, _ = next(iterator) forward = net.forward(x, training=True, update_stats=False) warmup_mse.append(net.predictor_step(forward["hiddens"], 0.1)) del calibration_forward, calibration_error, calibration_instruction train.g.set_state(loader_state) loader_state_restored = torch.equal(train.g.get_state(), loader_state) audit_forward = net.forward( calibration_x, training=True, update_stats=False) post_warmup_ratio = net.predictor_traffic_residual_rms_ratio( audit_forward["hiddens"]) if args.predictor_mode == "nlms20": del forward del audit_forward initial_state = network_state(net) if nonfinite_groups(initial_state): raise AssertionError("network is nonfinite before task training") trajectory = [] first_any_nonfinite_step = None first_any_nonfinite = {} first_nonfinite_by_group = {} first_training_failure_step = None first_training_failure_groups = {} first_nonfinite_metrics = [] first_nonfinite_by_metric = {} for step, (x, y) in enumerate(train): if step >= args.max_steps: break result = conv_kp_mixed_traffic_step( net, x, y, config, step, args.rule, predictor_every=args.predictor_every, neutral_projection=args.neutral_projection) state = network_state(net) bad = nonfinite_groups(state) for name, indices in bad.items(): first_nonfinite_by_group.setdefault(name, { "step": step + 1, "indices": indices}) if bad and first_any_nonfinite_step is None: first_any_nonfinite_step = step + 1 first_any_nonfinite = bad metric_names = [ key for key in ( "loss", "teaching_rms", "instruction_rms", "raw_apical_rms", "innovation_rms", "traffic_rms") if not math.isfinite(float(result[key]))] if (result["predictor_mse"] is not None and not math.isfinite(float(result["predictor_mse"]))): metric_names.append("predictor_mse") for name in metric_names: first_nonfinite_by_metric.setdefault(name, step + 1) active_metric_names = [ name for name in metric_names if name in ( "loss", "teaching_rms", "instruction_rms", "predictor_mse")] active_bad = {name: indices for name, indices in bad.items() if name not in TRAINING_INACTIVE_GROUPS} trajectory.append({ "step": step + 1, "batch_loss": result["loss"], "teaching_rms": result["teaching_rms"], "instruction_rms": result["instruction_rms"], "raw_apical_rms": result["raw_apical_rms"], "innovation_rms": result["innovation_rms"], "traffic_rms": result["traffic_rms"], "predictor_updated": result["did_predictor_update"], "predictor_mse": result["predictor_mse"], "neutral_projection": result["neutral_projection"], "parameter_state": state, }) if active_bad or active_metric_names: first_training_failure_step = step + 1 first_training_failure_groups = active_bad first_nonfinite_metrics = active_metric_names break numeric = [] for row in trajectory: numeric.extend(float(row[key]) for key in ( "batch_loss", "teaching_rms", "instruction_rms", "raw_apical_rms", "innovation_rms", "traffic_rms")) if row["predictor_mse"] is not None: numeric.append(float(row["predictor_mse"])) output = { "protocol": ("kp_dynamic_neutral_projection_diagnosis_v1" if args.neutral_projection else "kp_mixed_traffic_nonfinite_diagnosis_v1"), "scope": "training_only_no_validation_or_test_evaluation", "provenance": provenance(), "rule": args.rule, "predictor_mode": args.predictor_mode, "predictor_every": args.predictor_every, "stability_margin": args.stability_margin, "neutral_projection": args.neutral_projection, "max_steps": args.max_steps, "split": split, "traffic_calibration": calibration, "predictor_warmup": { "steps": len(warmup_mse), "first_mse": warmup_mse[0], "last_mse": warmup_mse[-1], "post_warmup_traffic_residual_rms_ratio": post_warmup_ratio, "task_loader_state_restored": loader_state_restored, "closed_form_fit": closed_form_fit, }, "initial_parameter_state": initial_state, "trajectory": trajectory, "first_any_nonfinite_step": first_any_nonfinite_step, "first_any_nonfinite_groups": first_any_nonfinite, "first_nonfinite_by_group": first_nonfinite_by_group, "first_nonfinite_by_metric": first_nonfinite_by_metric, "first_training_failure_step": first_training_failure_step, "first_training_failure_groups": first_training_failure_groups, "first_nonfinite_metrics": first_nonfinite_metrics, "trajectory_metrics_finite": all(math.isfinite(value) for value in numeric), "validation_evaluations": 0, "test_evaluations": 0, } os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) with open(args.out, "w") as handle: json.dump(output, handle, indent=2, sort_keys=True) handle.write("\n") print(json.dumps({ "out": args.out, "first_any_nonfinite_step": first_any_nonfinite_step, "first_any_nonfinite_groups": first_any_nonfinite, "first_training_failure_step": first_training_failure_step, "first_training_failure_groups": first_training_failure_groups, "first_nonfinite_metrics": first_nonfinite_metrics, "steps_recorded": len(trajectory), }, indent=2)) if __name__ == "__main__": main()