Note | NanoGPT: Reinventing the Wheel of Transformers

The AI revolution all started with the release of ChatGPT, which has since become a household name. The underlying technology behind GPT models, known as transformers, is the key to today’s AI boom. To understand transformers, your best bet is to build a small-scale version of them. Here I follow the NanoGPT implementation to build a simple GPT model from scratch. Detailed tutorials can be found in Karpathy’s YouTube series.

1. Autograd and Backpropagation

The first step when we build a GPT model is to reproduce the autograd and BP algorithm, which is the core of deep learning. Interestingly, their first letters are respectively “a” and “b”, a coincidence that I find amusing since it somehow implies the fact that the two concepts are critically important to ML just like the “alphabet” in the language system :).

To begin with, a value class is necessary to store the data and, most importantly, the history of operations that produced it. The history is crucial for backpropagation, as it allows us to compute gradients with respect to the inputs. A Value class has methods for basic arithmetic operations, and each operation creates a new Value object that keeps track of its parents and the operations performed.

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0.0
        self._prev = set(_children) # Store the previous nodes (i.e., the numerical inputs to the operation) of the computation graph
        self._op = _op # Store the operation that produced this value (for visualization purposes)
        self._backward = lambda: None # A placeholder for the backward function, which will be defined later

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')

        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward

        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')

        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward

        return out

    def backward(self):
        topo = []
        visited = set()

        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)

        build_topo(self)

        self.grad = 1.0
        for v in reversed(topo):
            v._backward()

Notably, the recursive backward() method performs a post-order DFS of the computation DAG to generate a topological odering of the nodes, ensuring that each node’s gradient is computed only after all its children have been processed. But you may wonder why we need to use a DAG instead of a tree. The reason is that the same value might be used in multiple operations, leading to shared nodes (or more precisely, a common child shared by multiple parents … sounds weird XD) in the computation graph. A tree structure would not allow for this sharing, while a DAG does. Thus a DAG is much more efficient in terms of memory and computation, as it avoids redundant copies of the same nodes by allowing for the reuse of intermediate results.

There is another question: why do we use a post-order DFS and then reverse the order of the nodes? Actually, using a pre-order DFS like Kahn’s algorithm to traverse the graph and compute the gradients is also an approach, but it would require much more complex bookkeeping to manage the dependencies between the nodes, which would make the implementation significantly harder. By using a post-order DFS, beginners like me can avoid this complexity and ensure that each node’s gradient is computed in a correct order.

What’s more, the use of += in the _backward() method is crucial for handling shared nodes. Since it is common for a node to have multiple parents, its gradient needs to be accumulated from all the paths that lead to it. The += operator allows us to sum the contributions from each parent, ensuring that the final gradient reflects the total influence of all the operations which depend on that node.

For instance, consider the following example:

    A
   / \
  B   C
   \ /
    D

Our goal is to compute $\frac{\mathrm{d}D}{\mathrm{d}A}$.

If we use = instead of +=, the gradient of D would only reflect the contribution from one of its parents (either A-B or A-C), leading to an incorrect result. By using +=, we ensure that gradient accumulates contributions from both paths, giving us the correct total gradient.