"""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= 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()