code/fed_real.py

"""FedDPO / DecDPO on DistilGPT-2 (82M) with real SHP preference pairs,
non-IID split across N=5 clients by SHP domain (the paper's heterogeneity).

Design notes that make the theorems actually testable:
 * Each client's batch sequence is fixed by its own seed, so with FULL
   participation (S=N) the whole round is deterministic. Corollary 5.2 predicts
   the partial-participation variance term vanishes at S=N -- here that becomes
   an EXACT zero-variance check, not a noisy approximation.
 * Client sampling uses a separate RNG, so S<N runs differ only in WHICH
   clients are sampled -- isolating the 1/S term from data-order noise.
"""
import json, time, copy, itertools
import numpy as np, torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from dpo_real import batch_logps, dpo_loss, DEV, MODEL, BETA

N_CLIENTS, E_DEF, R_DEF, BS = 5, 2, 40, 4
RES = {}


def build_clients(n_per_domain=90):
    from datasets import load_dataset
    ds = load_dataset("stanfordnlp/SHP", split="train", streaming=True)
    tok = AutoTokenizer.from_pretrained(MODEL); tok.pad_token = tok.eos_token
    doms, want = {}, N_CLIENTS
    for r in itertools.islice(ds, 400000):
        d = r["domain"]
        if d not in doms and len(doms) >= want: continue
        b = doms.setdefault(d, [])
        if len(b) >= n_per_domain: 
            if all(len(v) >= n_per_domain for v in doms.values()) and len(doms) >= want: break
            continue
        a, bb = r["human_ref_A"], r["human_ref_B"]
        w, l = (a, bb) if r["labels"] == 1 else (bb, a)
        p = "Q: " + r["history"][:400] + "\nA:"
        pw = tok(p + " " + w[:400], truncation=True, max_length=128)["input_ids"]
        pl = tok(p + " " + l[:400], truncation=True, max_length=128)["input_ids"]
        npr = len(tok(p, truncation=True, max_length=128)["input_ids"])
        if len(pw) > npr + 2 and len(pl) > npr + 2:
            b.append({"w": pw, "l": pl, "np": npr})
    names = sorted(doms)[:N_CLIENTS]
    return [doms[d] for d in names], names, tok


def flat(m): return torch.cat([p.detach().reshape(-1) for p in m.parameters()])
def setflat(m, v):
    i = 0
    for p in m.parameters():
        n = p.numel(); p.data.copy_(v[i:i+n].view_as(p)); i += n


def local_train(model, ref, data, E, lr, pad, rng):
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    for _ in range(E):
        idx = rng.integers(0, len(data), size=BS)
        loss, _ = dpo_loss(model, ref, [data[i] for i in idx], pad)
        opt.zero_grad(); loss.backward(); opt.step()


def evaluate(model, ref, clients, pad, nb=3):
    tot, acc, k = 0.0, 0.0, 0
    with torch.no_grad():
        for d in clients:
            for j in range(nb):
                l, a = dpo_loss(model, ref, d[j*BS:(j+1)*BS], pad)
                tot += l.item(); acc += a; k += 1
    return tot/k, acc/k


def fed_run(base, ref, clients, tok, S, R=R_DEF, E=E_DEF, lr=2e-5, q_max=0, seed=0):
    m = copy.deepcopy(base).to(DEV)
    th = flat(m).clone()
    samp = np.random.default_rng(10_000 + seed)
    crngs = [np.random.default_rng(777 + i) for i in range(len(clients))]  # fixed per client
    dly = np.random.default_rng(50_000 + seed)
    buf = {}
    for r in range(R):
        sel = samp.choice(len(clients), size=S, replace=False) if S < len(clients) else np.arange(len(clients))
        for i in sel:
            setflat(m, th)
            local_train(m, ref, clients[i], E, lr, tok.pad_token_id, crngs[i])
            d = flat(m) - th
            k = r + (int(dly.integers(0, q_max+1)) if q_max else 0)
            buf.setdefault(k, []).append(d)
        ds = buf.pop(r, [])
        if ds: th = th + torch.stack(ds).mean(0)
    setflat(m, th)
    return m, th


def metropolis(adj):
    n = adj.shape[0]; deg = adj.sum(1); W = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            if i != j and adj[i, j]: W[i, j] = 1.0/(1+max(deg[i], deg[j]))
        W[i, i] = 1 - W[i].sum()
    return W, float(np.sort(np.abs(np.linalg.eigvals(W)))[::-1][1])


def dec_run(base, ref, clients, tok, W, R=R_DEF, E=E_DEF, lr=2e-5):
    n = len(clients)
    m = copy.deepcopy(base).to(DEV)
    TH = torch.stack([flat(m).clone() for _ in range(n)])
    crngs = [np.random.default_rng(777+i) for i in range(n)]
    Wt = torch.tensor(W, dtype=TH.dtype, device=TH.device)
    for r in range(R):
        new = []
        for i in range(n):
            setflat(m, TH[i]); local_train(m, ref, clients[i], E, lr, tok.pad_token_id, crngs[i])
            new.append(flat(m).clone())
        TH = Wt @ torch.stack(new)
    cons = float(torch.norm(TH - TH.mean(0, keepdim=True), dim=1).mean())
    setflat(m, TH.mean(0))
    return m, cons


