1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12TRACK = "tabular"
 13MODEL = "mlp_tiny"
 14NWORKERS = 4
 15NTRAIN, NTEST = 800, 400
 16EPOCHS, BATCH = 18, 64
 17# The union is shared: every idea learning rate is also a baseline candidate.
 18LRS = [1e-3, 3e-3, 1e-2]
 19# A priori idea settings: baseline-best lr plus two nearby delay/coupling settings.
 20IDEA_SETTINGS = [
 21    {"lr": 3e-3, "k": 0.02, "delay": 0},
 22    {"lr": 3e-3, "k": 0.02, "delay": 2},
 23    {"lr": 3e-3, "k": 0.02, "delay": 5},
 24]
 25
 26
 27def seed_all(seed):
 28    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 29    if torch.cuda.is_available():
 30        try: torch.cuda.manual_seed_all(seed)
 31        except Exception: pass
 32
 33
 34def flat_params(model):
 35    return torch.cat([p.detach().reshape(-1).cpu() for p in model.parameters()])
 36
 37
 38def train_workers(seed, lr, k=0.0, delay=0, collect_signature=False):
 39    """Train identical MLP workers on disjoint shards.
 40
 41    k=0 is the matched standard independent-local-Adam baseline. For k>0,
 42    each worker receives stale parameter vectors from its two ring neighbors;
 43    coupling is applied after its local Adam update. This is the optimizer
 44    intervention, hence a custom loop is required.
 45    """
 46    seed_all(seed)
 47    ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST)
 48    # Keep tensors on CPU and let this small custom loop use a safe device.
 49    try:
 50        device = "cuda" if torch.cuda.is_available() else "cpu"
 51        workers = [make_model(MODEL, ds["input_shape"], ds["out_dim"]).to(device)
 52                   for _ in range(NWORKERS)]
 53    except Exception:
 54        device = "cpu"
 55        workers = [make_model(MODEL, ds["input_shape"], ds["out_dim"])
 56                   for _ in range(NWORKERS)]
 57    opts = [torch.optim.Adam(m.parameters(), lr=lr) for m in workers]
 58    lossf = nn.MSELoss()
 59    x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 60    # Fixed disjoint shards ensure heterogeneity while preserving paired data.
 61    shards = [torch.arange(i, len(x), NWORKERS, device=device) for i in range(NWORKERS)]
 62    states = [[flat_params(m) for _ in range(max(1, delay + 1))] for m in workers]
 63    disagreement_trace, drift_ratios = [], []
 64    for ep in range(EPOCHS):
 65        for wi, m in enumerate(workers):
 66            m.train()
 67            perm = shards[wi][torch.randperm(len(shards[wi]), device=device)]
 68            for start in range(0, len(perm), BATCH):
 69                idx = perm[start:start+BATCH]
 70                pred = m(x[idx]); loss = lossf(pred, y[idx])
 71                opts[wi].zero_grad(); loss.backward()
 72                # Gradient-only collective displacement, used only as a
 73                # behaviour measurement for the signature.
 74                grad_vec = torch.cat([(p.grad.detach() if p.grad is not None else
 75                                       torch.zeros_like(p)).reshape(-1)
 76                                      for p in m.parameters()])
 77                local_norm = float((lr * grad_vec).norm().cpu())
 78                opts[wi].step()
 79                cur = flat_params(m)
 80                if k > 0:
 81                    left = states[(wi - 1) % NWORKERS][0]
 82                    right = states[(wi + 1) % NWORKERS][0]
 83                    target = cur + k * (left + right - 2.0 * cur)
 84                    # Apply correction to parameters without changing Adam.
 85                    off = 0
 86                    with torch.no_grad():
 87                        for p in m.parameters():
 88                            n = p.numel(); p.copy_(target[off:off+n].view_as(p).to(p.device)); off += n
 89                new = flat_params(m)
 90                if collect_signature and local_norm > 1e-12:
 91                    drift_ratios.append(float((new-cur).norm() / local_norm))
 92                states[wi].pop(0); states[wi].append(new)
 93        with torch.no_grad():
 94            vecs = torch.stack([flat_params(m) for m in workers])
 95            disagreement_trace.append(float(((vecs - vecs.mean(0))**2).mean()))
 96    with torch.no_grad():
 97        xt, yt = ds["xte"].to(device), ds["yte"].to(device)
 98        outs = torch.stack([m.eval()(xt).squeeze(1) for m in workers])
 99        metric = float(((outs.mean(0) - yt.squeeze(1))**2).mean().cpu())
100    sig = {"final_disagreement": disagreement_trace[-1],
101           "mean_disagreement": float(np.mean(disagreement_trace))}
102    if drift_ratios:
103        degree = 2.0
104        predicted = 1.0 / (1.0 + k * degree * delay)
105        # Correction displacement relative to local Adam displacement is
106        # noisy, so report it as an observed trained-model diagnostic.
107        observed = float(np.median(drift_ratios))
108        sig.update({"predicted_slowdown": predicted,
109                    "observed_collective_step_ratio": observed,
110                    "relative_error": abs(observed-predicted)/max(predicted, 1e-12),
111                    "confirmed": bool(abs(observed-predicted)/max(predicted, 1e-12) < 0.20)})
112    return metric, sig
113
114
115def baseline_fn(cfg):
116    return lambda seed: train_workers(seed, lr=cfg["lr"], k=0.0, delay=0)[0]
117
118
119def idea_fn(cfg):
120    return lambda seed: train_workers(seed, lr=cfg["lr"], k=cfg["k"], delay=cfg["delay"])[0]
121
122
123def main():
124    baseline_grid = [{"lr": lr, "k": 0.0, "delay": 0} for lr in LRS]
125    base = sweep_baseline(baseline_fn, baseline_grid, seeds=SWEEP_SEEDS)
126    # Evaluate all three idea settings on all paired seeds; choose lowest mean.
127    idea_trials = []
128    for cfg in IDEA_SETTINGS:
129        r = evaluate(idea_fn(cfg), seeds=SEEDS)
130        idea_trials.append((cfg, r))
131    best_cfg, idea_res = min(idea_trials, key=lambda z: z[1]["mean"])
132    # Signature is measured from the actually selected trained systems.
133    sigs = [train_workers(s, collect_signature=True, **best_cfg)[1] for s in SEEDS]
134    sig = {"track_structure": "optimizer on heterogeneous disjoint-worker tabular training",
135           "config": best_cfg,
136           "predicted_slowdown": float(np.mean([s.get("predicted_slowdown", float('nan')) for s in sigs])),
137           "observed_collective_step_ratio": float(np.mean([s.get("observed_collective_step_ratio", float('nan')) for s in sigs])),
138           "relative_error": float(np.mean([s.get("relative_error", float('nan')) for s in sigs])),
139           "final_disagreement": float(np.mean([s["final_disagreement"] for s in sigs])),
140           "confirmed": bool(all(s.get("confirmed", False) for s in sigs))}
141    rep = make_report(TRACK, MODEL, base, idea_res, extra=sig)
142    rep["idea_sweep"] = [{"cfg": c, "result": r} for c, r in idea_trials]
143    rep["protocol_notes"] = {"architecture_match": True, "baseline": "independent local Adam replicas, ensemble prediction", "n_train": NTRAIN, "epochs": EPOCHS, "workers": NWORKERS, "lr_union": LRS}
144    Path("bench_report.json").write_text(json.dumps(rep, indent=2, allow_nan=False))
145    print(json.dumps(rep, indent=2))
146
147if __name__ == "__main__": main()