Micro-GPT art project by Andrej Karpathy¶
„The most atomic way to train and run inference for a GPT in pure, dependency-free Python.
This file is the complete algorithm.
Everything else is just efficiency.“
@karpathy
We implement a tiny GPT-like model from scratch using only the Python standard library.¶
Educational mindset:
keep everything readable
explain the ‚why‘ (what problem we are solving) more than the ‚how‘ (performance tricks)
do not worry about speed; clarity matters for beginners.
[1]:
import os # os.path.exists: check whether we already have the dataset locally
import math # math.log / math.exp: used for softmax and the negative log-likelihood loss
import random # random: shuffling the dataset and sampling tokens during generation
# Deterministic randomness makes results reproducible (important for learning + debugging).
random.seed(42) # Let there be order among chaos
Let there be a Dataset docs: a list of documents, where each document is a string.
This notebook learns on a list of human names and then generates (hallucinated) new names.
Why do we need a dataset?
A neural network doesn’t come with knowledge. It learns patterns by repeatedly seeing examples.
[2]:
# We use a small plain-text dataset. If it isn't present yet, download it.
if not os.path.exists('input.txt'):
import urllib.request
# The file format is simple: one name per line.
names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
urllib.request.urlretrieve(names_url, 'input.txt')
# Read the dataset file into memory.
# - `line.strip()` removes whitespace/newlines around the name
# - `if line.strip()` skips empty lines (defensive programming)
docs = [line.strip() for line in open('input.txt') if line.strip()]
# Shuffle the list so training sees a mixed stream of examples.
random.shuffle(docs)
print(f"num docs: {len(docs)}")
# print(docs[:10]) # removed: extra output not needed for the lesson
num docs: 32033
[3]:
# Let there be a Tokenizer to translate strings to sequences of integers ("tokens") and back
#
# Tokenizer idea (character-level):
# - collect all unique characters that appear in the dataset
# - assign each character an integer id (a 'token'), happens implicitly with list entrance of uchars, used during tokenization
#
# Why character-level?
# It keeps the problem small: we only need to predict the next character.
uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1
# BOS (Beginning Of Sequence) is a special token that marks the start of a sequence.
# We use it so the model learns a meaningful way to begin generation.
BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token
# vocab_size is the number of possible tokens the model can choose from.
vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS
print(f"vocab size: {vocab_size}")
vocab size: 27
[4]:
# Let there be Autograd to recursively apply the chain rule through a computation graph
class Value:
'''A tiny scalar value that supports automatic differentiation.
Each `Value` stores:
- `data`: the numeric result of this node (a single float)
- `grad`: the gradient of the final loss with respect to `data`
- `_children`: references to input nodes this value depends on
- `_local_grads`: local derivatives used for the chain rule
Why this exists:
Training requires gradients (how to change parameters to reduce loss).
Real ML frameworks build computation graphs + derivatives automatically.
This `Value` is the minimal, educational version of that mechanism.
'''
__slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage
def __init__(self, data, children=(), local_grads=()):
self.data = data # scalar value of this node calculated during forward pass
self.grad = 0 # derivative of the loss w.r.t. this node, calculated in backward pass
self._children = children # children of this node in the computation graph
self._local_grads = local_grads # local derivative of this node w.r.t. its children
def __add__(self, other):
other = other if isinstance(other, Value) else Value(other)
return Value(self.data + other.data, (self, other), (1, 1))
def __mul__(self, other):
other = other if isinstance(other, Value) else Value(other)
return Value(self.data * other.data, (self, other), (other.data, self.data))
def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))
def log(self): return Value(math.log(self.data), (self,), (1/self.data,))
def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))
def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))
def __neg__(self): return self * -1
def __radd__(self, other): return self + other
def __sub__(self, other): return self + (-other)
def __rsub__(self, other): return other + (-self)
def __rmul__(self, other): return self * other
def __truediv__(self, other): return self * other**-1
def __rtruediv__(self, other): return other * self**-1
def backward(self):
'''Compute gradients using reverse-mode automatic differentiation.
Forward pass builds a computation graph (nodes + edges).
Backward pass traverses that graph in reverse topological order and applies
the chain rule to propagate `grad` values from the output (usually the loss)
back to all parameters.
'''
topo = [] # nodes in topological order: children come before parents
visited = set() # avoid adding the same node multiple times
def build_topo(v):
# Depth-first search to collect nodes in the order we need for backprop.
if v not in visited:
visited.add(v)
for child in v._children:
build_topo(child)
topo.append(v) # append after children => topological order
build_topo(self)
# Seed the gradient at the output node.
# If `self` is the loss, then d(loss)/d(loss) = 1.
self.grad = 1
# Propagate gradients backwards.
for v in reversed(topo):
# Chain rule: each child's grad accumulates the local derivative times
# the gradient of the current node.
for child, local_grad in zip(v._children, v._local_grads):
child.grad += local_grad * v.grad
[5]:
# Initialize the parameters, to store the knowledge of the model.
#
# In a transformer/GPT, the parameters are the weight matrices used in:
# - embeddings (turn token ids into vectors, plus position information)
# - self-attention (mix information from previous positions)
# - the MLP (a feed-forward network that adds non-linearity)
# Hyperparameters (kept small for a teaching notebook).
n_layer = 1 # depth of the transformer neural network (number of layers)
n_embd = 16 # width of the network (embedding dimension)
block_size = 16 # maximum context length of the attention window (longest name is ~15 characters)
n_head = 4 # number of attention heads
head_dim = n_embd // n_head # derived dimension of each head
# Helper to create a matrix filled with small random numbers.
# We wrap every scalar with `Value(...)` so autograd can compute gradients.
matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
# state_dict collects *all* parameters by name.
# - wte: token embeddings
# - wpe: position embeddings
# - lm_head: maps final hidden state -> logits over the vocabulary
state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}
# Create additional weights for each transformer layer.
for i in range(n_layer):
state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)
state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)
state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)
state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)
state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)
# Flatten params into a single list[Value] so the optimizer can update them.
params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]
print(f"num params: {len(params)}")
num params: 4192
[6]:
# Define the model architecture: a function mapping tokens and parameters to logits over what comes next
# Follow GPT-2, blessed among the GPTs, with minor differences: layernorm -> rmsnorm, no biases, GeLU -> ReLU
def linear(x, w):
'''A tiny fully-connected (linear) layer.
x is a vector of length `nin`.
w is a matrix of shape (nout, nin).
The result is a vector of length `nout`.
'''
return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]
def softmax(logits):
'''Convert unnormalized scores (logits) into probabilities.
Softmax makes the outputs interpretable as:
P(token=i | context)
'''
max_val = max(val.data for val in logits)
exps = [(val - max_val).exp() for val in logits]
total = sum(exps)
return [e / total for e in exps]
def rmsnorm(x):
'''Root Mean Square normalization.
Neural networks are easier to train when activations have a reasonable scale.
RMSNorm is a simpler alternative to LayerNorm.
'''
ms = sum(xi * xi for xi in x) / len(x)
scale = (ms + 1e-5) ** -0.5
return [xi * scale for xi in x]
def gpt(token_id, pos_id, keys, values):
'''Run the mini-GPT for a single (token, position).
`keys` and `values` are caches that grow as we move from left to right.
That means attention is naturally causal here (no future tokens are in the cache).
'''
tok_emb = state_dict['wte'][token_id] # token embedding
pos_emb = state_dict['wpe'][pos_id] # position embedding
x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding
x = rmsnorm(x) # note: not redundant due to backward pass via the residual connection
for li in range(n_layer):
# 1) Multi-head Attention block
x_residual = x
x = rmsnorm(x)
q = linear(x, state_dict[f'layer{li}.attn_wq'])
k = linear(x, state_dict[f'layer{li}.attn_wk'])
v = linear(x, state_dict[f'layer{li}.attn_wv'])
keys[li].append(k)
values[li].append(v)
x_attn = []
for h in range(n_head):
hs = h * head_dim
q_h = q[hs:hs+head_dim]
k_h = [ki[hs:hs+head_dim] for ki in keys[li]]
v_h = [vi[hs:hs+head_dim] for vi in values[li]]
# Attention scores: how well the current query matches each past key.
attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]
# Turn scores into probabilities over past positions.
attn_weights = softmax(attn_logits)
# Weighted sum of value vectors gives the head output.
head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]
x_attn.extend(head_out)
x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])
x = [a + b for a, b in zip(x, x_residual)]
# 2) MLP block
x_residual = x
x = rmsnorm(x)
x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
# ReLU is the non-linearity: without it, the MLP would stay linear.
x = [xi.relu() for xi in x]
x = linear(x, state_dict[f'layer{li}.mlp_fc2'])
x = [a + b for a, b in zip(x, x_residual)]
# Final projection: map hidden state to logits over the vocabulary.
logits = linear(x, state_dict['lm_head'])
return logits
[7]:
# Let there be Adam, the blessed optimizer and its buffers.
#
# Training loop repeats:
# 1) forward pass -> compute loss
# 2) backward pass -> compute gradients
# 3) optimizer step -> update parameters
#
# Adam updates parameters using running averages of gradients.
learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
# First moment estimate (mean of gradients)
m = [0.0] * len(params) # first moment buffer
# Second moment estimate (mean of squared gradients)
v = [0.0] * len(params) # second moment buffer
[8]:
# Training loop: repeat in sequence
#
# Goal: next-character prediction.
# We take one training name, create a sequence of token ids,
# and repeatedly ask the model: 'Given the prefix so far, what comes next?'
num_steps = 1000 # number of training steps
for step in range(num_steps):
# Take a single document (name) and turn it into training tokens.
#
# Why tokenize?
# Machine learning code usually works on numbers.
# Here: each character is mapped to an integer token id.
#
# Why BOS at both ends?
# - First BOS marks 'start of sequence'.
# - Last BOS provides an explicit end marker.
doc = docs[step % len(docs)]
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS] # here, the element index is used as integer id (see above)
# We can only use a limited amount of context.
# `block_size` is the maximum sequence length the model will consider.
#
# We subtract 1 because for each position we need a *next* token as target.
n = min(block_size, len(tokens) - 1)
# Forward pass:
# We iterate over positions and, for each one, do:
# token -> model logits -> probabilities -> one-step loss
#
# Why 'build up the computation graph'?
# Our `Value` class records dependencies during the forward pass.
# Later, `loss.backward()` uses that recorded graph to compute gradients via
# the chain rule.
# Reset the attention caches for this new sample.
# (They will grow as we generate more characters.)
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
losses = []
for pos_id in range(n):
# Current token is the model input.
# The model should predict the next token as the supervised target.
token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
# Forward pass: get unnormalized scores (logits) for the next token.
logits = gpt(token_id, pos_id, keys, values)
probs = softmax(logits)
# Loss for this position (cross entropy / negative log likelihood).
# We take -log(P(correct next token)).
# If the correct token is unlikely, the loss becomes large.
loss_t = -probs[target_id].log()
losses.append(loss_t)
# Average loss across all positions in this example.
# The average makes the learning signal scale nicely with sequence length.
loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.
# Backward pass:
# `loss.backward()` fills `.grad` for every parameter in `params`.
#
# Remember: our `Value` objects stored the computation graph during the forward pass.
# Backward uses that graph to apply the chain rule.
loss.backward()
# Optimizer step (Adam): update parameters based on gradients.
#
# Adam uses moving averages (moments) to adapt the update size.
# It also helps prevent training from being overly sensitive to noise.
#
# We additionally decay the learning rate so updates become smaller over time.
lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay
for i, p in enumerate(params):
# Update first moment estimate (mean gradient).
# This tracks the 'direction' we should move parameters in.
m[i] = beta1 * m[i] + (1 - beta1) * p.grad
# Update second moment estimate (mean squared gradient).
# This tells us how noisy/variable the gradient is.
v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
# Bias correction for the first moment.
# At the start, the moving average is biased towards 0.
m_hat = m[i] / (1 - beta1 ** (step + 1))
# Bias correction for the second moment.
# This also compensates for initialization at 0.
v_hat = v[i] / (1 - beta2 ** (step + 1))
# Final Adam update.
# - divide by sqrt(v_hat): adapts step size per parameter
# - eps_adam avoids division by zero
p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
# Clear gradients so they don't accumulate across steps.
p.grad = 0
print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')
step 1000 / 1000 | loss 2.6497
[9]:
# Inference: generate new text from the trained model.
#
# During training we always know the correct next character.
# During inference we *don't* know it, so we sample from the model's predicted probabilities.
# Temperature controls how random the sampling is.
# - temperature < 1 makes the model more confident (less random)
# - temperature > 1 makes it more random (more diverse)
temperature = 0.5 # in (0, 1], control the "creativity" of generated text, low to high
print("\n--- inference (new, hallucinated names) ---")
for sample_idx in range(20):
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
# Start generation by forcing the first token to be BOS.
token_id = BOS
# We'll collect generated characters until we hit BOS or the context is full.
sample = []
for pos_id in range(block_size):
logits = gpt(token_id, pos_id, keys, values)
# Convert logits -> probabilities.
# We divide by `temperature` to control randomness.
probs = softmax([l / temperature for l in logits])
# Sample one token id according to the probability distribution.
# We sample instead of taking argmax so generation can be diverse.
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
# BOS acts like an end-of-sequence marker.
# If we sample BOS, we stop generating further characters.
if token_id == BOS:
break
# Convert token id back to a character and add it to our output.
sample.append(uchars[token_id])
print(f"sample {sample_idx+1:2d}: {''.join(sample)}")
--- inference (new, hallucinated names) ---
sample 1: kamon
sample 2: ann
sample 3: karai
sample 4: jaire
sample 5: vialan
sample 6: karia
sample 7: yeran
sample 8: anna
sample 9: areli
sample 10: kaina
sample 11: konna
sample 12: keylen
sample 13: liole
sample 14: alerin
sample 15: earan
sample 16: lenne
sample 17: kana
sample 18: lara
sample 19: alela
sample 20: anton