path: root/experiments/run.py

"""
SDIL main training / diagnostics driver.

Trains one of {bp, dfa, sdil} (sdil with ablation flags) on MNIST/FashionMNIST,
logging the quantities that actually test the hypothesis:
  - train loss, test accuracy
  - per-hidden-layer cos(innovation r_l, -grad h_l)      <- the headline metric
  - cos(raw apical a_l, -grad)  and  cos(A_l c, -grad)   <- residualization ablation
  - single-step loss-decrease ratio vs exact GD

Everything is JSON-logged for later plotting.
"""
import argparse
import json
import math
import os
import subprocess
import sys
import time

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step,
                       neutral_p_update)
from sdil.baselines import BPNet, dfa_config, evaluate
from sdil.local_baselines import FANet
from sdil import probes
from sdil.data import (get_dataset_splits, onehot, make_hierarchical,
                       make_teacher_student, make_tentmap,
                       split_training_loader)


REAL_DATASETS = ("mnist", "fmnist", "cifar10")
SYNTHETIC_DATASETS = ("teacher", "hierarchical", "tentmap")


def code_provenance():
    """Best-effort source revision metadata for reproducible result files."""
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    try:
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"], cwd=root, check=True,
            capture_output=True, text=True).stdout.strip()
        dirty = bool(subprocess.run(
            ["git", "status", "--porcelain", "--untracked-files=no"], cwd=root,
            check=True, capture_output=True, text=True).stdout.strip())
        return {"git_commit": commit, "git_dirty": dirty}
    except (OSError, subprocess.CalledProcessError):
        return {"git_commit": None, "git_dirty": None}


def device_sync(device):
    if str(device).startswith("cuda") and torch.cuda.is_available():
        torch.cuda.synchronize()


def reset_peak_memory(device):
    if str(device).startswith("cuda") and torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(torch.device(device))


def hardware_report(device):
    report = {"device": str(device), "torch_version": torch.__version__}
    if str(device).startswith("cuda") and torch.cuda.is_available():
        cuda_device = torch.device(device)
        props = torch.cuda.get_device_properties(cuda_device)
        report.update({
            "cuda_device_name": props.name,
            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
            "device_total_memory_bytes": props.total_memory,
            "peak_memory_allocated_bytes": torch.cuda.max_memory_allocated(cuda_device),
            "peak_memory_reserved_bytes": torch.cuda.max_memory_reserved(cuda_device),
        })
    else:
        report.update({
            "cuda_device_name": None,
            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
            "device_total_memory_bytes": None,
            "peak_memory_allocated_bytes": None,
            "peak_memory_reserved_bytes": None,
        })
    return report


def fixed_training_probe(train_loader, probe_bs, device):
    """Return a deterministic diagnostic batch without advancing data order.

    All project loaders expose their in-memory tensors.  Slicing those tensors
    avoids consuming the loader's shuffle generator and, importantly, keeps a
    held-out validation/test split out of gradient-alignment diagnostics.
    """
    if not hasattr(train_loader, "x") or not hasattr(train_loader, "y"):
        raise TypeError("diagnostic probes require an in-memory training loader")
    n = min(probe_bs, len(train_loader.x))
    if n <= 0:
        raise ValueError("diagnostic probe requires at least one training example")
    return train_loader.x[:n].to(device), train_loader.y[:n].to(device)


