"""Pure-NumPy GPT: forward + hand-written backward, Adam, checkpoint/resume. Param names match gary-4 so the int8 inference engine is a near drop-in. Trained AND served in numpy -- 'numpy, that's it.'""" import numpy as np C = 0.7978845608028654 # sqrt(2/pi) A = 0.044715 def gelu(x): return 0.5 * x * (1.0 + np.tanh(C * (x + A * x**3))) def dgelu(x): u = C * (x + A * x**3) t = np.tanh(u) return 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t*t) * C * (1.0 + 3.0*A*x*x) def ln_fwd(x, w, b, eps=1e-5): mu = x.mean(-1, keepdims=True) xc = x - mu var = (xc*xc).mean(-1, keepdims=True) inv = 1.0/np.sqrt(var+eps) xhat = xc*inv return xhat*w + b, (xhat, inv, w) def ln_bwd(dy, cache): xhat, inv, w = cache N = xhat.shape[-1] dw = (dy*xhat).reshape(-1, N).sum(0) db = dy.reshape(-1, N).sum(0) dxhat = dy*w dx = inv*(dxhat - dxhat.mean(-1,keepdims=True) - xhat*(dxhat*xhat).mean(-1,keepdims=True)) return dx, dw, db def softmax(x): e = np.exp(x - x.max(-1, keepdims=True)) return e/e.sum(-1, keepdims=True) def init_params(V, E, H, L, BLK, seed=0): rng = np.random.default_rng(seed) P = {} P["tok.weight"] = rng.normal(0,0.02,(V,E)).astype(np.float32) P["pos.weight"] = rng.normal(0,0.02,(BLK,E)).astype(np.float32) sc = 0.02/np.sqrt(2*L) for i in range(L): p=f"blocks.{i}." P[p+"ln1.weight"]=np.ones(E,np.float32); P[p+"ln1.bias"]=np.zeros(E,np.float32) P[p+"attn.in_proj_weight"]=rng.normal(0,0.02,(3*E,E)).astype(np.float32) P[p+"attn.in_proj_bias"]=np.zeros(3*E,np.float32) P[p+"attn.out_proj.weight"]=rng.normal(0,sc,(E,E)).astype(np.float32) P[p+"attn.out_proj.bias"]=np.zeros(E,np.float32) P[p+"ln2.weight"]=np.ones(E,np.float32); P[p+"ln2.bias"]=np.zeros(E,np.float32) P[p+"mlp.0.weight"]=rng.normal(0,0.02,(4*E,E)).astype(np.float32) P[p+"mlp.0.bias"]=np.zeros(4*E,np.float32) P[p+"mlp.2.weight"]=rng.normal(0,sc,(E,4*E)).astype(np.float32) P[p+"mlp.2.bias"]=np.zeros(E,np.float32) P["lnf.weight"]=np.ones(E,np.float32); P["lnf.bias"]=np.zeros(E,np.float32) return P def forward(P, X, cfg, targets=None): B,T = X.shape; E,H,L = cfg["E"],cfg["H"],cfg["L"]; hd=E//H mask = np.triu(np.full((T,T), -1e9, np.float32),1) cache={} x = P["tok.weight"][X] + P["pos.weight"][:T] # (B,T,E) cache["X"]=X; cache["T"]=T for i in range(L): p=f"blocks.{i}."; c={} h1,c["ln1"]=ln_fwd(x,P[p+"ln1.weight"],P[p+"ln1.bias"]); c["x_in"]=x qkv = h1 @ P[p+"attn.in_proj_weight"].T + P[p+"attn.in_proj_bias"] # (B,T,3E) c["h1"]=h1 q=qkv[...,:E]; k=qkv[...,E:2*E]; v=qkv[...,2*E:] def sh(z): return z.reshape(B,T,H,hd).transpose(0,2,1,3) # (B,H,T,hd) qh,kh,vh=sh(q),sh(k),sh(v) sc = (qh @ kh.transpose(0,1,3,2))/np.sqrt(hd) + mask # (B,H,T,T) pr = softmax(sc) oh = pr @ vh # (B,H,T,hd) o = oh.transpose(0,2,1,3).reshape(B,T,E) c["qh"],c["kh"],c["vh"],c["pr"]=qh,kh,vh,pr ao = o @ P[p+"attn.out_proj.weight"].T + P[p+"attn.out_proj.bias"] c["o"]=o x = x + ao h2,c["ln2"]=ln_fwd(x,P[p+"ln2.weight"],P[p+"ln2.bias"]); c["x_mid"]=x m0 = h2 @ P[p+"mlp.0.weight"].T + P[p+"mlp.0.bias"] g = gelu(m0) m2 = g @ P[p+"mlp.2.weight"].T + P[p+"mlp.2.bias"] c["h2"],c["m0"],c["g"]=h2,m0,g x = x + m2 cache[i]=c xf,cache["lnf"]=ln_fwd(x,P["lnf.weight"],P["lnf.bias"]); cache["xf"]=xf logits = xf @ P["tok.weight"].T # (B,T,V) if targets is None: return logits, cache V=logits.shape[-1] lg = logits.reshape(-1,V); tg=targets.reshape(-1) m = lg.max(1,keepdims=True) logp = lg - m - np.log(np.exp(lg-m).sum(1,keepdims=True)) loss = -logp[np.arange(len(tg)),tg].mean() cache["probs"]=np.exp(logp); cache["tg"]=tg; cache["Bn"]=len(tg) return loss, cache def backward(P, cfg, cache): E,H,L=cfg["E"],cfg["H"],cfg["L"]; hd=E//H G={k:np.zeros_like(v) for k,v in P.items()} V=P["tok.weight"].shape[0] probs=cache["probs"]; tg=cache["tg"]; Bn=cache["Bn"]; xf=cache["xf"] dlogits=probs.copy(); dlogits[np.arange(Bn),tg]-=1.0; dlogits/=Bn # (Bn,V) Xshape = cache["X"].shape; B,T=Xshape dxf2d = dlogits @ P["tok.weight"] # (Bn,E) G["tok.weight"] += dlogits.T @ xf.reshape(-1,E) # head part dxf = dxf2d.reshape(B,T,E) dx,dw,db = ln_bwd(dxf, cache["lnf"]); G["lnf.weight"]+=dw; G["lnf.bias"]+=db for i in reversed(range(L)): p=f"blocks.{i}."; c=cache[i] # mlp residual: x = x_mid_after_attn + m2 ; dx flows to both dm2 = dx G[p+"mlp.2.bias"]+=dm2.reshape(-1,E).sum(0) G[p+"mlp.2.weight"]+=dm2.reshape(-1,E).T @ c["g"].reshape(-1,4*E) dg = dm2 @ P[p+"mlp.2.weight"] # (B,T,4E) dm0 = dg * dgelu(c["m0"]) G[p+"mlp.0.bias"]+=dm0.reshape(-1,4*E).sum(0) G[p+"mlp.0.weight"]+=dm0.reshape(-1,4*E).T @ c["h2"].reshape(-1,E) dh2 = dm0 @ P[p+"mlp.0.weight"] dxln2,dw,db = ln_bwd(dh2,c["ln2"]); G[p+"ln2.weight"]+=dw; G[p+"ln2.bias"]+=db dx = dx + dxln2 # through residual # attn residual: x = x_in + ao dao = dx G[p+"attn.out_proj.bias"]+=dao.reshape(-1,E).sum(0) G[p+"attn.out_proj.weight"]+=dao.reshape(-1,E).T @ c["o"].reshape(-1,E) do = dao @ P[p+"attn.out_proj.weight"] # (B,T,E) doh = do.reshape(B,T,H,hd).transpose(0,2,1,3) # (B,H,T,hd) pr,qh,kh,vh = c["pr"],c["qh"],c["kh"],c["vh"] dvh = pr.transpose(0,1,3,2) @ doh dpr = doh @ vh.transpose(0,1,3,2) dsc = pr*(dpr - (dpr*pr).sum(-1,keepdims=True)) dsc /= np.sqrt(hd) dqh = dsc @ kh dkh = dsc.transpose(0,1,3,2) @ qh def unsh(z): return z.transpose(0,2,1,3).reshape(B,T,E) dq,dk,dv = unsh(dqh),unsh(dkh),unsh(dvh) dqkv = np.concatenate([dq,dk,dv],axis=-1) # (B,T,3E) G[p+"attn.in_proj_bias"]+=dqkv.reshape(-1,3*E).sum(0) G[p+"attn.in_proj_weight"]+=dqkv.reshape(-1,3*E).T @ c["h1"].reshape(-1,E) dh1 = dqkv @ P[p+"attn.in_proj_weight"] dxln1,dw,db = ln_bwd(dh1,c["ln1"]); G[p+"ln1.weight"]+=dw; G[p+"ln1.bias"]+=db dx = dx + dxln1 # embeddings demb = dx # (B,T,E) np.add.at(G["tok.weight"], cache["X"], demb) G["pos.weight"][:T] += demb.sum(0) return G class Adam: def __init__(self, P, lr=6e-4, b1=0.9, b2=0.95, wd=0.1, eps=1e-8): self.lr,self.b1,self.b2,self.wd,self.eps=lr,b1,b2,wd,eps self.m={k:np.zeros_like(v) for k,v in P.items()} self.v={k:np.zeros_like(v) for k,v in P.items()} self.t=0 def step(self, P, G, lr=None): self.t+=1; lr=lr or self.lr b1,b2=self.b1,self.b2 bc1=1-b1**self.t; bc2=1-b2**self.t for k in P: g=G[k] self.m[k]=b1*self.m[k]+(1-b1)*g self.v[k]=b2*self.v[k]+(1-b2)*(g*g) mh=self.m[k]/bc1; vh=self.v[k]/bc2 upd=mh/(np.sqrt(vh)+self.eps) if k.endswith("weight") and ("ln" not in k): upd=upd+self.wd*P[k] # decoupled weight decay on matmul weights P[k]-=lr*upd