This is Part II of porting Andrej Karpathy's microgpt to the data-parallel language Futhark. Part I covered the forward pass. Now: training and inference.
This part of the port is more involved than Part I, and doesn't map as cleanly 1-to-1, so I won't blog every line of code like I did in Part I, just excerpts. The full code (so far) is on GitHub. Examples of things left out of this post include: some code to do random number generation in Futhark (it's not built in), and some data layout fiddling needed at Futhark external entry points. I note this at the top mostly because if you want to make a fair comparison of how concise the final result is compared to the Python original, you probably do need to take into account the extra bookkeeping that I partly gloss over.
The first thing is to initialize the model. Ok, this is also basically bookkeeping, but it's at least parallelized.
Python:
matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}
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)
Futhark:
def mk_params (v: i64) (seed: i32) : params [v] =
let std = 0.08f32
let o0 = 0i64
let o1 = o0 + v * n_embd
let o2 = o1 + block_size * n_embd
let o3 = o2 + v * n_embd
let o4 = o3 + n_layer * n_embd * n_embd
let o5 = o4 + n_layer * n_embd * n_embd
let o6 = o5 + n_layer * n_embd * n_embd
let o7 = o6 + n_layer * n_embd * n_embd
let o8 = o7 + n_layer * (4 * n_embd) * n_embd
let g s o i j d = (rng_gauss (rng_for_index s (o + i*d + j)) std).1
in {
wte = tabulate_2d v n_embd
(\i j -> g seed o0 i j n_embd),
wpe = tabulate_2d block_size n_embd
(\i j -> g seed o1 i j n_embd),
lm_head = tabulate_2d v n_embd
(\i j -> g seed o2 i j n_embd),
attn_wq = tabulate_3d n_layer n_embd n_embd
(\l i j -> g seed (o3 + l*n_embd*n_embd) i j n_embd),
attn_wk = tabulate_3d n_layer n_embd n_embd
(\l i j -> g seed (o4 + l*n_embd*n_embd) i j n_embd),
attn_wv = tabulate_3d n_layer n_embd n_embd
(\l i j -> g seed (o5 + l*n_embd*n_embd) i j n_embd),
attn_wo = tabulate_3d n_layer n_embd n_embd
(\l i j -> g seed (o6 + l*n_embd*n_embd) i j n_embd),
mlp_fc1 = tabulate_3d n_layer (4 * n_embd) n_embd
(\l i j -> g seed (o7 + l*(4*n_embd)*n_embd) i j n_embd),
mlp_fc2 = tabulate_3d n_layer n_embd (4 * n_embd)
(\l i j -> g seed (o8 + l*n_embd*(4*n_embd)) i j (4*n_embd))
}
The extra code here is because, in Futhark, we conceptually flatten the whole parameter space into a single 1d array for the purposes of parallel random number generation: o0, o1, etc. are offsets into the flattened array where each layer's weights start. Python initializes each layer sequentially, relying on the global mutable state in the random module, but Futhark doesn't have global mutable state.
Another bit of bookkeeping (I promise we will get to real code soon). In Python, all model weights are stored in a flat list of Value objects, which will later make the Adam optimizer be a single loop. In Futhark, I chose (in Part I) to store parameters in a custom params record. Since Futhark lacks generic record mapping, we need two custom map functions to apply elementwise updates.
Futhark:
def params_map2 [v] (f: f32 -> f32 -> f32)
(a: params [v]) (b: params [v]) : params [v] = {
wte = map2 (map2 f) a.wte b.wte,
wpe = map2 (map2 f) a.wpe b.wpe,
lm_head = map2 (map2 f) a.lm_head b.lm_head,
attn_wq = map2 (map2 (map2 f)) a.attn_wq b.attn_wq,
attn_wk = map2 (map2 (map2 f)) a.attn_wk b.attn_wk,
attn_wv = map2 (map2 (map2 f)) a.attn_wv b.attn_wv,
attn_wo = map2 (map2 (map2 f)) a.attn_wo b.attn_wo,
mlp_fc1 = map2 (map2 (map2 f)) a.mlp_fc1 b.mlp_fc1,
mlp_fc2 = map2 (map2 (map2 f)) a.mlp_fc2 b.mlp_fc2
}
def params_map3 [v] (f: f32 -> f32 -> f32 -> f32)
(a: params [v]) (b: params [v]) (c: params [v])
: params [v] = {
wte = map3 (map3 f) a.wte b.wte c.wte,
wpe = map3 (map3 f) a.wpe b.wpe c.wpe,
lm_head = map3 (map3 f) a.lm_head b.lm_head c.lm_head,
attn_wq = map3 (map3 (map3 f)) a.attn_wq b.attn_wq c.attn_wq,
attn_wk = map3 (map3 (map3 f)) a.attn_wk b.attn_wk c.attn_wk,
attn_wv = map3 (map3 (map3 f)) a.attn_wv b.attn_wv c.attn_wv,
attn_wo = map3 (map3 (map3 f)) a.attn_wo b.attn_wo c.attn_wo,
mlp_fc1 = map3 (map3 (map3 f)) a.mlp_fc1 b.mlp_fc1 c.mlp_fc1,
mlp_fc2 = map3 (map3 (map3 f)) a.mlp_fc2 b.mlp_fc2 c.mlp_fc2
}
These could just be inlined (params_map2 is only used twice, and the map3 variant only once), but factoring them out like this will make the actual training code much clearer.
The main reason I even included this in the post (unlike the RNG stuff I left for GitHub) is that it higlights a tricky language design feature. It would be nice if the language did this for us, analogously to what deriving does in Haskell. I believe some parts of this are being tackled in the proposed AUTOMAP language feature, which will eventually be able to figure out that nested map3 (map3 (map3 ... stuff above from the shape of the arrays. However I don't think automatically deriving maps across custom record types is in scope for AUTOMAP.
Ok, now we're on to some real algorithms. And here is where we are least faithful to Karpathy's code, for reasons I argue are justified, but it's a bit of a philosophical question.
Karpathy spends a lot of the total Python linecount implementing a custom Value class with an automatic differentiation (autograd) engine that manually walks computation graphs. During the forward pass, it tracks every operation and builds the computation graph. When backward() is called, it recursively applies the chain rule back through this graph to calculate gradients. (This is also the biggest scaling problem with microgpt: the recursive gradient computation quickly blows the stack if you try to scale the network size even a little bit.)
An excerpted portion of his autograd machinery:
Python:
class Value:
def __init__(self, data, children=(), local_grads=()):
self.data = data # scalar value of this node
self.grad = 0 # derivative of the loss w.r.t. this node
self._children = children # children in the computation graph
self._local_grads = local_grads # local derivative 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))
# ... other operations (log, exp, relu, pow, etc.) ...
def backward(self):
topo = []
visited = set()
def build_topo(v):
if v not in visited:
visited.add(v)
for child in v._children:
build_topo(child)
topo.append(v)
build_topo(self)
self.grad = 1
for v in reversed(topo):
for child, local_grad in zip(v._children, v._local_grads):
child.grad += local_grad * v.grad
In Futhark, instead of porting any of that, we call a single built-in library function, vjp2, which computes a vector-Jacobian product. As it turns out, if you can compute vector-Jacobian products, you have automatic differentiation, so we're done. Explaining the math is beyond the scope of this post, but the explanation for Futhark is roughly the same as this explanation from the jax docs.
Is this cheating? Karpathy's goal was to implement an LLM in a way that's very explicit and self-contained. This port doesn't make the autograd engine's code explicit, which arguably violates that. I think it boils down to a question of what makes sense as a language feature versus application code. Obviously even in his self-contained implementation, Karpathy didn't roll his own object system. I think there's an argument that if you are working in a purely functional programming language with immutable data, Jacobian-vector products and vector-Jacobian products are reasonable things to include as a language-level feature. But I'll admit it's not a standard language feature like an object system is.
Now on to the core of the LLM training, starting with the loss function.
Python:
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
losses = []
for pos_id in range(n):
token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
logits = gpt(token_id, pos_id, keys, values)
probs = softmax(logits)
loss_t = -probs[target_id].log()
losses.append(loss_t)
loss = (1 / n) * sum(losses)
For Futhark, we need to write the loss function as a pure function, so we can later pass it to vjp2 to get gradients. To do that, compute_loss wraps the sequential computation in a function and returns a scalar loss. It is otherwise logically the same, except that instead of dynamically appending to Python lists, we pre-allocate a fixed KV cache.
Futhark:
def compute_loss [v]
(p: params [v])
(tokens: [block_size + 1]i64)
(n_pos: i64)
: f32 =
let keys_init = replicate n_layer
(replicate block_size (replicate n_embd 0f32))
let values_init = replicate n_layer
(replicate block_size (replicate n_embd 0f32))
let (total_loss, _, _) =
loop (loss_acc, keys, values) =
(0f32, keys_init, values_init)
for pos < n_pos do
let token_id = tokens[pos]
let target_id = tokens[pos + 1]
let (logits, keys', values') = gpt p token_id pos keys values
let probs = softmax logits
let loss_t = -(f32.log (f32.max probs[target_id] 1e-10))
in (loss_acc + loss_t, keys', values')
in total_loss / f32.i64 n_pos
Minor note: we pad the token sequence to a fixed block_size + 1 due to how vjp2 is implemented as a static compiler transformation. The closure passed to it cannot have existential sizes (sizes that depend on data), so we need to pad to a statically known maximum size.
Now we have enough to port the training loop. Since this is a big monolithic loop, I'll first pull out just the Adam optimizer update from the middle of it. In Python, we update each parameter individually by iterating through a flat list of Value objects (which already have their .grad attributes computed by loss.backward()).
Python:
lr_t = learning_rate * (1 - step / num_steps)
for i, p in enumerate(params):
m[i] = beta1 * m[i] + (1 - beta1) * p.grad
v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
m_hat = m[i] / (1 - beta1 ** (step + 1))
v_hat = v[i] / (1 - beta2 ** (step + 1))
p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
p.grad = 0
In Futhark, this is where we use our custom params_map2/params_map3 helpers to perform the elementwise momentum and parameter updates in parallel (both because we want parallelization, and because Futhark doesn't allow in-place list mutation anyway). This is also the one and only place we call vjp2 to compute the loss and gradients.
Futhark:
let (loss, grads) =
vjp2 (\p_ -> compute_loss p_ padded n_pos) p 1f32
let lr_t = learning_rate * (1f32 - f32.i64 step / f32.i64 num_steps)
let bc1 = 1f32 - beta1 ** (f32.i64 step + 1)
let bc2 = 1f32 - beta2 ** (f32.i64 step + 1)
let mom' = params_map2 (\m g -> beta1 * m + (1 - beta1) * g) mom grads
let vel' = params_map2 (\v_ g -> beta2 * v_ + (1 - beta2) * g * g) vel grads
let p' = params_map3
(\p_val m_val v_val ->
let m_hat = m_val / bc1
let v_hat = v_val / bc2
in p_val - lr_t * m_hat / (f32.sqrt v_hat + eps_adam)
) p mom' vel'
This part is actually a pretty satisfying almost-line-for-line port.
Now to wrap that Adam update inside the full loop in both languages.
Python:
num_steps = 1000 # number of training steps
for step in range(num_steps):
# Take single document, tokenize it, surround it with BOS special token on both sides
doc = docs[step % len(docs)]
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
n = min(block_size, len(tokens) - 1)
# Forward the token sequence through the model, building up the computation graph all the way to the loss
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
losses = []
for pos_id in range(n):
token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
logits = gpt(token_id, pos_id, keys, values)
probs = softmax(logits)
loss_t = -probs[target_id].log()
losses.append(loss_t)
loss = (1 / n) * sum(losses)
# Backward the loss, calculating the gradients with respect to all model parameters
loss.backward()
# Adam optimizer update: update the model parameters based on the corresponding gradients
lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay
for i, p in enumerate(params):
m[i] = beta1 * m[i] + (1 - beta1) * p.grad
v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
m_hat = m[i] / (1 - beta1 ** (step + 1))
v_hat = v[i] / (1 - beta2 ** (step + 1))
p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
p.grad = 0
print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')
Futhark:
let (p, _mom, _vel, final_loss) =
loop (p, mom, vel, _loss) = (copy p, copy mom, copy vel, 0f32)
for step < num_steps do
let doc_idx = step % n_docs
let doc = docs[doc_idx]
let doc_len = doc_lens[doc_idx]
let padded : [block_size + 1]i64 =
tabulate (block_size + 1) (\i ->
if i == 0 then bos
else if i - 1 < doc_len then doc[i - 1]
else bos
)
let n_pos = i64.min block_size (doc_len + 1)
-- Compute loss and gradients using automatic differentiation (VJP)
let (loss, grads) =
vjp2 (\p_ -> compute_loss p_ padded n_pos) p 1f32
-- Adam updates
let lr_t = learning_rate * (1f32 - f32.i64 step / f32.i64 num_steps)
let bc1 = 1f32 - beta1 ** (f32.i64 step + 1)
let bc2 = 1f32 - beta2 ** (f32.i64 step + 1)
let mom' = params_map2 (\m g -> beta1 * m + (1 - beta1) * g) mom grads
let vel' = params_map2 (\v_ g -> beta2 * v_ + (1 - beta2) * g * g) vel grads
let p' = params_map3
(\p_val m_val v_val ->
let m_hat = m_val / bc1
let v_hat = v_val / bc2
in p_val - lr_t * m_hat / (f32.sqrt v_hat + eps_adam)
) p mom' vel'
in (p', mom', vel', loss)
Both loops are inherently sequential. Futhark has to do a bit more work to satisfy the type system: padding to fixed array sizes and handling uniqueness types (the copy stuff). On the other hand, the Python version does tokenization inline in the training loop, while in Futhark, we factored that out to a host driver outside of the GPU kernel (not shown here, but subject of the next post), since string processing is not a good GPU fit.
Once the model is trained, we want to sample from it. In Python, generating a token is just a single call to the weighted-random version of random.choices:
Python:
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
In Futhark, we implement weighted random sampling by computing the cumulative distribution function (CDF) of the probabilities using scan, and find the first index where the CDF is greater than a random uniform float (this sounds tricky but is a pretty standard array-language recipe):
Futhark:
def weighted_sample [n] (probs: [n]f32) (rng_state: rng) : (rng, i64) =
let (rng_state, u) = rng_f32 rng_state
let cdf = scan (+) 0f32 probs
let idx = i64.sum (map (\c -> if c <= u then 1i64 else 0i64) cdf)
in (rng_state, idx)
And finally the full generation loops. The Python original generates 20 distinct samples of block_size tokens each. In Futhark we will turn the 20 into a num_samples parameter, but otherwise port fairly directly.
Python:
for sample_idx in range(20):
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
token_id = BOS
sample = []
for pos_id in range(block_size):
logits = gpt(token_id, pos_id, keys, values)
probs = softmax([l / temperature for l in logits])
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
if token_id == BOS:
break
sample.append(uchars[token_id])
print(f"sample {sample_idx+1:2d}: {''.join(sample)}")
Futhark:
tabulate num_samples (\sample_idx ->
let rng_state = rng_for_index infer_seed sample_idx
-- Fresh KV caches per sample
let keys_init = replicate n_layer
(replicate block_size (replicate n_embd 0f32))
let values_init = replicate n_layer
(replicate block_size (replicate n_embd 0f32))
-- Autoregressive generation
let (_rng, _tok, _keys, _vals, result) =
loop (rng_state, token_id, keys, vals, tokens) =
(rng_state, bos, keys_init, values_init,
replicate block_size (-1i64))
for pos_id < block_size do
if token_id == bos && pos_id > 0
then (rng_state, token_id, keys, vals, tokens)
else
let (logits, keys', vals') =
gpt p token_id pos_id keys vals
let probs = softmax (map (/ temperature) logits)
let (rng_state', next_token) =
weighted_sample probs rng_state
let tokens =
if next_token == bos then tokens
else tokens with [pos_id] = next_token
in (rng_state', next_token, keys', vals', tokens)
in result
)
Note that there are two loops, and one is still sequential while the other is now parallelized. Autoregressive generation of a single sequence is still sequential (loop in Futhark) because generated tokens in a sequence inherently depend on the previously generated tokens. However we can parallelize generating the num_samples separate sequences.
A nice little language feature to highlight here: tokens with [pos_id] = next_token is a functional in-place update.
The code so far is on GitHub, and is conceptually basically a complete port (the RNG code, miscellaneous boilerplate, and entry points I skipped discussing in this post are there as well). However it is not yet ready to run, because Futhark programs just compile to GPU kernels, and need an external driver program in some general-purpose language like Python. I have one for testing, but it's a bit of a mess so I haven't yet committed it to the repository. Part III of this series will discuss the driver program and include some benchmarks of how this ported microgpt scales.