@torch.no_grad()
def residual_lesion_report(net, evaluation_loader, probe_x, fraction):
    """Measure whether a trained residual network uses its late blocks.

    Interior hidden layers are the constant-width residual blocks after the
    input projection.  The lesion bypasses the final ``ceil(fraction * n)``
    such blocks at inference, while branch/skip RMS ratios are measured on a
    fixed training-prefix probe.  Training is never rerun after selecting the
    lesion, and the output readout is unchanged.
    """
    if not getattr(net, "residual", False):
        raise ValueError("residual lesion requires a residual network")
    if not 0.0 < fraction <= 1.0:
        raise ValueError(f"lesion fraction must be in (0, 1], got {fraction}")
    interior = list(range(1, net.L - 1))
    if not interior:
        raise ValueError("residual lesion requires at least one interior block")
    n_lesion = max(1, math.ceil(fraction * len(interior)))
    lesioned = interior[-n_lesion:]
    lesioned_set = set(lesioned)

    def forward(x, lesion=False, collect_ratios=False):
        h = x
        ratios = []
        for layer in range(net.L - 1):
            pre = h
            u = pre @ net.W[layer].t() + net.b[layer]
            act = net.act(u)
            if layer >= 1:
                branch = net.res_alpha * act
                if collect_ratios:
                    branch_rms = branch.square().mean().sqrt()
                    skip_rms = pre.square().mean().sqrt()
                    ratios.append(float(branch_rms / (skip_rms + 1e-12)))
                h = pre if lesion and layer in lesioned_set else pre + branch
            else:
                h = act
        logits = h @ net.W[-1].t() + net.b[-1]
        return logits, ratios

    correct = 0
    total = 0
    total_loss = 0.0
    for x, y in evaluation_loader:
        logits, _ = forward(x, lesion=True)
        total_loss += F.cross_entropy(logits, y, reduction="sum").item()
        correct += (logits.argmax(1) == y).sum().item()
        total += y.shape[0]
    _, ratios = forward(probe_x, collect_ratios=True)
    return {
        "interior_layers": interior,
        "lesioned_layers": lesioned,
        "branch_to_skip_rms": ratios,
        "lesion_eval_acc": correct / total,
        "lesion_eval_loss": total_loss / total,
    }


def calibration_work_per_event(net, cfg, pert_mode=None, pert_ndirs=None):
    """Hardware-independent causal-calibration work for one minibatch.

    Forward equivalents use affine multiply-add work as the denominator. The
    returned loss evaluations count calls producing per-example scalar losses;
    the caller multiplies by the actual minibatch size for scalar observations.
    """
    directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs
    mode = cfg.pert_mode if pert_mode is None else pert_mode
    if mode == "simultaneous":
        return {
            "batch_loss_evaluations": 2 * directions,
            # The implementation performs one duplicate clean forward plus one
            # batched +/- forward for every direction.
            "forward_equivalent_batches": 1.0 + 2.0 * directions,
            "perturbation_batch_expansion": 2 * directions,
        }

    full_work = sum(weight.numel() for weight in net.W)
    tail_work = 0
    for hidden_layer in range(net.L - 1):
        tail_work += sum(net.W[j].numel() for j in range(hidden_layer + 1, net.L))
    return {
        # One clean baseline loss plus +/- tail loss for every hidden layer.
        "batch_loss_evaluations": 1 + 2 * directions * (net.L - 1),
        "forward_equivalent_batches": 1.0 + 2.0 * directions * tail_work / full_work,
        "perturbation_batch_expansion": 1,
    }


