gpt_numpy.py

"""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