Rewriting an LLM in Pure Python
Greetings! Welcome to the latest in our series on the inner workings of a large language model. In our last post we explored the anatomy of an LLM and wrote our own tools to take one apart, store it, and put it back together. This was part of a larger effort we began in our first post of the series where we explored the workings of a piece of hardware called an FPGA. Our broader goal here will be to take an LLM apart and get it running on entirely new hardware both so that we can enrich our understanding of how the LLM works and so that we can benchmark it and see if an FPGA has the potential to run our network more efficiently.
In today's post we're going to go one level deeper and actually write some code ourselves that runs our LLM. The idea is to write a pure python implementation (no ML libraries) of our neural network so that we can see and feel, line by line, what our network is doing when it does its work. We'll start with pure python so that we can see how, in principle, the structure of a neural network is actually quite simple and in future posts we'll move on to GPU accelerated and finally FPGA accelerated implementations to understand how those pieces of hardware can improve performance.
Tiny Stories Revisited
Like in previous posts we'll look at TinyStories-33M as a good place to start. According to its original research paper it's a modest model that can still produce coherent English because it has been trained on short stories. The principles we get from dissecting this model should translate well to understanding much larger models like GPT-x.
To review, as a basic sanity check we'll "run" the model by making the following script
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL = "roneneldan/TinyStories-33M"
model = AutoModelForCausalLM.from_pretrained(MODEL)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
inputs = tokenizer("Once upon a time", return_tensors="pt")
print(inputs)
outputs = model.generate(**inputs, max_new_tokens=50)
print(outputs[0])
print(tokenizer.decode(outputs[0]))
and running it to produce this text
$ python llm_run_tiny_llm.py
Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red apple on the ground. She picked it up and took a bite. It was so juicy and delicious!So this gives us a basic black box look at what our model is. We give it some seed text, our "Once upon a time", and using the model our program is able to predict what would logically follow that text based on what the model has been trained on. As before, we can go one step further and get an overview of the parts of our model that allow us to do this.
from transformers import AutoModelForCausalLM
MODEL = "roneneldan/TinyStories-33M"
model = AutoModelForCausalLM.from_pretrained(MODEL)
print(model)
which gives us the broad outlines of our model we studied in depth in our previous post.
$ python llm_summarize_model.py
GPTNeoForCausalLM(
(transformer): GPTNeoModel(
(wte): Embedding(50257, 768)
(wpe): Embedding(2048, 768)
(drop): Dropout(p=0.0, inplace=False)
(h): ModuleList(
(0-3): 4 x GPTNeoBlock(
(ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(attn): GPTNeoAttention(
(attention): GPTNeoSelfAttention(
(attn_dropout): Dropout(p=0.0, inplace=False)
(resid_dropout): Dropout(p=0.0, inplace=False)
(k_proj): Linear(in_features=768, out_features=768, bias=False)
(v_proj): Linear(in_features=768, out_features=768, bias=False)
(q_proj): Linear(in_features=768, out_features=768, bias=False)
(out_proj): Linear(in_features=768, out_features=768, bias=True)
)
)
(ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(mlp): GPTNeoMLP(
(c_fc): Linear(in_features=768, out_features=3072, bias=True)
(c_proj): Linear(in_features=3072, out_features=768, bias=True)
(act): NewGELUActivation()
(dropout): Dropout(p=0.0, inplace=False)
)
)
)
(ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
)
(lm_head): Linear(in_features=768, out_features=50257, bias=False)
)Self Attention in Depth
The Attention Calculation
In today's post we'll be focusing on the self attention transformers of this model. I decided to stick with the one most interesting component rather than try and cover the whole model. I figure focused learning is better for our ends and the principles should apply all the same.
If you haven't read last post's primer on what this component is, I would start there. You'll get a summary of where this component comes from in the broader approaches to modeling language and get an explanation of the functional behavior of the component is supposed to be. Remember we are aiming for the self attention behavior
Splitting and Merging Heads
When we did our overview last time we imagined your self attention transformer would take a single token from our text and compare it to all of the text leading up to it to tell us how much attention we should pay to each bit of it. Imagine we have the text "The dog ran for many miles and then it". We can take "it", project it onto our query vector to make Q, project each token leading up to "it" onto K and get the resulting attention we should pay to each of these tokens.
As it turns out, a typical model won't just do this process all in one go, but will actually split each of these vectors into multiple "heads", do an attention calculation on each one, and then combine the results back.
Some parts of our token vectors are really only relevant for some things we want to pay attention to. Others are relevant for other things we want to pay attention to. Remember that matrix calcualtions in general grow quadratically (\(O(n^2)\)) so we get performance gains by breaking up the same matrix into smaller ones and as it turns out the nature of language lets us get away with this. Going back to our example:
The dog ran for many miles and then it
Imagine we have multiple heads that have each been trained to pay attention to different tokens for different reasons.
One head for instance might help enforce subject verb agreement. It might tell you to pay attention to the token that best ensures you have subject-verb agreement. In this case it would give a high attention rating for "it -> dog" because that's the only thing you need to ensure this part is done right. But we can take it one step further. Imagine our tokenized words (remember these are vectors) has all sorts of dimensions for things like emotional tone. These are irrelevant here, and this attention calculation can only work on the dimensions relevant to subject verb agreement, so just abstract things like: part of speech, number, gender, etc. Other heads could help you pay attention to things like the condition of the subject "it -> ran for many miles" or the expected next part of speech "it -> and then". Putting these all together our model can conclude a high likelihood that we want a verb, agrees with a singular neuter subject, and is an appropriate action for something that ran for many miles (so e.g. "rested"). Each of these different attention sub-matrices can be trained differently and adapt over time to serve different roles.
Setting up the Challenge
So my goal in this was to look at the existing code from the transformers package and reproduce it in pure python. The package is fully open source and its specific components for doing self attention are available here. Since these all use torch in order to get hardware assisted performance gains the challenge was to take this package and walk through it step by step to understand the calculations and reproduce them in pure python. It turned out to not be too complicated so I'll share the full implementation here and walk through it step by step to share what I found about how it works.
The Implementation
Starting From the End
In order to ensure I was in fact producing a component that matched the transformers package component exactly I wrote the code as a sort of iterative unit test. This allowed me to do the self attention calculations one part at a time, compare the results to the same calculation done from the standard package and ensure I got the same results. The end result looked like this. The key test to pay attention to is the last one test_self_attention_component_calculations where we go step by step and compare the results of both the serial/vanilla self attention calculation and the transformers/torch self attention calculation side-by-side.
import sys
import pytest
import torch
import llm_serialllm_neo_attn_block as block
from transformers import AutoModelForCausalLM
class TestTensorManipulator:
def test_reshape_tensor(self):
t = torch.tensor(range(16))
true_reshaped = t.view([2, 4, 2]).detach().cpu().tolist()
op = block.TensorManipulator()
reshaped = op.reshape_tensor(list(range(16)), [2, 4, 2])
assert reshaped == true_reshaped
true_reshaped_2 = t.view([1, 8, 2]).detach().cpu().tolist()
reshaped_2 = op.reshape_tensor(list(range(16)), [1, 8, 2])
assert reshaped_2 == true_reshaped_2
def test_reshape_big_tensor(self):
t = torch.tensor(range(24576))
true_reshaped = t.view([2, 16, 16, 48]).detach().cpu().tolist()
op = block.TensorManipulator()
reshaped = op.reshape_tensor(list(range(24576)), [2, 16, 16, 48])
assert reshaped == true_reshaped
def test_permute_tensor(self):
t = torch.tensor(range(16))
true_permuted = t.view([2, 4, 2]).permute([1, 0, 2]).detach().cpu().tolist()
op = block.TensorManipulator()
reshaped = op.reshape_tensor(list(range(16)), [2, 4, 2])
permuted = op.permute_tensor(reshaped, [1, 0, 2])
assert true_permuted == permuted
def test_transpose_tensor(self):
t = torch.tensor(range(16))
true_transposed = t.view([2, 4, 2]).transpose(-1, -2).detach().cpu().tolist()
op = block.TensorManipulator()
reshaped = op.reshape_tensor(list(range(16)), [2, 4, 2])
transposed = op.transpose_tensor(reshaped, -1, -2)
assert true_transposed == transposed
def test_matmul_2d(self):
t_base = torch.tensor(range(16))
t1 = t_base.view([2, 8])
t2 = t_base.view([8, 2])
product = torch.matmul(t1, t2).detach().cpu().tolist()
op = block.TensorManipulator()
vanilla_t1 = op.reshape_tensor(list(range(16)), [2, 8])
vanilla_t2 = op.reshape_tensor(list(range(16)), [8, 2])
vanilla_product = op.matmul_2d(vanilla_t1, vanilla_t2)
assert product == vanilla_product
def test_matmul(self):
# First do a basic 2d matrix
t_base = torch.tensor(range(16))
t1 = t_base.view([2, 8])
t2 = t_base.view([8, 2])
product = torch.matmul(t1, t2).detach().cpu().tolist()
op = block.TensorManipulator()
vanilla_t1 = op.reshape_tensor(list(range(16)), [2, 8])
vanilla_t2 = op.reshape_tensor(list(range(16)), [8, 2])
vanilla_product = op.matmul(vanilla_t1, vanilla_t2)
assert product == vanilla_product
# Now try a batched multiply
t_base = torch.tensor(range(128))
t1 = t_base.view([4, 2, 2, 8])
t2 = t_base.view([4, 2, 8, 2])
product = torch.matmul(t1, t2).detach().cpu().tolist()
op = block.TensorManipulator()
vanilla_t1 = op.reshape_tensor(list(range(128)), [4, 2, 2, 8])
vanilla_t2 = op.reshape_tensor(list(range(128)), [4, 2, 8, 2])
vanilla_product = op.matmul(vanilla_t1, vanilla_t2)
assert product == vanilla_product
def test_where(self):
t = torch.tensor(range(16)).view([2, 2, 4])
mask = torch.tensor([True, False] * 8).view([2, 2, 4])
expected = torch.where(mask, t, 0).detach().cpu().tolist()
op = block.TensorManipulator()
t_vanilla = op.reshape_tensor(list(range(16)), [2, 2, 4])
mask_vanilla = op.reshape_tensor([True, False] * 8, [2, 2, 4])
actual = op.where(mask_vanilla, t_vanilla, 0)
assert expected == actual
def test_slice(self):
t = torch.tensor(range(16)).view([2, 2, 4])
expected = t[:,0,0:4:2].detach().cpu().tolist()
op = block.TensorManipulator()
t_vanilla = op.reshape_tensor(list(range(16)), [2, 2, 4])
actual = op.slice_tensor(t_vanilla, (), (0,), (0, 4, 2))
assert expected == actual
class TestSerialNeoBlock:
def test_self_attention_component_calculations(self):
torch.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained(block.MODEL)
neo_block = model.transformer.h[0]
self_attn_block = neo_block.attn.attention
serial_attn_block = block.SerialSelfAttention.from_transformer(self_attn_block)
# Compare basic tensor multiplication on one dimension vector of Q,K,V,Out transformations
for proj_name in ["k_proj", "v_proj", "q_proj", "out_proj"]:
proj = getattr(self_attn_block, proj_name)
serial_proj = getattr(serial_attn_block, proj_name)
x = torch.randn(proj.in_features)
x_vanilla = block.vanilla_vector(x)
y1 = proj(x)
y2 = serial_proj(x_vanilla)
assert block.vectors_match(y1.detach().cpu().tolist(), y2)
# Compare attention transformations from the two transformers
batch_size, seq_len = 2, 16 # arbitrary for testing
attn_input = torch.randn(
batch_size,
seq_len,
self_attn_block.config.hidden_size,
)
vanilla_input = block.vanilla_vector(attn_input)
for proj_name in ["k_proj", "v_proj", "q_proj", "out_proj"]:
proj = getattr(self_attn_block, proj_name)
out = proj(attn_input)
serial_proj = getattr(serial_attn_block, proj_name)
test = serial_proj(vanilla_input)
# Compare basic tensor multiplication on batch vectors of Q,K,V,Out transformations
assert block.vectors_match(block.vanilla_vector(out), test)
# Compare splitting heads
split_out = self_attn_block._split_heads(
out,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
vanilla_split_out = serial_attn_block._split_heads(
test,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
assert block.vectors_match(block.vanilla_vector(split_out), vanilla_split_out)
# After we split heads we do the attention calculation on the split tensors
op = block.TensorManipulator()
# Calculate attention weights
key = self_attn_block.k_proj(attn_input)
split_key = self_attn_block._split_heads(
key,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
key_transposed = split_key.transpose(-1, -2)
vanilla_key = serial_attn_block.k_proj(vanilla_input)
vanilla_split_key = serial_attn_block._split_heads(
vanilla_key,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
vanilla_key_transposed = op.transpose_tensor(vanilla_split_key, -1, -2)
assert block.vectors_match(block.vanilla_vector(key_transposed), vanilla_key_transposed)
query = self_attn_block.q_proj(attn_input)
split_query = self_attn_block._split_heads(
query,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
vanilla_query = serial_attn_block.q_proj(vanilla_input)
vanilla_split_query = serial_attn_block._split_heads(
vanilla_query,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
assert block.vectors_match(block.vanilla_vector(split_query), vanilla_split_query)
value = self_attn_block.v_proj(attn_input)
split_value = self_attn_block._split_heads(
value,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
vanilla_value = serial_attn_block.v_proj(vanilla_input)
vanilla_split_value = serial_attn_block._split_heads(
vanilla_value,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
assert block.vectors_match(block.vanilla_vector(split_value), vanilla_split_value)
attn_weights = torch.matmul(split_query, key_transposed)
vanilla_attn_weights = op.matmul(vanilla_split_query, vanilla_key_transposed)
assert block.vectors_match(block.vanilla_vector(attn_weights), vanilla_attn_weights)
# Causal mask and attention weight calculations
query_length, key_length = split_query.size(-2), split_key.size(-2)
vanilla_query_length = op.tensor_shape(vanilla_split_query)[-2]
vanilla_key_length = op.tensor_shape(vanilla_split_key)[-2]
assert query_length == vanilla_query_length
assert key_length == vanilla_key_length
bias = self_attn_block.bias
causal_mask = bias[:, :, key_length - query_length : key_length, :key_length]
vanilla_causal_mask = op.slice_tensor(serial_attn_block.bias, (), (), (key_length - query_length,key_length),(0, key_length))
assert causal_mask.detach().cpu().tolist() == vanilla_causal_mask
mask_value = torch.tensor(torch.finfo(attn_weights.dtype).min, dtype=attn_weights.dtype, device=attn_weights.device)
# Hard coded to the same minimum float32 value used by torch
vanilla_mask_value = -3.4028234663852886e+38
attn_weights = torch.where(causal_mask, attn_weights, mask_value)
vanilla_attn_weights = op.where(vanilla_causal_mask, vanilla_attn_weights, vanilla_mask_value)
assert block.vectors_match(attn_weights.detach().cpu().tolist(), vanilla_attn_weights)
# softmax and dropout
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
vanilla_attn_weights = op.softmax(vanilla_attn_weights) # our vanilla implementation only supports normalizing over the last dimension so no dim provided
assert block.vectors_match(attn_weights.detach().cpu().tolist(), vanilla_attn_weights)
# Technically the dropout is 0 so this is a no-op but for completeness we do it anyway
# - subsequent vectors_match comparisons only work because dropout is 0 - otherwise
# we'd have non-deterministic zeroing out of the matrix entries which would create
# very mismatched values even between different calls of the same self attention block
attn_weights = self_attn_block.attn_dropout(attn_weights)
vanilla_attn_weights = op.dropout(serial_attn_block.attn_dropout, vanilla_attn_weights)
assert block.vectors_match_with_dropout(attn_weights.detach().cpu().tolist(), vanilla_attn_weights)
# Subsequent assertions only work because our dropout is 0 - otherwise we would have
# non-deterministic output. Make this an explicit prerequisite
assert 0.0 == serial_attn_block.attn_dropout
attn_output = torch.matmul(attn_weights, split_value)
vanilla_attn_output = op.matmul(vanilla_attn_weights, vanilla_split_value)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)
# merge heads
attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
new_shape = attn_output.size()[:-2] + (self_attn_block.num_heads * self_attn_block.head_dim,)
attn_output = attn_output.view(new_shape)
vanilla_attn_output = op.permute_tensor(vanilla_attn_output, [0, 2, 1, 3])
vanilla_new_shape = op.tensor_shape(vanilla_attn_output)[:-2] + (serial_attn_block.num_heads * serial_attn_block.head_dim,)
vanilla_attn_output = op.reshape_tensor(vanilla_attn_output, vanilla_new_shape)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)
# out projection and residual dropout
attn_output = self_attn_block.out_proj(attn_output)
vanilla_attn_output = serial_attn_block.out_proj(vanilla_attn_output)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)
# final test
actual_attn_output = self_attn_block(attn_input)[0]
assert block.vectors_match(actual_attn_output.detach().cpu().tolist(), vanilla_attn_output)
These tests were done against a set of utility functions I wrote here
"""
Iteration 3
Expands the serial linear transformation approach to have a serial
attention block as well
"""
import itertools
import json
import math
import random
import torch
from torch import nn
from torch.nn.modules.linear import Linear
from transformers import AutoModelForCausalLM
from transformers.models.gpt_neo import modeling_gpt_neo as neo
MODEL = "roneneldan/TinyStories-33M"
class SerialNeoBlock(object):
def __init__(self):
pass
@classmethod
def from_transformer(cls, neo_block):
return cls()
class TensorManipulator(object):
def tensor_shape(self, tensor):
if not isinstance(tensor, list):
return ()
return (len(tensor), ) + self.tensor_shape(tensor[0])
def flatten_tensor(self, tensor):
if not isinstance(tensor[0], list):
return tensor
return sum(map(self.flatten_tensor, tensor), start=[])
def reshape_tensor(self, tensor, shape):
flattened_tensor = self.flatten_tensor(tensor)
if len(shape) == 1:
if len(tensor) != shape[0]:
raise ValueError(
"Expected tensor of length {} but is {}".format(
shape[0], len(tensor)))
return tensor
sub_tensors = []
partitions = shape[0]
if len(flattened_tensor) % partitions:
raise ValueError(
"Partitions {} must divide into tensor size {} evenly".format(
partitions, len(tensor)))
partition_size = len(flattened_tensor) // partitions
for i in range(0, len(flattened_tensor), partition_size):
sub_tensors.append(
self.reshape_tensor(flattened_tensor[i:i + partition_size],
shape[1:]))
return sub_tensors
def permute_tensor(self, tensor, perm):
flattened = self.flatten_tensor(tensor)
shape = self.tensor_shape(tensor)
permuted = []
for coord in self.generate_coordinates(shape):
new_coord = self.permute_coordinate(coord, perm)
self.insert_coordinate(permuted, new_coord,
self.get_coordinate(tensor, coord))
return permuted
def generate_coordinates(self, shape):
return itertools.product(*map(range, shape))
def get_coordinate(self, tensor, coordinate):
if len(coordinate) == 1:
return tensor[coordinate[0]]
return self.get_coordinate(tensor[coordinate[0]], coordinate[1:])
def insert_coordinate(self, tensor, coordinate, value):
index = coordinate[0]
if len(coordinate) == 1:
if len(tensor) <= index:
tensor.extend([None] * ((index - len(tensor)) + 1))
tensor[index] = value
return
if len(tensor) <= index:
tensor.extend([[]] * ((index - len(tensor)) + 1))
self.insert_coordinate(tensor[index], coordinate[1:], value)
def get_at_coordinate(self, tensor, coordinate):
if len(coordinate) == 0:
return tensor
return self.get_at_coordinate(tensor[coordinate[0]], coordinate[1:])
def permute_coordinate(self, coordinate, perm):
new_coord = [None] * len(coordinate)
for i, j in enumerate(perm):
new_coord[j] = coordinate[i]
return tuple(new_coord)
def transpose_tensor(self, tensor, dim0, dim1):
normalized_dims = list(range(len(self.tensor_shape(tensor))))
normalized_dim0 = normalized_dims[dim0]
normalized_dim1 = normalized_dims[dim1]
normalized_dims[dim0] = normalized_dim1
normalized_dims[dim1] = normalized_dim0
return self.permute_tensor(tensor, normalized_dims)
def matmul(self, a, b):
a_shape = self.tensor_shape(a)
b_shape = self.tensor_shape(b)
if len(a_shape) < 2 or len(b_shape) < 2:
raise ValueError("Invalid input shapes, must be at least two-dimensional", a_shape, b_shape)
batch_size = a_shape[:-2]
batch_size_b = b_shape[:-2]
if batch_size != batch_size_b:
raise ValueError("Mismatched batch sizes", batch_size, batch_size_b)
if len(batch_size) == 0:
return self.matmul_2d(a, b)
result = []
for coordinate in self.generate_coordinates(batch_size):
self.insert_coordinate(
result,
coordinate,
self.matmul_2d(
self.get_at_coordinate(a, coordinate),
self.get_at_coordinate(b, coordinate),
),
)
return result
def matmul_2d(self, a, b):
a_shape = self.tensor_shape(a)
b_shape = self.tensor_shape(b)
if len(a_shape) != 2:
raise ValueError("Tensor a must by an NxM matrix")
if len(b_shape) != 2:
raise ValueError("Tensor b must by an NxM matrix")
if a_shape[1] != b_shape[0]:
raise ValueError(f"Second dimension of matrix a ({a_shape[1]}) must match first dimension of matrix b {b_shape[0]}")
b_transposed = self.transpose_tensor(b, 0, 1)
return [
[
sum(a_entry * b_entry for a_entry, b_entry in zip(a_row, b_col))
for b_col in b_transposed
]
for a_row in a
]
def where(self, condition, tensor, value):
if isinstance(condition, list) and isinstance(tensor, list):
if len(condition) == 1 and len(tensor) != 1:
# Handle broadcasting (a single condition is applied
# to an entire dimension of the tensor)
condition = condition * len(tensor)
if isinstance(condition, list):
return [self.where(c, t, value) for c, t in zip(condition, tensor)]
if condition:
return tensor
return value
def _single_slice(self, mat, slc):
if len(slc) == 0:
return mat
if len(slc) == 1:
return mat[slc[0]]
if len(slc) == 2:
return mat[slc[0]:slc[1]]
if len(slc) == 3:
return mat[slc[0]:slc[1]:slc[2]]
raise ValueError("Invalid slice", slc)
def slice_tensor(self, mat, *slices):
first, remainder = slices[0], slices[1:]
sliced = self._single_slice(mat, first)
if not remainder:
return sliced
if len(first) == 1:
return self.slice_tensor(sliced, *remainder)
return [self.slice_tensor(elem, *remainder) for elem in sliced]
def softmax(self, t):
if isinstance(t[0], list):
return list(map(self.softmax, t))
exps = list(map(math.exp, t))
total = sum(exps)
return [ex / total for ex in exps]
def _random_bool(self, p):
return random.random() < p
def dropout(self, p, t):
if isinstance(t[0], list):
return [self.dropout(p, elm) for elm in t]
return [0 if self._random_bool(p) else elm for elm in t]
class SerialSelfAttention(object):
def __init__(self, k_proj, v_proj, q_proj, out_proj, num_heads, head_dim, bias, attn_dropout):
self.k_proj = k_proj
self.v_proj = v_proj
self.q_proj = q_proj
self.out_proj = out_proj
self.num_heads = num_heads
self.head_dim = head_dim
self.bias = bias
self.attn_dropout = attn_dropout
def __call__(self, x):
query = self.q_proj(x)
key = self.k_proj(x)
value = self.v_proj(x)
query = self._split_heads(query, self.num_heads, self.head_dim)
def _split_heads(self, tensor, num_heads, attn_head_size):
manipulator = TensorManipulator()
new_shape = manipulator.tensor_shape(tensor)[:-1] + (num_heads,
attn_head_size)
new_tensor = manipulator.reshape_tensor(tensor, new_shape)
return manipulator.permute_tensor(new_tensor, (0, 2, 1, 3))
@classmethod
def from_transformer(cls, neo_self_attention):
return cls(
k_proj=SerialLinear.from_transformer(neo_self_attention.k_proj),
v_proj=SerialLinear.from_transformer(neo_self_attention.v_proj),
q_proj=SerialLinear.from_transformer(neo_self_attention.q_proj),
out_proj=SerialLinear.from_transformer(
neo_self_attention.out_proj),
num_heads=neo_self_attention.num_heads,
head_dim=neo_self_attention.head_dim,
bias=neo_self_attention.bias.detach().cpu().tolist(),
attn_dropout = neo_self_attention.attn_dropout.p,
)
class SerialLinear(object):
def __init__(self, weights, biases):
self.weights = weights
self.biases = biases
def __call__(self, x):
if len(x) == 0:
return x
if isinstance(x[0], list):
return list(map(self, x))
out = [
sum([wi * xi for wi, xi in zip(row, x)]) for row in self.weights
]
if self.biases is None:
return out
return [outi + bi for outi, bi in zip(out, self.biases)]
@classmethod
def from_transformer(cls, neo_linear):
if neo_linear.bias is None:
return cls(vanilla_vector(neo_linear.weight), neo_linear.bias)
return cls(vanilla_vector(neo_linear.weight), vanilla_vector(neo_linear.bias))
def vanilla_vector(torch_vector):
return torch_vector.detach().cpu().tolist()
def vectors_match(y1, y2):
if isinstance(y1, list) != isinstance(y2, list):
return False
if len(y1) != len(y2):
return False
if len(y1) == 0:
return len(y2) == 0
if isinstance(y1[0], list):
return all(vectors_match(l1, l2) for l1, l2 in zip(y1, y2))
return all(
math.isclose(y2i, y1i, abs_tol=1e-5) for y2i, y1i in zip(y2, y1))
def vectors_match_with_dropout(y1, y2):
if isinstance(y1, list) != isinstance(y2, list):
return False
if len(y1) != len(y2):
return False
if len(y1) == 0:
return len(y2) == 0
if isinstance(y1[0], list):
return all(vectors_match(l1, l2) for l1, l2 in zip(y1, y2))
return all(
math.isclose(y2i, y1i, abs_tol=1e-5) or y2i == 0 or y1i == 0 for y2i, y1i in zip(y2, y1))
def main():
pass
if __name__ == "__main__":
main()
So let's go through it.
The Broad Strokes
So from a high level I realized we can divide our attention calculation into three parts. We first split our attention heads, apply the individual attention calcuations, and then merge our heads back. This is more or less how the existing library code divided up the work and it makes sense. In essence this means we want to do three independent attention computations and then combine the results.
Deep Dive 1: Splitting Heads
From a high level, the code for splitting heads was a simple operation. I basically reproduced it in a few lines of my own.
def _split_heads(self, tensor, num_heads, attn_head_size):
manipulator = TensorManipulator()
new_shape = manipulator.tensor_shape(tensor)[:-1] + (num_heads,
attn_head_size)
new_tensor = manipulator.reshape_tensor(tensor, new_shape)
return manipulator.permute_tensor(new_tensor, (0, 2, 1, 3))But it's not the most intuitive code. As noted before, we're basically taking an existing vector (our input token projected onto one of our three projection vectors) and breaking it up into multiple vectors. The entries are all there in the vector already, more or less in the arrangement we want to have, but we reshape our vector to be split up into multiple ones. So from my perspective, I just treated a vector as a list of lists of lists... of numbers. You can flatten these out into a single list and then reshape however you like.
def tensor_shape(self, tensor):
if not isinstance(tensor, list):
return ()
return (len(tensor), ) + self.tensor_shape(tensor[0])
def flatten_tensor(self, tensor):
if not isinstance(tensor[0], list):
return tensor
return sum(map(self.flatten_tensor, tensor), start=[])
def reshape_tensor(self, tensor, shape):
flattened_tensor = self.flatten_tensor(tensor)
if len(shape) == 1:
if len(tensor) != shape[0]:
raise ValueError(
"Expected tensor of length {} but is {}".format(
shape[0], len(tensor)))
return tensor
sub_tensors = []
partitions = shape[0]
if len(flattened_tensor) % partitions:
raise ValueError(
"Partitions {} must divide into tensor size {} evenly".format(
partitions, len(tensor)))
partition_size = len(flattened_tensor) // partitions
for i in range(0, len(flattened_tensor), partition_size):
sub_tensors.append(
self.reshape_tensor(flattened_tensor[i:i + partition_size],
shape[1:]))
return sub_tensors
def permute_tensor(self, tensor, perm):
flattened = self.flatten_tensor(tensor)
shape = self.tensor_shape(tensor)
permuted = []
for coord in self.generate_coordinates(shape):
new_coord = self.permute_coordinate(coord, perm)
self.insert_coordinate(permuted, new_coord,
self.get_coordinate(tensor, coord))
return permutedDeep Dive 2: Attention Calculation
The next step was doing the attention calculation on each of these split heads. So remember this is
Working from the outside in I transposed the key vector and matrix multiplied it with the query vector.
key = self_attn_block.k_proj(attn_input)
split_key = self_attn_block._split_heads(
key,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
key_transposed = split_key.transpose(-1, -2)
vanilla_key = serial_attn_block.k_proj(vanilla_input)
vanilla_split_key = serial_attn_block._split_heads(
vanilla_key,
serial_attn_block.num_heads,
serial_attn_block.head_dim,
)
vanilla_key_transposed = op.transpose_tensor(vanilla_split_key, -1, -2)attn_weights = torch.matmul(split_query, key_transposed)
vanilla_attn_weights = op.matmul(vanilla_split_query, vanilla_key_transposed)I had to remember by college linear algebra and reverse engineer a bit of how torch handles batching of multiplication requests but it ended up simple enough.
def matmul(self, a, b):
a_shape = self.tensor_shape(a)
b_shape = self.tensor_shape(b)
if len(a_shape) < 2 or len(b_shape) < 2:
raise ValueError("Invalid input shapes, must be at least two-dimensional", a_shape, b_shape)
batch_size = a_shape[:-2]
batch_size_b = b_shape[:-2]
if batch_size != batch_size_b:
raise ValueError("Mismatched batch sizes", batch_size, batch_size_b)
if len(batch_size) == 0:
return self.matmul_2d(a, b)
result = []
for coordinate in self.generate_coordinates(batch_size):
self.insert_coordinate(
result,
coordinate,
self.matmul_2d(
self.get_at_coordinate(a, coordinate),
self.get_at_coordinate(b, coordinate),
),
)
return result
def matmul_2d(self, a, b):
a_shape = self.tensor_shape(a)
b_shape = self.tensor_shape(b)
if len(a_shape) != 2:
raise ValueError("Tensor a must by an NxM matrix")
if len(b_shape) != 2:
raise ValueError("Tensor b must by an NxM matrix")
if a_shape[1] != b_shape[0]:
raise ValueError(f"Second dimension of matrix a ({a_shape[1]}) must match first dimension of matrix b {b_shape[0]}")
b_transposed = self.transpose_tensor(b, 0, 1)
return [
[
sum(a_entry * b_entry for a_entry, b_entry in zip(a_row, b_col))
for b_col in b_transposed
]
for a_row in a
]Next we had a bunch of droupout calculations which as best as I can tell these introduced some random zeros into our calculations so that we wouldn't become overly reliant on any particular weights on them and could allow for some flexibility in our inputs/outputs. The configured dropout for our particular LLM was non-existent so we can go ahead and skip this in the summary.
And finally we did the remainder of the attention calculation via softmax and factoring in the value component.
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
vanilla_attn_weights = op.softmax(vanilla_attn_weights) # our vanilla implementation only supports normalizing over the last dimension so no dim provided
assert block.vectors_match(attn_weights.detach().cpu().tolist(), vanilla_attn_weights)
# Technically the dropout is 0 so this is a no-op but for completeness we do it anyway
# - subsequent vectors_match comparisons only work because dropout is 0 - otherwise
# we'd have non-deterministic zeroing out of the matrix entries which would create
# very mismatched values even between different calls of the same self attention block
attn_weights = self_attn_block.attn_dropout(attn_weights)
vanilla_attn_weights = op.dropout(serial_attn_block.attn_dropout, vanilla_attn_weights)
assert block.vectors_match_with_dropout(attn_weights.detach().cpu().tolist(), vanilla_attn_weights)
# Subsequent assertions only work because our dropout is 0 - otherwise we would have
# non-deterministic output. Make this an explicit prerequisite
assert 0.0 == serial_attn_block.attn_dropout
attn_output = torch.matmul(attn_weights, split_value)
vanilla_attn_output = op.matmul(vanilla_attn_weights, vanilla_split_value)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)Deep Dive 3: Merging Heads
The last part of this was merging the heads back in so that we would take our separate calculations and recombine them into a single vector output. This was basically the reverse of the splitting operation we did before and more-or-less relied on the same matrix operations.
attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
new_shape = attn_output.size()[:-2] + (self_attn_block.num_heads * self_attn_block.head_dim,)
attn_output = attn_output.view(new_shape)
vanilla_attn_output = op.permute_tensor(vanilla_attn_output, [0, 2, 1, 3])
vanilla_new_shape = op.tensor_shape(vanilla_attn_output)[:-2] + (serial_attn_block.num_heads * serial_attn_block.head_dim,)
vanilla_attn_output = op.reshape_tensor(vanilla_attn_output, vanilla_new_shape)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)
# out projection and residual dropout
attn_output = self_attn_block.out_proj(attn_output)
vanilla_attn_output = serial_attn_block.out_proj(vanilla_attn_output)
assert block.vectors_match(attn_output.detach().cpu().tolist(), vanilla_attn_output)
# final test
actual_attn_output = self_attn_block(attn_input)[0]
assert block.vectors_match(actual_attn_output.detach().cpu().tolist(), vanilla_attn_output)Next Steps
And voila! We reproduced a critical section of our LLM using just vanilla python! While not exactly performant, in still a matter of seconds this purely serial version of our matrix operations was able to have exactly the same output as a standard self attention module. This exercise definitely gave me a greater appreciation for the nuance of what goes into modern more sophisticated language models but also served as a real reminder that at the end of the day LLMs really are just numbers in a matrix. In our next post we'll see how to accelerate these calculations as we get deeper into what GPU really is and finally we'll go full circle and try to get these same computations running on an FPGA. So stay tuned for the next exciting developments!