def build(args, device):
    sizes = [args.n_in] + [args.width] * args.depth + [args.n_out]
    if args.mode == "bp":
        net = BPNet(sizes, act=args.act, device=device, seed=args.seed,
                    w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual),
                    predictor_mode=args.predictor_mode,
                    vectorizer_mode=args.vectorizer_mode)
        cfg = SDILConfig(eta=args.eta, momentum=args.momentum)
        return net, cfg
    if args.mode == "fa":
        net = FANet(sizes, act=args.act, device=device, seed=args.seed,
                    w_scale=args.w_scale, nuis_rho=0.0,
                    residual=bool(args.residual),
                    predictor_mode=args.predictor_mode,
                    vectorizer_mode=args.vectorizer_mode,
                    b_scale=args.feedback_scale)
        return net, SDILConfig(eta=args.eta, momentum=args.momentum)
    net = SDILNet(sizes, act=args.act, device=device, seed=args.seed,
                  w_scale=args.w_scale, a_scale=args.a_scale,
                  nuis_rho=args.nuis_rho, feedback=args.feedback,
                  residual=bool(args.residual), predictor_mode=args.predictor_mode,
                  vectorizer_mode=args.vectorizer_mode,
                  traffic_mode=args.traffic_mode, nuis_seed=args.traffic_seed)
    if args.mode == "dfa":
        cfg = dfa_config(eta=args.eta, momentum=args.momentum)
    elif args.mode == "sdil":
        cfg = SDILConfig(
            eta=args.eta, eta_A=args.eta_A, eta_P=args.eta_P,
            use_residual=bool(args.use_residual), learn_A=bool(args.learn_A),
            learn_P=bool(args.learn_P), pert_sigma=args.pert_sigma,
            pert_every=args.pert_every, pert_ndirs=args.pert_ndirs,
            pert_mode=args.pert_mode,
            momentum=args.momentum, settle_steps=args.settle_steps,
            kappa=args.kappa, feedback=args.feedback,
            p_update_on_neutral=bool(args.p_neutral),
            normalize_delta=bool(args.normalize_delta),
            raw_scale_control=args.raw_scale_control,
            vectorizer_optimizer=args.vectorizer_optimizer,
            vectorizer_eps=args.vectorizer_eps)
    elif args.mode == "nodepert":
        if args.pert_every != 1:
            raise ValueError("direct node perturbation requires --pert_every 1")
        cfg = SDILConfig(
            eta=args.eta, use_residual=False, learn_A=False, learn_P=False,
            pert_sigma=args.pert_sigma, pert_every=args.pert_every,
            pert_ndirs=args.pert_ndirs,
            pert_mode=args.pert_mode, momentum=args.momentum,
            normalize_delta=bool(args.normalize_delta), direct_node_pert=True)
    else:
        raise ValueError(args.mode)
    return net, cfg


def load_task(args, device):
    """Load a real dataset or construct one fixed synthetic task.

    ``task_seed`` controls examples and the target function; ``seed`` controls
    only the student initialization.  A depth sweep therefore compares models
    on exactly the same compositional problem instead of silently changing the
    teacher between model seeds.
    """
    if args.dataset in REAL_DATASETS:
        train, validation, test, n_in, n_out, split = get_dataset_splits(
            args.dataset, batch_size=args.batch_size, device=device,
            train_limit=args.train_examples or None,
            val_examples=args.val_examples, split_seed=args.split_seed)
        if args.eval_split == "validation":
            if validation is None:
                raise ValueError("--eval_split validation requires --val_examples > 0")
            evaluation = validation
        else:
            evaluation = test
        split["evaluation_split"] = args.eval_split
        return train, evaluation, n_in, n_out, split
    common = dict(
        n_train=args.task_train_examples,
        n_test=args.task_test_examples,
        seed=args.task_seed,
        batch_size=args.batch_size,
        device=device,
    )
    if args.dataset == "teacher":
        task = make_teacher_student(
            n_in=args.task_n_in, n_classes=args.task_classes,
            t_depth=args.teacher_depth, t_width=args.teacher_width,
            residual=bool(args.teacher_residual), **common)
    elif args.dataset == "hierarchical":
        task = make_hierarchical(
            levels=args.task_levels, n_classes=args.task_classes, **common)
    elif args.dataset == "tentmap":
        task = make_tentmap(
            levels=args.task_levels, n_in=args.task_n_in, **common)
    else:
        raise ValueError(f"unknown dataset: {args.dataset}")
    train, test, n_in, n_out = task
    validation = None
    split_details = {}
    if args.val_examples:
        train, validation, split_details = split_training_loader(
            train, args.val_examples, args.split_seed, args.batch_size)
    if args.eval_split == "validation":
        if validation is None:
            raise ValueError("--eval_split validation requires --val_examples > 0")
        evaluation = validation
    else:
        evaluation = test
    split = {
        "dataset": args.dataset,
        "task_seed": args.task_seed,
        "train_examples": len(train.x),
        "test_examples": args.task_test_examples,
        "evaluation_split": args.eval_split,
        "synthetic_generator": True,
    }
    split.update(split_details)
    return train, evaluation, n_in, n_out, split


