| |
| """ |
| DGPT — from-scratch decoder-only Transformer (inference-only reimplementation). |
| |
| This is a forward-pass-only extraction of the architecture defined in the |
| original training notebook (NanoScratchGPT / DGPT). The original notebook |
| also hand-implements the backward pass and an AdamW optimizer for training; |
| those are intentionally NOT reproduced here since this file is for loading |
| `model.npz` and generating text only. |
| |
| Locked architecture (see configs/v1-base.json): |
| vocab_size = 6000 |
| block_size = 256 |
| d_model = 384 |
| n_layer = 6 |
| n_head = 6 |
| head_dim = 64 |
| d_ff = 1536 |
| activation = GELU (tanh approximation, GPT-2 style) |
| norm = Pre-LN |
| positions = learned |
| lm_head = tied to token embedding, no output bias |
| |
| No PyTorch, no Hugging Face Transformers. NumPy only. |
| """ |
|
|
| import numpy as np |
|
|
|
|
| def stable_softmax(x, axis=-1): |
| x_max = np.max(x, axis=axis, keepdims=True) |
| e = np.exp(x - x_max) |
| return e / np.sum(e, axis=axis, keepdims=True) |
|
|
|
|
| def gelu(x): |
| c = np.float32((2.0 / np.pi) ** 0.5) |
| inner = c * (x + 0.044715 * x ** 3) |
| return 0.5 * x * (1.0 + np.tanh(inner)) |
|
|
|
|
| class Linear: |
| def __init__(self, W, b=None): |
| self.W = W |
| self.b = b |
|
|
| def __call__(self, x): |
| out = x @ self.W |
| if self.b is not None: |
| out = out + self.b |
| return out |
|
|
|
|
| class LayerNorm: |
| def __init__(self, gamma, beta, eps=1e-5): |
| self.gamma = gamma |
| self.beta = beta |
| self.eps = eps |
|
|
| def __call__(self, x): |
| mu = x.mean(axis=-1, keepdims=True) |
| var = x.var(axis=-1, keepdims=True) |
| x_hat = (x - mu) / np.sqrt(var + self.eps) |
| return self.gamma * x_hat + self.beta |
|
|
|
|
| class CausalSelfAttention: |
| def __init__(self, params, prefix, n_head): |
| self.Wq = Linear(params[f"{prefix}.Wq.W"], params[f"{prefix}.Wq.b"]) |
| self.Wk = Linear(params[f"{prefix}.Wk.W"], params[f"{prefix}.Wk.b"]) |
| self.Wv = Linear(params[f"{prefix}.Wv.W"], params[f"{prefix}.Wv.b"]) |
| self.Wo = Linear(params[f"{prefix}.Wo.W"], params[f"{prefix}.Wo.b"]) |
| self.n_head = n_head |
|
|
| def __call__(self, x): |
| B, T, C = x.shape |
| H = self.n_head |
| hd = C // H |
|
|
| Q, K, V = self.Wq(x), self.Wk(x), self.Wv(x) |
|
|
| def split_heads(t): |
| return t.reshape(B, T, H, hd).transpose(0, 2, 1, 3) |
|
|
| Qh, Kh, Vh = split_heads(Q), split_heads(K), split_heads(V) |
|
|
| scale = np.float32(1.0 / (hd ** 0.5)) |
| scores = np.matmul(Qh, Kh.transpose(0, 1, 3, 2)) * scale |
|
|
| mask = np.triu(np.ones((T, T), dtype=bool), k=1) |
| scores = np.where(mask, np.float32(-1e9), scores) |
|
|
| A = stable_softmax(scores, axis=-1) |
| ctx = np.matmul(A, Vh) |
| ctx_merged = ctx.transpose(0, 2, 1, 3).reshape(B, T, C) |
| return self.Wo(ctx_merged) |
|
|
|
|
| class FeedForward: |
| def __init__(self, params, prefix): |
| self.fc1 = Linear(params[f"{prefix}.fc1.W"], params[f"{prefix}.fc1.b"]) |
| self.fc2 = Linear(params[f"{prefix}.fc2.W"], params[f"{prefix}.fc2.b"]) |
|
|
| def __call__(self, x): |
| return self.fc2(gelu(self.fc1(x))) |
|
|
|
|
| class TransformerBlock: |
| def __init__(self, params, prefix, n_head): |
| self.ln1 = LayerNorm(params[f"{prefix}.ln1.gamma"], params[f"{prefix}.ln1.beta"]) |
| self.attn = CausalSelfAttention(params, f"{prefix}.attn", n_head) |
| self.ln2 = LayerNorm(params[f"{prefix}.ln2.gamma"], params[f"{prefix}.ln2.beta"]) |
| self.ffn = FeedForward(params, f"{prefix}.ffn") |
|
|
| def __call__(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.ffn(self.ln2(x)) |
| return x |
|
|
|
|
| class DGPT: |
| """Inference-only DGPT. Load parameters with `DGPT.from_params(params, config)`.""" |
|
|
| def __init__(self, params, config): |
| self.config = config |
| self.vocab_size = config["vocab_size"] |
| self.block_size = config["block_size"] |
| self.d_model = config["d_model"] |
| self.n_layer = config["n_layer"] |
| self.n_head = config["n_head"] |
|
|
| self.tok_emb_W = params["tok_emb.W"] |
| self.pos_emb_W = params["pos_emb.W"] |
| self.blocks = [ |
| TransformerBlock(params, f"blocks.{i}", self.n_head) |
| for i in range(self.n_layer) |
| ] |
| self.ln_f = LayerNorm(params["ln_f.gamma"], params["ln_f.beta"]) |
|
|
| def num_parameters(self): |
| return sum(v.size for v in self.__dict__.get("_raw_params", {}).values()) |
|
|
| def forward(self, idx): |
| """idx: int array (B, T) with T <= block_size. Returns logits (B, T, vocab_size).""" |
| B, T = idx.shape |
| assert T <= self.block_size, "sequence length exceeds block_size" |
|
|
| tok = self.tok_emb_W[idx] |
| pos = self.pos_emb_W[:T] |
| x = tok + pos[None, :, :] |
|
|
| for blk in self.blocks: |
| x = blk(x) |
|
|
| x = self.ln_f(x) |
| logits = x @ self.tok_emb_W.T |
| return logits |
|
|
