It’s an interesting journey to revisit Kaparthy’s GPT-2 creation focusing on codes and math.

One of the first pieces is understanding how automatic differentiation works. In a minimal neural network framework, every value keeps track of where it came from. The system remembers that out was created from self and other, so during backpropagation it knows how to propagate gradients backward. This tiny example is the foundation of frameworks like PyTorch.

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

PyTorch’s nn.Module: The Building Block of Modern Neural Networks. Moving from a toy autograd engine to PyTorch, everything is organized around the base nn.Module. It manages the entire parameter tree of a neural network. A GPT model is essentially a hierarchy of modules:

GPT
├── Token Embedding
├── Position Embedding
├── Transformer Blocks
│ ├── LayerNorm
│ ├── Causal Self Attention
│ └── MLP
└── Output Head

In coding Q, K, V, first project the original dimensions/n_embd to be tripled, say 768 to 2304. So then to compute Attention(Q,K,V) = softmax(QKᵀ / √d) × V

Note, Causal Attention is applied, which is the mechanism of decoder or autoregressive or causal mask, all is referring to that token only can see previous and itself, not the future tokens.

There is Flash Attention changes the computation strategy. Instead of materializing the entire attention matrix, it: reduces memory movement computes attention in smaller blocks; keeps intermediate values in fast GPU memory.

PyTorch provides register_buffer(). A buffer is a tensor that belongs to the model but is not a parameter, tbd…