Mini GPT Training Loop with Custom Autograd: A Hands-On Portfolio Project

# Install: pip install numpy
import numpy as np

# Vocabulary: characters of a short corpus
corpus = "hello world, this is a minimal gpt training loop."
vocab = sorted(set(corpus))
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for ch, i in stoi.items()}
vocab_size = len(vocab)
print("Vocab size:", vocab_size)   # 30

# --- custom_autograd.py ---
class Tensor:
    """A minimal autograd tensor wrapping a numpy array."""
    def __init__(self, data, requires_grad=False):
        self.data = np.array(data, dtype=np.float64)
        self.requires_grad = requires_grad
        self.grad = None
        self._backward = lambda: None   # placeholder

    def __repr__(self):
        return f"Tensor(data={self.data[:3]}..., grad={self.grad})"

    def __add__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data + other.data, requires_grad=self.requires_grad or other.requires_grad)
        # Record a function that will compute gradients later
        out._backward = lambda: (
            self._grad_accumulate(other, lambda: out.grad) if self.requires_grad else None,
            other._grad_accumulate(self, lambda: out.grad) if other.requires_grad else None,
        )
        return out

    def _grad_accumulate(self, other, grad_getter):
        g = grad_getter()
        if self.grad is None:
            self.grad = g
        else:
            self.grad = self.grad + g
        if other.grad is None:
            other.grad = g
        else:
            other.grad = other.grad + g

    def backward(self):
        # Execute the graph in reverse order (simple topological sort omitted for brevity)
        self._backward()
        if self.grad is not None:
            self.grad = np.zeros_like(self.data)  # reset after first backward

# Example usage later in the model

embed_weight = Tensor(np.random.randn(vocab_size, 16) * 0.01)
pos_embed_weight = Tensor(np.random.randn(32, 16) * 0.01)   # max seq len = 32
# Simple linear layer weights for logits
W_lin = Tensor(np.random.randn(16, vocab_size) * 0.01)
b_lin = Tensor(np.zeros(vocab_size))

def softmax(x):
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / e_x.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    d_k = K.shape[-1]
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)
    attn_weights = softmax(scores)
    return attn_weights @ V

def transformer_block(x, embed_w, pos_w, n_heads=2):
    # x shape: (batch, seq_len, dim)
    # Linear projections
    q = x @ embed_w.data
    k = x @ embed_w.data
    v = x @ embed_w.data
    # Split into heads (simplified: just add pos embed)
    q = q + pos_w[: x.shape[1], :]
    k = k + pos_w[: x.shape[1], :]
    v = v + pos_w[: x.shape[1], :]
    out = attention(q, k, v)
    # FFN
    w1 = Tensor(np.random.randn(16, 32) * 0.01)
    w2 = Tensor(np.random.randn(32, 16) * 0.01)
    h = out @ w1.data
    h = np.maximum(0, h)  # ReLU
    out = h @ w2.data
    return out

def train_step(inputs, targets):
    # inputs, targets: int arrays of token ids, shape (batch, seq_len)
    batch, seq_len = inputs.shape
    # Embed inputs
    x = np.zeros((batch, seq_len, 16), dtype=np.float64)
    for t in range(seq_len):
        for b in range(batch):
            x[b, t, :] = embed_weight.data[inputs[b, t]]

    # Add positional embeddings
    for t in range(seq_len):
        x[:, t, :] += pos_embed_weight.data[t]

    # Pass through transformer
    x = transformer_block(x, embed_weight, pos_embed_weight)

    # Final linear projection to logits
    logits = x @ W_lin.data + b_lin.data   # (batch, seq_len, vocab_size)

    # Cross‑entropy loss (ignore position weighting for brevity)
    # Convert logits to probabilities per token
    probs = np.exp(logits) / np.exp(logits).sum(axis=-1, keepdims=True)
    # Gather probabilities of target tokens
    batch_idx = np.arange(batch)[:, None]
    loss = -np.mean(np.log(probs[batch_idx, np.arange(seq_len), targets] + 1e-10))

    # Backward: we need gradients of loss w.r.t. all trainable tensors
    # Simplified: compute gradient of loss w.r.t. logits, then backprop through
    # the linear layer, transformer, and embeddings using the custom tape.
    # Here we just call a hand‑rolled backward that accumulates grads.
    d_logits = probs.copy()
    d_logits[batch_idx, np.arange(seq_len), targets] -= 1
    d_logits /= batch * seq_len

    # Gradient w.r.t. linear weights
    d_W = np.zeros_like(W_lin.data)
    d_b = np.zeros_like(b_lin.data)
    for t in range(seq_len):
        d_W += x[:, t, :].reshape(-1, 1).T @ d_logits[:, t, :].reshape(1, -1)
        d_b += d_logits[:, t, :].sum(axis=0)
    W_lin.grad = d_W
    b_lin.grad = d_b

    # Backprop through linear: gradient flows to x
    d_x = d_logits @ W_lin.data.T   # (batch, seq_len, 16)

    # Remove positional contribution (simple version)
    for t in range(seq_len):
        d_x[:, t, :] -= pos_embed_weight.data[t]  # placeholder; real tape would track

    # Now backprop through transformer (omitted for brevity – you would call
    # each sub‑component’s .backward() that eventually reaches embed_weight,
    # pos_embed_weight, etc.)
    # For this demo we’ll just print the loss.
    return loss

# Prepare tiny dataset from corpus
def build_dataset(corpus, stoi, block_size=8):
    ids = [stoi[c] for c in corpus]
    # sliding windows
    X, Y = [], []
    for i in range(len(ids) - block_size):
        X.append(ids[i:i+block_size])
        Y.append(ids[i+1:i+block_size+1])
    return np.array(X), np.array(Y)

X, Y = build_dataset(corpus, stoi, block_size=8)
dataset = list(zip(X, Y))

# Hyper‑params
lr = 0.01
epochs = 150

for epoch in range(epochs):
    total_loss = 0.0
    for xb, yb in dataset:
        loss = train_step(xb, yb)
        total_loss += loss.item() if hasattr(loss, 'item') else loss
        # Simple SGD update using custom grads
        W_lin.data -= lr * W_lin.grad
        b_lin.data -= lr * b_lin.grad
        # Reset grads (our engine sets them to zero after backward)
        W_lin.grad = None
        b_lin.grad = None
    print(f"Epoch {epoch+1}/{epochs} | Avg Loss: {total_loss/len(dataset):.4f}")