1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6torch.set_num_threads(2)
  7
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import sweep_baseline, evaluate, make_report, get_dataset
 10
 11META = {"name": "congestion_moe_sequence", "domain": "moe-routing", "description": "Registered multi-token expert regression track"}
 12
 13DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 14# CUDA initialization is known to be flaky on this host; train_model-style fallback.
 15def safe_device():
 16    if DEVICE != "cuda": return "cpu"
 17    try:
 18        torch.zeros(1, device="cuda").sum().item()
 19        return "cuda"
 20    except Exception:
 21        return "cpu"
 22
 23
 24def project_simplex(v, z=1.0):
 25    # Batched Euclidean projection, differentiable almost everywhere.
 26    u, _ = torch.sort(v, dim=-1, descending=True)
 27    cssv = torch.cumsum(u, dim=-1) - z
 28    ind = torch.arange(1, v.shape[-1] + 1, device=v.device, dtype=v.dtype)
 29    cond = u - cssv / ind > 0
 30    rho = cond.sum(dim=-1).clamp_min(1).long() - 1
 31    theta = cssv.gather(-1, rho.unsqueeze(-1)).squeeze(-1) / (rho + 1).to(v.dtype)
 32    return (v - theta.unsqueeze(-1)).clamp_min(0)
 33
 34
 35class TinyMoE(nn.Module):
 36    def __init__(self, d=4, h=8, experts=4, temperature=0.7, idea=False, steps=4, eta=0.8):
 37        super().__init__()
 38        self.experts = nn.ModuleList([nn.Sequential(nn.Linear(d, h), nn.Tanh(), nn.Linear(h, h), nn.Tanh()) for _ in range(experts)])
 39        self.router = nn.Linear(d, experts)
 40        self.head = nn.Sequential(nn.Linear(h, h), nn.Tanh(), nn.Linear(h, 1))
 41        self.temperature, self.idea, self.steps, self.eta = temperature, idea, steps, eta
 42        self.last_load = None
 43        self.last_soft_load = None
 44
 45    def forward(self, x):
 46        # Sequence-level groups: each sequence is one routing player; token outputs
 47        # are mixed using the same group allocation, preserving multi-token structure.
 48        b, l, d = x.shape
 49        logits = self.router(x).mean(dim=1) / max(self.temperature, 1e-5)
 50        a = torch.exp(logits - logits.detach().amax(dim=-1, keepdim=True))
 51        soft = torch.softmax(logits, dim=-1)
 52        if not self.idea:
 53            alloc = soft
 54        else:
 55            # b_i is a positive capacity prior; equal capacity is standard in this tiny bench.
 56            cap = torch.full((a.shape[-1],), 0.35, device=x.device, dtype=x.dtype)
 57            alloc = torch.full_like(a, 1.0 / a.shape[-1])
 58            for _ in range(self.steps):
 59                D = cap + alloc.sum(dim=0)
 60                g = a * (D.unsqueeze(0) - alloc) / (D.unsqueeze(0).square() + 1e-8)
 61                old_pay = (a * alloc / D.unsqueeze(0)).sum()
 62                step = self.eta
 63                proposal = project_simplex(alloc + step * g)
 64                new_pay = (a * proposal / (cap + proposal.sum(dim=0)).unsqueeze(0)).sum()
 65                # conservative backtracking, while retaining gradients through accepted proposal
 66                for _ in range(5):
 67                    if new_pay.detach() >= old_pay.detach() - 1e-7: break
 68                    step *= 0.5
 69                    proposal = project_simplex(alloc + step * g)
 70                    new_pay = (a * proposal / (cap + proposal.sum(dim=0)).unsqueeze(0)).sum()
 71                alloc = proposal
 72        tok = torch.stack([e(x.reshape(-1, d)).reshape(b, l, -1) for e in self.experts], dim=2)
 73        pooled = (tok * alloc[:, None, :, None]).sum(dim=2).mean(dim=1)
 74        self.last_load = alloc.detach().sum(0).cpu().numpy()
 75        self.last_soft_load = soft.detach().sum(0).cpu().numpy()
 76        return self.head(pooled)
 77
 78
 79def run_one(seed, lr, temperature, idea, epochs=4, return_model=False):
 80    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 81    ds = get_dataset("congestion_moe_sequence", seed, 240, 100)
 82    dev = safe_device()
 83    model = TinyMoE(temperature=temperature, idea=idea).to(dev)
 84    xtr, ytr = torch.tensor(ds["xtr"], dtype=torch.float32, device=dev), torch.tensor(ds["ytr"], dtype=torch.float32, device=dev)
 85    xte, yte = torch.tensor(ds["xte"], dtype=torch.float32, device=dev), torch.tensor(ds["yte"], dtype=torch.float32, device=dev)
 86    opt = torch.optim.Adam(model.parameters(), lr=lr)
 87    model.train()
 88    for ep in range(epochs):
 89        g = torch.Generator(device="cpu"); g.manual_seed(seed * 100 + ep)
 90        ix = torch.randperm(len(xtr), generator=g, device="cpu").to(dev)
 91        for start in range(0, len(ix), 128):
 92            j = ix[start:start+128]
 93            loss = (model(xtr[j]) - ytr[j]).square().mean()
 94            opt.zero_grad(); loss.backward(); opt.step()
 95    model.eval()
 96    with torch.no_grad(): metric = float((model(xte) - yte).square().mean().cpu())
 97    if return_model: return metric, model
 98    del model
 99    if dev == "cuda": torch.cuda.empty_cache()