def train(args):
    device = args.device
    torch.manual_seed(args.seed)
    train_loader, eval_loader, n_in, n_out, split = load_task(args, device)
    args.n_in, args.n_out = n_in, n_out
    net, cfg = build(args, device)
    # Reset after persistent data/model allocation so the peak includes the
    # resident training state plus transient perturbation-batch expansion.
    reset_peak_memory(device)

    # Diagnostics use a fixed training prefix without advancing the shuffled
    # training iterator.  Held-out data are reserved for metric evaluation.
    px = py = poh = None
    if args.diagnostics != "none" or args.residual_lesion_fraction > 0:
        px, py = fixed_training_probe(train_loader, args.probe_bs, device)
        if args.diagnostics != "none":
            poh = onehot(py, n_out, device=device)

    log = {"args": vars(args), "split": split, "provenance": code_provenance(),
           "diagnostic_protocol": {
               "probe_source": "training_prefix" if px is not None else None,
               "probe_examples": len(px) if px is not None else 0,
               "schedule": args.diagnostics_schedule,
           },
           "steps": [], "final": {}}
    step = 0
    prev_error = None
    device_sync(device)
    t0 = time.time()
    train_wall_s = 0.0
    eval_wall_s = 0.0
    diagnostics_wall_s = 0.0
    inline_diagnostics_wall_s = 0.0
    warmup_examples = 0
    ordinary_forward_examples = 0
    perturbation_events = 0
    calibration_batch_loss_evaluations = 0
    calibration_example_loss_evaluations = 0
    calibration_forward_equivalent_examples = 0.0
    perturbation_batch_expansion = 0
    feedback_warmup_examples = 0
    feedback_warmup_events = 0

    # predictor warmup on neutral-period (c=0) drive, so P cancels the apical
    # nuisance before task plasticity relies on the residual (no-op when rho=0).
    if args.mode == "sdil" and args.learn_P and args.p_warmup_steps > 0 and args.nuis_rho > 0:
        device_sync(device)
        warmup_t0 = time.time()
        it = iter(train_loader)
        for _ in range(args.p_warmup_steps):
            try:
                wx, _ = next(it)
            except StopIteration:
                it = iter(train_loader)
                wx, _ = next(it)
            wx = wx.to(device)
            warmup_examples += wx.shape[0]
            neutral_p_update(net, wx, args.p_warmup_eta)
        device_sync(device)
        warmup_wall_s = time.time() - warmup_t0
    else:
        warmup_wall_s = 0.0

    # Feedback-first timescale separation. Causal perturbations fit A on a
    # stationary forward network before noisy predictions can move W into a
    # bad basin. The output readout is frozen as well. This is supervised
    # calibration work, not free initialization, so every forward and scalar
    # loss observation is included in the same hardware-independent ledger.
    if args.a_warmup_steps > 0:
        if args.mode != "sdil" or not args.learn_A:
            raise ValueError("--a_warmup_steps requires SDIL with --learn_A 1")
        device_sync(device)
        feedback_warmup_t0 = time.time()
        loader_state = (train_loader.g.get_state().clone()
                        if hasattr(train_loader, "g") else None)
        rng_devices = ([torch.cuda.current_device()]
                       if str(device).startswith("cuda") and torch.cuda.is_available()
                       else [])
        # Warmup length must not silently change the minibatch order or random
        # directions used by the subsequent joint phase. fork_rng and restoring
        # the loader generator isolate those nuisance differences while keeping
        # the learned A parameters.
        try:
            with torch.random.fork_rng(devices=rng_devices):
                it = iter(train_loader)
                for _ in range(args.a_warmup_steps):
                    try:
                        wx, wy = next(it)
                    except StopIteration:
                        it = iter(train_loader)
                        wx, wy = next(it)
                    wx, wy = wx.to(device), wy.to(device)
                    wyoh = onehot(wy, n_out, device=device)
                    apical_calibration_step(
                        net, wx, wy, wyoh, cfg,
                        pert_mode=args.a_warmup_mode,
                        pert_ndirs=args.a_warmup_ndirs)
                    feedback_warmup_examples += wx.shape[0]
                    feedback_warmup_events += 1
                    event = calibration_work_per_event(
                        net, cfg, pert_mode=args.a_warmup_mode,
                        pert_ndirs=args.a_warmup_ndirs)
                    perturbation_events += 1
                    calibration_batch_loss_evaluations += event["batch_loss_evaluations"]
                    calibration_example_loss_evaluations += (
                        event["batch_loss_evaluations"] * wx.shape[0])
                    calibration_forward_equivalent_examples += (
                        event["forward_equivalent_batches"] * wx.shape[0])
                    perturbation_batch_expansion = max(
                        perturbation_batch_expansion,
                        event["perturbation_batch_expansion"])
        finally:
            if loader_state is not None:
                train_loader.g.set_state(loader_state)
        device_sync(device)
        feedback_warmup_wall_s = time.time() - feedback_warmup_t0
    else:
        feedback_warmup_wall_s = 0.0

    for epoch in range(args.epochs):
        device_sync(device)
        train_t0 = time.time()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            ordinary_forward_examples += x.shape[0]
            yoh = onehot(y, n_out, device=device)
            if args.mode == "bp":
                loss = net.bp_step(x, y, cfg.eta, momentum=cfg.momentum)
            elif args.mode == "fa":
                loss = net.fa_step(x, y, yoh, cfg.eta, momentum=cfg.momentum)
            else:
                loss, aux = sdil_step(net, x, y, yoh, cfg, step, prev_error=prev_error)
                prev_error = aux["error"]
                if args.mode in ("sdil", "nodepert") and aux["did_pert"]:
                    event = calibration_work_per_event(net, cfg)
                    perturbation_events += 1
                    calibration_batch_loss_evaluations += event["batch_loss_evaluations"]
                    calibration_example_loss_evaluations += (
                        event["batch_loss_evaluations"] * x.shape[0])
                    calibration_forward_equivalent_examples += (
                        event["forward_equivalent_batches"] * x.shape[0])
                    perturbation_batch_expansion = max(
                        perturbation_batch_expansion,
                        event["perturbation_batch_expansion"])

            if step % args.log_every == 0:
                rec = {"step": step, "epoch": epoch, "train_loss": float(loss)}
                inline_diagnostics = (args.diagnostics != "none"
                                      and args.diagnostics_schedule == "inline")
                if inline_diagnostics:
                    device_sync(device)
                    diagnostics_t0 = time.time()
                if args.mode == "fa" and inline_diagnostics:
                    rec.update(probes.fa_alignment_report(net, px, py, poh))
                elif args.mode == "nodepert" and inline_diagnostics:
                    rec.update(probes.nodepert_alignment_report(net, px, py, cfg))
                elif args.mode != "bp" and inline_diagnostics:
                    al = probes.alignment_report(net, px, py, poh, cfg)
                    rec["cos_r_negg"] = al["cos_r_negg"]
                    rec["cos_innovation_negg"] = al["cos_innovation_negg"]
                    rec["cos_apical_negg"] = al["cos_apical_negg"]
                    rec["cos_Ac_negg"] = al["cos_Ac_negg"]
                    rec["r_norm"] = al["r_norm"]
                    rec["traffic_norm"] = al["traffic_norm"]
                    rec["traffic_residual_norm"] = al["traffic_residual_norm"]
                    rec["traffic_r2"] = al["traffic_r2"]
                    if (args.mode == "sdil" and args.diagnostics == "full"
                            and step % (args.log_every * 5) == 0):
                        rec["ldr"] = probes.loss_decrease_ratio(net, px, py, poh, cfg, step)
                if inline_diagnostics:
                    device_sync(device)
                    elapsed = time.time() - diagnostics_t0
                    diagnostics_wall_s += elapsed
                    inline_diagnostics_wall_s += elapsed
                log["steps"].append(rec)
            step += 1
            if args.max_steps and step >= args.max_steps:
                break
        if args.max_steps and step >= args.max_steps:
            device_sync(device)
            train_wall_s += time.time() - train_t0
            break

        device_sync(device)
        train_wall_s += time.time() - train_t0

        should_evaluate = args.eval_every > 0 and (epoch + 1) % args.eval_every == 0
        record = {"epoch_end": epoch, "step": step}
        msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f}"
        if should_evaluate:
            device_sync(device)
            eval_t0 = time.time()
            acc, tloss = evaluate(net, eval_loader)
            device_sync(device)
            eval_wall_s += time.time() - eval_t0
            metric = "val_acc" if args.eval_split == "validation" else "test_acc"
            msg += f" {metric} {acc:.4f}"
            record.update({"eval_split": args.eval_split,
                           "eval_acc": acc, "eval_loss": tloss})
            if args.eval_split == "test":
                record.update({"test_acc": acc, "test_loss": tloss})
            else:
                record.update({"val_acc": acc, "val_loss": tloss})
        if (args.diagnostics != "none"
                and args.diagnostics_schedule == "inline"):
            device_sync(device)
            diagnostics_t0 = time.time()
            if args.mode == "fa":
                al = probes.fa_alignment_report(net, px, py, poh)
                meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"])
                msg += f" mean_cos(fa,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_fa_negg']]}"
            elif args.mode == "nodepert":
                al = probes.nodepert_alignment_report(net, px, py, cfg)
                meancos = sum(al["cos_q_negg"]) / len(al["cos_q_negg"])
                msg += f" mean_cos(q,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_q_negg']]}"
            elif args.mode != "bp":
                al = probes.alignment_report(net, px, py, poh, cfg)
                meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"])
                msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}"
            device_sync(device)
            diagnostics_wall_s += time.time() - diagnostics_t0
        print(msg, flush=True)
        log["steps"].append(record)

    device_sync(device)
    eval_t0 = time.time()
    acc, tloss = evaluate(net, eval_loader)
    device_sync(device)
    eval_wall_s += time.time() - eval_t0
    log["final"] = {"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss,
                    "wall_s": None}
    if args.eval_split == "test":
        log["final"].update({"test_acc": acc, "test_loss": tloss})
    else:
        log["final"].update({"val_acc": acc, "val_loss": tloss})
    if args.mode == "fa" and args.diagnostics != "none":
        device_sync(device)
        diagnostics_t0 = time.time()
        log["final"].update(probes.fa_alignment_report(net, px, py, poh))
        device_sync(device)
        diagnostics_wall_s += time.time() - diagnostics_t0
    elif args.mode == "nodepert" and args.diagnostics != "none":
        device_sync(device)
        diagnostics_t0 = time.time()
        log["final"].update(probes.nodepert_alignment_report(net, px, py, cfg))
        device_sync(device)
        diagnostics_wall_s += time.time() - diagnostics_t0
    elif args.mode != "bp" and args.diagnostics != "none":
        device_sync(device)
        diagnostics_t0 = time.time()
        al = probes.alignment_report(net, px, py, poh, cfg)
        log["final"]["cos_r_negg"] = al["cos_r_negg"]
        log["final"]["cos_innovation_negg"] = al["cos_innovation_negg"]
        log["final"]["cos_apical_negg"] = al["cos_apical_negg"]
        log["final"]["cos_Ac_negg"] = al["cos_Ac_negg"]
        log["final"]["r_norm"] = al["r_norm"]
        log["final"]["g_norm"] = al["g_norm"]
        log["final"]["traffic_norm"] = al["traffic_norm"]
        log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"]
        log["final"]["traffic_r2"] = al["traffic_r2"]
        device_sync(device)
        diagnostics_wall_s += time.time() - diagnostics_t0
    if args.residual_lesion_fraction > 0:
        device_sync(device)
        lesion_t0 = time.time()
        report = residual_lesion_report(
            net, eval_loader, px, args.residual_lesion_fraction)
        device_sync(device)
        eval_wall_s += time.time() - lesion_t0
        report["lesion_acc_drop"] = acc - report["lesion_eval_acc"]
        log["final"]["residual_lesion"] = report
        log["lesion_protocol"] = {
            "fraction_of_interior_blocks": args.residual_lesion_fraction,
            "selection": "final_contiguous_interior_blocks",
            "probe_source": "training_prefix",
            "probe_examples": len(px),
            "evaluation_split": args.eval_split,
        }
    device_sync(device)
    log["final"]["wall_s"] = time.time() - t0
    log["timing"] = {
        "warmup_wall_s": warmup_wall_s,
        "feedback_warmup_wall_s": feedback_warmup_wall_s,
        "training_loop_wall_s": train_wall_s,
        "diagnostics_wall_s": diagnostics_wall_s,
        "inline_diagnostics_wall_s": inline_diagnostics_wall_s,
        "optimizer_wall_s_excluding_inline_diagnostics": max(
            0.0, train_wall_s - inline_diagnostics_wall_s),
        "evaluation_wall_s": eval_wall_s,
    }
    log["cost"] = {
        "train_steps": step,
        "ordinary_training_forward_examples": ordinary_forward_examples,
        "predictor_warmup_forward_examples": warmup_examples,
        "feedback_warmup_forward_examples": feedback_warmup_examples,
        "feedback_warmup_perturbation_events": feedback_warmup_events,
        "perturbation_events": perturbation_events,
        "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations,
        "calibration_example_loss_evaluations": calibration_example_loss_evaluations,
        "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples,
        "training_forward_equivalent_examples": (
            ordinary_forward_examples + warmup_examples + feedback_warmup_examples
            + calibration_forward_equivalent_examples),
        "max_perturbation_batch_expansion": perturbation_batch_expansion,
    }
    log["hardware"] = hardware_report(device)
    os.makedirs(args.outdir, exist_ok=True)
    outpath = os.path.join(args.outdir, f"{args.tag}.json")
    with open(outpath, "w") as f:
        json.dump(log, f)
    print(f"[{args.tag}] DONE {args.eval_split}_acc={acc:.4f} -> {outpath}", flush=True)
    return log