def main():
    t0 = time.time()
    clients, names, tok = build_clients()
    print("clients:", [(n, len(d)) for n, d in zip(names, clients)], flush=True)
    base = AutoModelForCausalLM.from_pretrained(MODEL)
    ref = AutoModelForCausalLM.from_pretrained(MODEL).to(DEV).eval()
    for p in ref.parameters(): p.requires_grad_(False)
    pad = tok.pad_token_id
    l0, a0 = evaluate(base.to(DEV), ref, clients, pad)
    RES["init"] = {"loss": l0, "acc": a0, "clients": dict(zip(names, [len(d) for d in clients]))}
    print("init loss=%.4f acc=%.3f" % (l0, a0), flush=True)

    # ---- Claim 2: participation sweep, variance of the global model ----
    rows = []
    for S in (1, 2, 3, 5):
        ths, ls = [], []
        for s in range(3):
            m, th = fed_run(base, ref, clients, tok, S, seed=s)
            l, a = evaluate(m, ref, clients, pad); ths.append(th); ls.append(l)
            print("   S=%d seed=%d loss=%.4f acc=%.3f (%.0fs)" % (S, s, l, a, time.time()-t0), flush=True)
        TH = torch.stack(ths)
        var = float((TH - TH.mean(0, keepdim=True)).pow(2).sum(1).mean())
        fpc = (N_CLIENTS-S)/(S*(N_CLIENTS-1))
        rows.append({"S": S, "param_variance": var, "fpc": round(fpc, 5),
                     "ratio": (var/fpc if fpc > 0 else None),
                     "loss_mean": float(np.mean(ls)), "loss_sd": float(np.std(ls))})
        print("  S=%d Var=%.4e fpc=%.4f loss=%.4f+-%.4f" % (S, var, fpc, np.mean(ls), np.std(ls)), flush=True)
    RES["claim2_participation"] = {"rows": rows}
    json.dump(RES, open("fed_real_results.json", "w"), indent=1)

    # ---- Claim 3: staleness ----
    rows = []
    for q in (0, 2, 5):
        ls = []
        for s in range(2):
            m, _ = fed_run(base, ref, clients, tok, N_CLIENTS, q_max=q, seed=s)
            l, a = evaluate(m, ref, clients, pad); ls.append(l)
        rows.append({"q_max": q, "loss_mean": float(np.mean(ls)), "loss_sd": float(np.std(ls))})
        print("  q_max=%d loss=%.5f+-%.5f (%.0fs)" % (q, np.mean(ls), np.std(ls), time.time()-t0), flush=True)
    RES["claim3_staleness"] = {"rows": rows,
        "monotone": all(rows[i+1]["loss_mean"] >= rows[i]["loss_mean"]-1e-9 for i in range(len(rows)-1))}
    json.dump(RES, open("fed_real_results.json", "w"), indent=1)

    # ---- Claim 5: topology ----
    n = N_CLIENTS
    tops = {}
    ring = np.zeros((n, n), int)
    for i in range(n): ring[i, (i+1) % n] = ring[(i+1) % n, i] = 1
    tops["ring"] = ring
    star = np.zeros((n, n), int); star[0, 1:] = star[1:, 0] = 1; tops["star"] = star
    path = np.zeros((n, n), int)
    for i in range(n-1): path[i, i+1] = path[i+1, i] = 1
    tops["path"] = path
    tops["complete"] = np.ones((n, n), int) - np.eye(n, dtype=int)
    rows = []
    for nm, adj in tops.items():
        W, rho = metropolis(adj)
        m, cons = dec_run(base, ref, clients, tok, W)
        l, a = evaluate(m, ref, clients, pad)
        rows.append({"topology": nm, "rho": round(rho, 4),
                     "one_over_1_minus_rho2": round(1/(1-rho**2), 3),
                     "consensus_error": cons, "loss": l, "acc": a})
        print("  %-9s rho=%.4f cons=%.4e loss=%.4f (%.0fs)" % (nm, rho, cons, l, time.time()-t0), flush=True)
    x = np.log([r["one_over_1_minus_rho2"] for r in rows if r["consensus_error"] > 1e-9])
    y = np.log([r["consensus_error"] for r in rows if r["consensus_error"] > 1e-9])
    if len(x) > 2:
        sl, ic = np.polyfit(x, y, 1)
        RES["claim5_topology"] = {"rows": rows, "loglog_slope": round(float(sl), 4),
            "r2": round(float(1-np.var(y-(sl*x+ic))/np.var(y)), 4),
            "rank_corr_positive": bool(np.corrcoef(x, y)[0, 1] > 0)}
    else:
        RES["claim5_topology"] = {"rows": rows}
    RES["wall_clock_s"] = time.time()-t0
    json.dump(RES, open("fed_real_results.json", "w"), indent=1)
    print("DONE %.0fs" % (time.time()-t0), flush=True)


if __name__ == "__main__":
    main()