100    return metric
101
102
103def main():
104    # Union parity: all idea learning rates are also baseline candidates; baseline's
105    # decisive temperature knob is swept as well.
106    lrs = [0.0015, 0.003, 0.006]
107    temps = [0.5, 0.7, 1.0]
108    base_grid = [{"lr": lr, "temperature": t} for lr in lrs for t in temps]
109    idea_grid = [{"lr": lr, "temperature": 0.7} for lr in lrs]
110    def base_factory(cfg): return lambda seed: run_one(seed, cfg["lr"], cfg["temperature"], False)
111    base = sweep_baseline(base_factory, base_grid)
112    # Required three idea settings, evaluated on all eight paired seeds.
113    idea_runs = []
114    for cfg in idea_grid:
115        r = evaluate(lambda s, c=cfg: run_one(s, c["lr"], c["temperature"], True))
116        idea_runs.append({"cfg": cfg, "result": r})
117    best = min(idea_runs, key=lambda q: q["result"]["mean"])
118    idea = best["result"]
119    # Signature is measured from trained systems on the same held-out task inputs.
120    sig_rows = []
121    for s in range(8):
122        bm, bmodel = run_one(s, base["best_cfg"]["lr"], base["best_cfg"]["temperature"], False, return_model=True)
123        im, imodel = run_one(s, best["cfg"]["lr"], best["cfg"]["temperature"], True, return_model=True)
124        sig_rows.append({"seed": s, "baseline_load_cv": float(np.std(bmodel.last_load)/(np.mean(bmodel.last_load)+1e-9)), "idea_load_cv": float(np.std(imodel.last_load)/(np.mean(imodel.last_load)+1e-9)), "baseline_mse": bm, "idea_mse": im})
125    base_c = base["full"]; cmp = __import__("bench").compare_results(base_c, idea)
126    observed_cv_delta = float(np.mean([r["idea_load_cv"]-r["baseline_load_cv"] for r in sig_rows]))
127    signature = {"prediction": "congestion-aware routing reduces expert-load CV while retaining a fixed row budget", "predicted_direction": "negative", "observed_load_cv_delta": observed_cv_delta, "per_seed": sig_rows, "confirmed": bool(observed_cv_delta < 0)}
128    report = make_report("congestion_moe_sequence", "tiny_moe", base, idea, {"idea_sweep": idea_runs, **signature})
129    report["custom_track"] = {"name": META["name"], "file": "custom_moe_track.py", "domain": META["domain"]}
130    report["device"] = safe_device(); report["epochs"] = 4; report["batch"] = 128
131    Path("bench_report.json").write_text(json.dumps(report, indent=2))
132    print(json.dumps(report, indent=2))
133
134if __name__ == "__main__": main()