transformer.py

1from dataclasses import dataclass2import numpy as np3from ops import (cross_entropy_forward, cross_entropy_backward, layer_norm_backward,4                 layer_norm_forward, linear_backward, relu_backward, relu_forward,5                 softmax_forward, softmax_backward, xavier_uniform)67@dataclass(frozen=True)8class ModelConfig:9    vocab_size: int10    context_len: int11    emb_dim: int12    num_heads: int13    ff_dim: int14    num_layers: int15    ln_eps: float = 1e-51617def init_params(cfg: ModelConfig, rng):18    d, F, V = cfg.emb_dim, cfg.ff_dim, cfg.vocab_size1920    p = {21        "tok_emb": xavier_uniform(V, d, rng), # lookup table with one row per vocabulary token22        "pos_emb": xavier_uniform(cfg.context_len, d, rng) # second table with one row per position, weighted average of values23        # input to first block is tok_emb + pos_emb24        # this makes cat at first position and fifth position different in a sentence25    }2627    for l in range(cfg.num_layers):28        b = f"b{l}." # key29        p[b + "ln1_g"], p[b + "ln1_b"] = np.ones(d), np.zeros(d) # layer norm3031        for n in "qkvo":32            p[b + "w" + n] = xavier_uniform(d, d, rng)33            p[b + "b" + n] = np.zeros(d)3435        p[b + "ln2_g"], p[b + "ln2_b"] = np.ones(d), np.zeros(d)36        p[b + "ff1_w"], p[b + "ff1_b"] = xavier_uniform(d, F, rng), np.zeros(F)37        p[b + "ff2_w"], p[b + "ff2_b"] = xavier_uniform(F, d, rng), np.zeros(d)3839    p["lnf_g"], p["lnf_b"] = np.ones(d), np.zeros(d)40    p["head_w"], p["head_b"] = xavier_uniform(d, V, rng), np.zeros(V)4142    return p4344def count_params(params):45    return int(sum(v.size for v in params.values()))4647# X0[t] = E[i_t] + P[t]48# pos_emb has one row for every position the model could ever see (context_len)49# training window has 16 tokens, suppose 3 tokens are given for generation50# tok_emb[ids]  ->  (T, d)        e.g. (3, 32)51# pos_emb       ->  (Tmax, d)     (32, 32)    # shapes don't match52def embed_forward(params, ids):53    return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]5455def embed_backward(dX0, ids, grads):56    np.add.at(grads["tok_emb"], ids, dX0) # grads["tok_emb"][ids[t]] += dX0[t]57    grads["pos_emb"][: len(ids)] += dX05859# (T,d)->(H,T,dh)60def split_heads(M, H):61    T, d = M.shape62    return M.reshape(T, H, d // H).transpose(1, 0, 2)6364# (H,T,dh)->(T,d)65def merge_heads(Mh):66    H, T, dh = Mh.shape67    return Mh.transpose(1, 0, 2).reshape(T, H * dh)6869# takes one sequence of normalized token vectors and returns a new vector built from itself and earlier positions70def attention_forward(Xn, params, p, cfg: ModelConfig):71    T, d = Xn.shape72    H = cfg.num_heads73    dh = d // H # width of one head7475    Q = Xn @ params[p+"wq"] + params[p+"bq"]    # (T,d), QUERY76    K = Xn @ params[p+"wk"] + params[p+"bk"]    # KEY77    V = Xn @ params[p+"wv"] + params[p+"bv"]    # VALUE7879    Qh, Kh, Vh = split_heads(Q, H), split_heads(K, H), split_heads(V, H)8081    S = (Qh @ Kh.transpose(0, 2, 1)) / np.sqrt(dh)82    S = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, S)83    A = softmax_forward(S)84    O = merge_heads(A @ Vh)85    out = O @ params[p + "wo"] + params[p + "bo"]86    return out, dict(Xn=Xn, Qh=Qh, Kh=Kh, Vh=Vh, A=A, O=O)8788# receives the gradient on attention output and returns gradient on its input89# given   dOut = ∂L/∂out         (T,d)   what arrives from the residual branch90# fills   grads[p+"wo","bo","wq","bq","wk","bk","wv","bv"]91# returns dXn  = ∂L/∂Xn          (T,d)   goes on to layer-norm 1's backward92def attention_backward(dOut, c, params, p, cfg: ModelConfig, grads):93    H = cfg.num_heads94    dh = cfg.emb_dim // H95    grads[p + "wo"] += c["O"].T @ dOut96    grads[p + "bo"] += dOut.sum(axis=0)97    dOh = split_heads(dOut @ params[p + "wo"].T, H)98    dVh = c["A"].transpose(0, 2, 1) @ dOh 99    dA = dOh @ c["Vh"].transpose(0, 2, 1)100    dS = softmax_backward(c["A"], dA)101    dQh = dS @ c["Kh"] / np.sqrt(dh)102    dKh = dS.transpose(0, 2, 1) @ c["Qh"] / np.sqrt(dh)103    dQ, dK, dV = merge_heads(dQh), merge_heads(dKh), merge_heads(dVh)104    dXn = np.zeros_like(c["Xn"])105    for n, dM in (("q", dQ), ("k", dK), ("v", dV)):106        grads[p + "w" + n] += c["Xn"].T @ dM107        grads[p + "b" + n] += dM.sum(axis=0)108        dXn += dM @ params[p + "w" + n].T109    return dXn110111# attention moves information between positions112def ffn_forward(Xn, params, p):113    Z = Xn @ params[p + "ff1_w"] + params[p + "ff1_b"]114    A = relu_forward(Z)115116    return A @ params[p + "ff2_w"] + params[p + "ff2_b"], dict(Xn=Xn, Z=Z, A=A)117118def ffn_backward(dOut, c, params, p, grads):119    dA, dW2, db2 = linear_backward(c["A"], params[p + "ff2_w"], dOut)120    grads[p + "ff2_w"] += dW2121    grads[p + "ff2_b"] += db2122    dZ = relu_backward(c["Z"], dA)123    dXn, dW1, db1 = linear_backward(c["Xn"], params[p + "ff1_w"], dZ)124125    grads[p + "ff1_w"] += dW1126    grads[p + "ff1_b"] += db1127128    return dXn129130# one transformer block in pre-LN form131def block_forward(X, params, l: int, cfg: ModelConfig):132    p = f"b{l}."133    Xn, c_ln1 = layer_norm_forward(X, params[p + "ln1_g"], params[p + "ln1_b"], cfg.ln_eps) # make a normalized copy for attention to read134    att, c_att = attention_forward(Xn, params, p, cfg) # each position gathers info from earlier positions135    X1 = X + att136    Xn2, c_ln2 = layer_norm_forward(X1, params[p + "ln2_g"], params[p + "ln2_b"], cfg.ln_eps) # normalized copy again, for the feed-forward137    ff, c_ff = ffn_forward(Xn2, params, p) # ach position thinks on its own138139    return X1 + ff, dict(ln1=c_ln1, att=c_att, ln2=c_ln2, ff=c_ff)140141def block_backward(dY, c, params, l, cfg: ModelConfig, grads):142    p = f"b{l}."143    dXn2 = ffn_backward(dY, c["ff"], params, p, grads)144    dX1_branch, dg, db = layer_norm_backward(dXn2, c["ln2"], params[p + "ln2_g"])145    grads[p + "ln2_g"] += dg146    grads[p + "ln2_b"] += db147    dX1 = dY + dX1_branch148    dXn = attention_backward(dX1, c["att"], params, p, cfg, grads)149    dX_branch, dg, db = layer_norm_backward(dXn, c["ln1"], params[p + "ln1_g"])150    grads[p + "ln1_g"] += dg151    grads[p + "ln1_b"] += db152    return dX1 + dX_branch153154# whole model in the forward direction155def forward(params, ids, cfg: ModelConfig):156    assert len(ids) <= cfg.context_len, "sequence longer than the position table"157158    X = embed_forward(params, ids)159    block_caches = []160161    for l in range(cfg.num_layers):162        X, c = block_forward(X, params, l, cfg)163        block_caches.append(c)164165    Xf, c_lnf = layer_norm_forward(X, params["lnf_g"], params["lnf_b"], cfg.ln_eps)166    logits = Xf @ params["head_w"] + params["head_b"]167168    return logits, dict(blocks=block_caches, lnf=c_lnf, Xf=Xf)169170# runs forward, computes the loss, then the whole backward pass171def loss_and_grads(params, ids, targets, cfg: ModelConfig):172    logits, cache = forward(params, ids, cfg)173    loss, probs = cross_entropy_forward(logits, targets)174    grads = {k: np.zeros_like(v) for k, v in params.items()}175    dZ = cross_entropy_backward(probs, targets)176    dXf, dW, db = linear_backward(cache["Xf"], params["head_w"], dZ)177178    grads["head_w"] += dW179    grads["head_b"] += db180181    dX, dg, db = layer_norm_backward(dXf, cache["lnf"], params["lnf_g"])182183    grads["lnf_g"] += dg184    grads["lnf_b"] += db185186    for l in reversed(range(cfg.num_layers)):187        dX = block_backward(dX, cache["blocks"][l], params, l, cfg, grads)188    embed_backward(dX, ids, grads)189190    return loss, grads