def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--mode", default="sdil",
                   choices=["bp", "fa", "dfa", "sdil", "nodepert"])
    p.add_argument("--dataset", default="mnist",
                   choices=list(REAL_DATASETS + SYNTHETIC_DATASETS))
    p.add_argument("--depth", type=int, default=3)          # hidden layers
    p.add_argument("--width", type=int, default=256)
    p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"])
    p.add_argument("--residual", type=int, default=0)       # skip connections (deep no-BN)
    p.add_argument("--residual_lesion_fraction", type=float, default=0.0,
                   help="after training, bypass this fraction of final interior residual blocks")
    p.add_argument("--epochs", type=int, default=15)
    p.add_argument("--batch_size", type=int, default=128)
    p.add_argument("--train_examples", type=int, default=0,
                   help="0 uses the full training split")
    p.add_argument("--val_examples", type=int, default=0,
                   help="stratified validation examples held out from training")
    p.add_argument("--split_seed", type=int, default=2027)
    p.add_argument("--eval_split", default="test", choices=["validation", "test"])
    p.add_argument("--task_seed", type=int, default=0,
                   help="fixed target/data seed, separate from student --seed")
    p.add_argument("--task_train_examples", type=int, default=50000)
    p.add_argument("--task_test_examples", type=int, default=10000)
    p.add_argument("--task_levels", type=int, default=8)
    p.add_argument("--task_n_in", type=int, default=128)
    p.add_argument("--task_classes", type=int, default=10)
    p.add_argument("--teacher_depth", type=int, default=8)
    p.add_argument("--teacher_width", type=int, default=64)
    p.add_argument("--teacher_residual", type=int, default=1)
    p.add_argument("--eta", type=float, default=0.05)
    p.add_argument("--eta_A", type=float, default=0.02)
    p.add_argument("--eta_P", type=float, default=0.002)
    p.add_argument("--momentum", type=float, default=0.9)
    p.add_argument("--w_scale", type=float, default=1.0)
    p.add_argument("--a_scale", type=float, default=1.0)
    p.add_argument("--feedback_scale", type=float, default=1.0)
    p.add_argument("--pert_sigma", type=float, default=1e-2)
    p.add_argument("--pert_every", type=int, default=4)
    p.add_argument("--pert_ndirs", type=int, default=4)
    p.add_argument("--pert_mode", default="layerwise", choices=["layerwise", "simultaneous"])
    p.add_argument("--use_residual", type=int, default=1)
    p.add_argument("--raw_scale_control", default="none",
                   choices=["none", "match_innovation_norm"])
    p.add_argument("--learn_A", type=int, default=1)
    p.add_argument("--learn_P", type=int, default=1)
    p.add_argument("--p_neutral", type=int, default=1)      # P update on neutral (c=0) drive
    p.add_argument("--p_warmup_steps", type=int, default=200)  # pre-task neutral P warmup
    p.add_argument("--p_warmup_eta", type=float, default=0.05)
    p.add_argument("--a_warmup_steps", type=int, default=0,
                   help="A-only causal-calibration prefix with all forward weights frozen")
    p.add_argument("--a_warmup_ndirs", type=int, default=4)
    p.add_argument("--a_warmup_mode", default="layerwise",
                   choices=["layerwise", "simultaneous"])
    p.add_argument("--nuis_rho", type=float, default=0.0)
    p.add_argument("--traffic_seed", type=int, default=1234,
                   help="ordinary apical-traffic projection seed; independent of model/data seeds")
    p.add_argument("--traffic_mode", default="soma",
                   choices=["none", "soma", "topdown", "mixed"])
    p.add_argument("--predictor_mode", default="diagonal", choices=["diagonal", "full"])
    p.add_argument("--vectorizer_mode", default="linear",
                   choices=["linear", "soma_gated", "context_gated"])
    p.add_argument("--vectorizer_optimizer", default="sgd", choices=["sgd", "nlms"])
    p.add_argument("--vectorizer_eps", type=float, default=1e-6)
    p.add_argument("--normalize_delta", type=int, default=0)
    p.add_argument("--settle_steps", type=int, default=0)
    p.add_argument("--kappa", type=float, default=0.0)
    p.add_argument("--feedback", default="error", choices=["error", "error_deriv"])
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--log_every", type=int, default=50)
    p.add_argument("--diagnostics", default="full", choices=["none", "alignment", "full"])
    p.add_argument("--diagnostics_schedule", default="inline",
                   choices=["inline", "final"],
                   help="run probes throughout training or only after final evaluation")
    p.add_argument("--eval_every", type=int, default=1,
                   help="0 evaluates only once at the end")
    p.add_argument("--max_steps", type=int, default=0)      # 0 = no cap (smoke only)
    p.add_argument("--probe_bs", type=int, default=512)
    p.add_argument("--outdir", default="results")
    p.add_argument("--tag", default="sdil_run")
    return p.parse_args()


if __name__ == "__main__":
    train(get_args())