| """Finite-difference check of the from-scratch autograd. If analytic grads match |
| numeric grads, the MoE+NCA backward is correct and training can be trusted.""" |
| import numpy as np |
| from garyneuron import init_params, forward |
| from data import make_batch |
|
|
| def run(aux=0.0): |
| cfg = dict(S=4, d=6, he=6, K=3, topk=2, steps=3, p_update=0.6, |
| Vin=10, Vout=10, aux=aux) |
| P = init_params(cfg, seed=0) |
| A, B, Y = make_batch(5, cfg["S"], np.random.default_rng(1)) |
| MASK_SEED = 42 |
| def fwd(): |
| return forward(P, A, B, Y, cfg, np.random.default_rng(MASK_SEED), train=True)[0] |
| tot = fwd() |
| for v in P.values(): v.g = None |
| tot.backward() |
| worst = 0.0 |
| for name in ["Wr", "br", "e0.W2", "e1.W1", "emb", "posemb", "Wo", "bo"]: |
| p = P[name]; flat = p.d.reshape(-1); gf = p.g.reshape(-1) |
| idxs = np.random.default_rng(7).integers(0, flat.size, min(6, flat.size)) |
| maxerr = 0.0 |
| for i in idxs: |
| old = float(flat[i]); eps = 1e-3 |
| flat[i] = old + eps; lp = fwd().d |
| flat[i] = old - eps; lm = fwd().d |
| flat[i] = old |
| num = (lp - lm) / (2 * eps); ana = float(gf[i]) |
| err = abs(num - ana) / (abs(num) + abs(ana) + 1e-8) |
| maxerr = max(maxerr, err) |
| worst = max(worst, maxerr) |
| print(f" {name:8s} max-rel-err {maxerr:.2e}") |
| print(f"aux={aux} WORST {worst:.2e} -> {'PASS' if worst < 2e-2 else 'FAIL'}") |
| return worst |
|
|
| if __name__ == "__main__": |
| print("gradcheck (no aux):"); w1 = run(0.0) |
| print("gradcheck (with load-balance aux):"); w2 = run(0.02) |
|
|
