Plato’s Cave (Jan Pietersz Saenredam after Cornelis Cornelisz van Haarlem 1604)

II. Neural Networks

That is, although using the precise tools of mathematics to specify architectures, loss functions, and optimizers, the essence of the discipline is more akin to synthetic chemistry or microbiology which “grow” chemicals and cells. The deep learning revolution started bearing fruit throughout the 2010s with magical technology such as Translate, Siri, and Autopilot, but most recently with ChatGPT and Claude whose core loss function you already understand to be the cross entropy loss $\mathcal{L}$ on next token prediction, which corresponds to maximizing the likelihood $p(d|\mathbf{w})$. It’s amazing how close you are to deep neural nets and the transformer architecture powering ChatGPT without even knowing. This is what we’ll cover in this second part of SITP.

In §3. Self-Supervised Sequence Learning with Deep Networks, you will increase the expressivity of the linear models implemented in Part I with non-linearities to get a class of models known as deep neural networks. We will start from feedforward neural networksFFN (FFN), and move on to other architecture variants with convolutional neural networksCNN (CNN), recurrent neural networksRNN (RNN), and finally, generative pre-trained transformersGPT (GPT), specifically by reproducing GPT-2Presented in Language Models are Unsupervised Multitask Learners (Raford et al. 2019) following karpathy/nanogpt. These various neural network architectures express and implement different inductive biases, which were all explored during the 2012-2019 time period of what is coloquially known as the age of researchage of research.

Then, in §4. The Engine of Automatic Differentiation, you will evolve the implementation of borscht from a numerical linear algebra library like numpy from §2. The Multidimensional Parallelized Tensor into a full blown batteries-included deep learning framework like torch which will be able to run the training and inference of the said nets just trained in §3. This means implementing the optimizersoptimizers for neural network primitives whose evaluations are accelerated with massively parallel processors known as GPUsGPU and whose gradients are automatically evaluated with an automatic differentationautomatic differentiation engine.

After completing Part II. Neural Networks, you will be ready for Part III. Scaling Networks of the book where we finally achieve our quest of building our own ChatGPT following karpathy/nanochat and subsequently, PyTorch 2 with borscht.compile() and borscht.distributed(). Part three follows the 2020-2025 time period of what is colloquially known as the the age of scalingage of scaling where researchers focused on scaling up the generality of transformer architecture by adding additional phases after pre-trainingpre-training stochastic parrots, namely the mid-trainingmid-training and post-trainingpost-training phases which respectively produce helpful assistantsassistants with reinforcement learning from human feedback and slow thinking reasonersreasoners with verifiable rewards.

But before we get there we have some work to do.
Let’s go.

II. Table of Contents

3. Self-Supervised Sequence Learning with Deep Networks

$\hookleftarrow$ Table of Contents

3.1 Learning Representations with FFNs

$\hookleftarrow$ Table of ContentsYoshua Bengio (Turing Award 2018)A Neural Probabilistic Language Model (2000)

In which we increase the expressivity of the architecture from linear to non-linear models, namely by composing many layers of logistic regression to form of a neural network.

3.1.1 From Localist to Distributional Representations

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

Linear ReLU Embed Conv2D Attention

3.1.2 From Linear Perceptrons to Non-Linear Multi Layer Perceptrons

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

The next natural question to ask then is whether the logistic regression model is considered a deep neural network? The answer is that technically yes, it can be considered a degenerative deep neural network with $0$ hidden layersIn the same way that a list can be considered a degenerative binary tree or graph. These so-called hidden layers automate the construction of the representation through learning, so that the model not only discovers the mapping from representation to output, but also the representation itself. The functions of deep neural networks will take the form of $f(\mathbf{x}):=g\circ \varphi (\mathbf{x})$ with $g$ being a linear classifier on feature extractor $\varphi :{\mathbb{R}}^{d_1}\to {\mathbb{R}}^{d_n}$, $\varphi (\mathbf{x}):=h^{[l]}\circ h^{[l-1]}\circ \cdots h^{[1]}\circ h^{[0]}(\mathbf{x})$ where $l$ is the number of compositional layers, and each intermediate function $h^{[i]}$ has the form of $h^{[i]}:{\mathbb{R}}^{d_i}\to {\mathbb{R}}^{d_{i+1}}$. Each hidden layer successively and gradually lifts the complexity and abstraction of the data’s representationChris Olah, cofounder of Anthropic and the lead of it’s interpretability research wrote an excellent article on how software 2.0’s representation learning loosely correspond to software 1.0’s types in Neural Networks, Types, and Functional Programming.

Together, these two aspects of learning non-linear, representations form the essence of deep learning.

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

But what is a neural network? | Deep learning chapter 1 (3Blue1Brown 2017)

Let’s now turn out attention to the function bodies of these $h^{[i]}$’s with a deep neural netork of $1$ hidden layer, carrying out the task of price regression and sentiment classification so that $f$ has the form $f:{\mathbb{R}}^d\to \mathbb{R}$. With the statistical learning foundation from part one, we will simply present the forward pass $f(\mathbf{x})$, the loss function $\mathcal{L}(d,\theta )$, and the backward pass ${\ell }^{\prime }(y,\hat{y},\theta )$.

Forward Pass

The functions of deep neural networks will take the form of $f(\mathbf{x}):=g\circ \varphi (\mathbf{x})$ with $g$ being a linear classifier on feature extractor

$$\begin{array}{rl}\varphi & :{\mathbb{R}}^{d_1}\to {\mathbb{R}}^{d_n}\\ \varphi (\mathbf{x})& :=h^{[l]}\circ h^{[l-1]}\circ \cdots h^{[1]}\circ h^{[0]}(\mathbf{x})\end{array}$$ where $l$ is the number of compositional layers, and each intermediate function $h^{[i]}$ has the form of $h^{[i]}:{\mathbb{R}}^{d_i}\to {\mathbb{R}}^{d_{i+1}}$.

Loss Function

Backward Pass

"""
model: Neural Language Models (Bengio et al. 2003) URL: https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdf

Dimension key:
# windows
B: batch size
T: sequence length

# input/output
V: vocabulary size
E: embedding dimension
D: model dimension
"""
import torch
import matplotlib.pyplot as plt
%matplotlib inline
# from jaxtyping import ...
g = torch.Generator().manual_seed(1337) # for .randn()

B, T = 32, 3
V, E, D = 27, 10, 200

# step: 0/200000, loss 27.63208770751953
# -> expected loss = nll = p(c) = -torch.tensor(1/V=27).log() = 3.2958
# -> self.W = torch.randn() is sampling from N(0, 1)
# -> self.W * [gain/sqrt(D_in)] (torch.init_kaimingnormal())

# residuals + normalization + Adam/RMSprop has made initialization less fragile
# -> b/c initialization is fragile/intractable with *deep* neural networks

class Linear:
    def __init__(self, D_in, D_out, bias=True):
        self.W_DiDo = torch.randn((D_in, D_out), generator=g) * (5/3)/D_in**0.5 # kaiming init (He et al. 2015)
        self.b_Do = torch.zeros(D_out) if bias else None

    def __call__(self, X_Di):
        self.X_Do = X_Di @ self.W_DiDo
        if self.b_Do is not None:
            self.X_Do += self.b_Do
        self.out = self.X_Do
        return self.X_Do

    def parameters(self):
        return [self.W_DiDo] + ([] if self.b_Do is None else [self.b_Do])

class BatchNorm1D:
  
  def __init__(self, dim, eps=1e-5, momentum=0.1):
    self.eps = eps
    self.momentum = momentum
    self.training = True
    # parameters (trained with backprop)
    self.gamma = torch.ones(dim)
    self.beta = torch.zeros(dim)
    # buffers (trained with a running 'momentum update')
    self.running_mean = torch.zeros(dim)
    self.running_var = torch.ones(dim)
  
  def __call__(self, x):
    # calculate the forward pass
    if self.training:
      xmean = x.mean(0, keepdim=True) # batch mean
      xvar = x.var(0, keepdim=True) # batch variance
    else:
      xmean = self.running_mean
      xvar = self.running_var
    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance
    self.out = self.gamma * xhat + self.beta
    # update the buffers
    if self.training:
      with torch.no_grad():
        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
    return self.out
  
  def parameters(self):
    return [self.gamma, self.beta]

class Tanh:
    def __call__(self, X_BD):
        self.X_BD = torch.tanh(X_BD)
        # plt.hist(self.X_BD.view(-1).tolist(), 50); # distribution of weights
        # plt.imshow(self.X_BD.abs() > 0.99, cmap='gray', interpolation='nearest') # vanishing gradients
        self.out = self.X_BD
        return self.X_BD
    
    def parameters(self):
        return []

model = [
    Linear(T * E, D, bias=False), BatchNorm1D(D), Tanh(),
    Linear(D, D, bias=False), BatchNorm1D(D), Tanh(),
    Linear(D, V, bias=False), BatchNorm1D(V)
]

C = torch.randn((V,E), generator=g)
params = [C] + [p for l in model for p in l.parameters()]
for p in params:
    p.requires_grad = True

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.1.3 From Direct to Iterative Optimization

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

In §1. Self-Supervised Sequence Learning with Single-Layer Networks, we became familar with two approaches to model $\hat{f}$, namely the bigram and logistic regression models corresponding to a stochastic matrix $P$ and a squashed weighted sum $\sigma (w^{\top }\mathbf{x})$ respectively, which in turn are composed with decision boundary to decide on the final output $v\in V$. That is, $f:=$ and $f_{\theta }:=$ (expand). Particularly with the parameterized logistic regression $f_{\theta }$, we specified a loss function $\mathcal{L}(\theta )$ (expand, loss.mle) of such parameters and optimized it with iterative reweighted least squares. Before we move onto increasing the expressivity of our models $\hat{f}$ from single layer networks to multi layer networks, let us first introduce a new method of optimizing the loss function, namely gradient descentgradient descent which is the fundamental optmization technique for training neural nets. It wouldn’t be an overstatement to call it the $F=ma$ of deep learning.

The intuition behind gradient descent is quite simpleIn the case of ${\mathbb{R}}^2$ and ${\mathbb{R}}^3$ which correspond to our physical intuition. We will see how this breaks down in higher dimensions ${\mathbb{R}}^d$..

${\theta }_{i+1}:={\theta }_i-\eta \nabla \mathcal{L}(\theta )$

(pseudocode) (borschtcode borscht.backwards())

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.2 Iterative Optimization

$\hookleftarrow$ Table of ContentsIsaac NewtonPhilosophiæ Naturalis Principia Mathematica (1687)

In which we transition from direct to iterative optimization methods, given that neural networks we will encounter are highly non-convex.

3.2.1 Taylor Series and their Derivatives, Gradients, and Hessians

3.2.2 First Order Optimization with Steepest Descent

3.2.3 First Order Optimization with Gradient Descent

3.2.4 Second Order Optimization with Newton’s Method

3.2.5 Second Order Optimization with Gauss Newton Method

3.3 Learning Representations with CNNs

$\hookleftarrow$ Table of ContentsYann LeCun (Turing Award 2018)Convolutional Networks for Images, Speech, and Time-Series (1995)

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.4 Learning Representations with RNNs

$\hookleftarrow$ Table of Contents

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

class RNNCell(nn.Module):
  """
  the job of a 'Cell' is to:
  take input at current time step x_{t} and the hidden state at the
  previous time step h_{t-1} and return the resulting hidden state
  h_{t} at the current timestep
  """
  def __init__(self, config):
    super().__init__()
    self.xh_to_h = nn.Linear(config.n_embd + config.n_embd2, config.n_embd2)

  def forward(self, xt, hprev):
    xh = torch.cat([xt, hprev], dim=1)
    ht = F.tanh(self.xh_to_h(xh))
    return ht

class RNN(nn.Module):
  def __init__(self, config, cell_type):
    super().__init__()
    self.block_size = config.block_size
    self.vocab_size = config.vocab_size
    self.start = nn.Parameter(torch.zeros(1, config.n_embd2)) # the starting hidden state
    self.wte = nn.Embedding(config.vocab_size, config.n_embd) # token embeddings table
    if cell_type == 'rnn':
        self.cell = RNNCell(config)
    elif cell_type == 'gru':
        self.cell = GRUCell(config)
    self.lm_head = nn.Linear(config.n_embd2, self.vocab_size)

  def get_block_size(self):
    return self.block_size

  def forward(self, idx, targets=None):
    device = idx.device
    b, t = idx.size()

    # embed all the integers up front and all at once for efficiency
    emb = self.wte(idx) # (b, t, n_embd)

    # sequentially iterate over the inputs and update the RNN state each tick
    hprev = self.start.expand((b, -1)) # expand out the batch dimension
    hiddens = []
    for i in range(t):
      xt = emb[:, i, :] # (b, n_embd)
      ht = self.cell(xt, hprev) # (b, n_embd2)
      hprev = ht
      hiddens.append(ht)

    # decode the outputs
    hidden = torch.stack(hiddens, 1) # (b, t, n_embd2)
    logits = self.lm_head(hidden)

    # if we are given some desired targets also calculate the loss
    loss = None
    if targets is not None:
      loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

    return logits, loss

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.5 Learning Representations with GPTs

$\hookleftarrow$ Table of ContentsAttention Is All You Need (2017)Ashish Vaswani

In which we transition from long-range depencies of recurrent neural nets to the ___ of transformer-based neural nets.

3.5.1 Transformers Architecture Overview

Stanford CS229 Machine Learning Lecture 14: Transformers, In-Context Learning

Stanford CME295 Transformers & LLMs | Lecture 1 - Transformer

Stanford CS25: V2 I Introduction to Transformers w/ Andrej Karpathy

3.5.2 Embedding $W_E$ and Unembedding $W_U$

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

Transformers, the tech behind LLMs | Deep Learning Chapter 5 (3Blue1Brown 2024)

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.5.3 Attention: Keys $W_K$, Queries $W_Q$, and Values $W_V$

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

class NewGELU(nn.Module):
  """
  Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
  Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
  """
  def forward(self, x):
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

class CausalSelfAttention(nn.Module):
  """
  A vanilla multi-head masked self-attention layer with a projection at the end.
  It is possible to use torch.nn.MultiheadAttention here but I am including an
  explicit implementation here to show that there is nothing too scary here.
  """

  def __init__(self, config):
    super().__init__()
    assert config.n_embd % config.n_head == 0
    # key, query, value projections for all heads, but in a batch
    self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
    # output projection
    self.c_proj = nn.Linear(config.n_embd, config.n_embd)
    # causal mask to ensure that attention is only applied to the left in the input sequence
    self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                  .view(1, 1, config.block_size, config.block_size))
    self.n_head = config.n_head
    self.n_embd = config.n_embd

  def forward(self, x):
    B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

    # calculate query, key, values for all heads in batch and move head forward to be the batch dim
    q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
    k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
    q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
    v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

    # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
    att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
    att = F.softmax(att, dim=-1)
    y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
    y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

    # output projection
    y = self.c_proj(y)
    return y

class Block(nn.Module):
  """ an unassuming Transformer block """

  def __init__(self, config):
    super().__init__()
    self.ln_1 = nn.LayerNorm(config.n_embd)
    self.attn = CausalSelfAttention(config)
    self.ln_2 = nn.LayerNorm(config.n_embd)
    self.mlp = nn.ModuleDict(dict(
        c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),
        c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),
        act     = NewGELU(),
    ))
    m = self.mlp
    self.mlpf = lambda x: m.c_proj(m.act(m.c_fc(x))) # MLP forward

  def forward(self, x):
    x = x + self.attn(self.ln_1(x))
    x = x + self.mlpf(self.ln_2(x))
    return x

class Transformer(nn.Module):
  """ Transformer Language Model, exactly as seen in GPT-2 """

  def __init__(self, config):
    super().__init__()
    self.block_size = config.block_size

    self.transformer = nn.ModuleDict(dict(
        wte = nn.Embedding(config.vocab_size, config.n_embd),
        wpe = nn.Embedding(config.block_size, config.n_embd),
        h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
        ln_f = nn.LayerNorm(config.n_embd),
    ))
    self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

    # report number of parameters (note we don't count the decoder parameters in lm_head)
    n_params = sum(p.numel() for p in self.transformer.parameters())
    print("number of parameters: %.2fM" % (n_params/1e6,))

  def get_block_size(self):
    return self.block_size

  def forward(self, idx, targets=None):
    device = idx.device
    b, t = idx.size()
    assert t <= self.block_size, f"Cannot forward sequence of length {t}, block size is only {self.block_size}"
    pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)

    # forward the GPT model itself
    tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
    pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)
    x = tok_emb + pos_emb
    for block in self.transformer.h:
        x = block(x)
    x = self.transformer.ln_f(x)
    logits = self.lm_head(x)

    # if we are given some desired targets also calculate the loss
    loss = None
    if targets is not None:
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

    return logits, loss

Intermezzo Three: The Language of Neural Networks

I3.1 Inner Product Spaces, Normed Spaces, Hilbert Spaces

I3.2 Derivatives in Hilbert Spaces

I3.3 Superposition

4. The Engine of Automatic Differentiation

$\hookleftarrow$ Table of Contents

After §1. Self-Supervised Sequence Learning with Single-Layer Networks and §3. Self-Supervised Sequence Learning with Deep Networks, we have now built foundation in the architecture of deep neural networks, which culminated in implementing the training and inference loops for nanogpt in §3.5 Learning Representations with GPTs.

First, we understand that language models such as nanogpt, Chat, and Claude are specified as next-token predictors $f:V^{\ast }\to V$ which although decide on a final classifying token $\mathbf{y}\in V$ for the given context $({\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_t)\in V^{\ast }$ (todo $V^t$?), have a step of inference by modeling language stochastically with a conditional random variable $Y|X_1,X_2,\dots ,X_N$ endowed with a probability distribution $p_{Y|X_1,X_2,\dots ,X_N}(\mathbf{y}|{\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_t)$ where such distribution corresponds to an array $\mathbf{y}\in {\mathbb{R}}^V$ which can be represented by a borscht.Tensor with .shape of (V). This was the first step in our transition from classical software 1.0 to the jazz of software 2.0, namely, by transitioning from variables to random variables in §1.1 From Certain to Uncertain Knowledge.

Second, we understanding that such language models probabilistically specified as autoregressive next-token predictors $f$ are implemented by estimating the parameterization of their functions with weights $\mathbf{w}$ and bias $b$ so that $f_{\theta }$ where $\theta :=\{\mathbf{w},b\}$ is a function that takes in a one-hot vector $\mathbf{x}$ as input (expand, context?), evaluates the weighted sum with $\mathbf{w}$ which is the standard inner product of the dot product, and then either uses the dot product directly in the case of linear regression with §1.3 Parameterizing Classification with Linear Regression or further composes the output with a squashing function $\sigma (\cdot )$ in §1.4 Parameterizing Classification with Logistic Regression. In either case, both models $f$ are specified dimensionally with linear algebra and implemented with the optimization of differential calculus.

Third, we increased the expressivity of such implementations by transitioning from localist representations of one hot vectors to distributional representations of feature map $\varphi (\mathbf{x})$. We started with a vanilla multi-layer perceptron also known as a feedforward neural net.

Fourth, we explored various inductive biases, exploring FFNs, CNNs, RNNs, culminating in GPTs with their core attention mechanism.

Moreover, after §2. The Parallelized Multidimensional Tensor, we also have foundation in the systems of deep neural networks, which culminated in implementing the training and inference loops for nanogpt in §3.5 Learning Representations with GPTs.

First, layer -1 implemented our own borscht.Tensor by §2.1 Virtualizing Shapes with Strides

and layer 0 §2.2 Accelerating Basic Linear Algebra on CPUs.

automatic differentiationautomatic differentiation optimizersoptimizers network primitivesnetwork primitives parallel programmingparallel programming massively parallel processorsmassively parallel processors graphics processing unit (GPU)GPU CUDA RustCUDA Rust PTXPTX SASSSASS

Finish strong young hacker.

4.1 Gradient-Based Optimization with Automatic Differentation

$\hookleftarrow$ Table of ContentsMatthew JohnsonAutograd tutorial

In which we implement the engine of automatic differentiation, the magic spell AI researchers cast to receive $\nabla \mathcal{L}(\theta )$ just from the specifcation of $\mathcal{L}$.

4.1.1 Symbolic, Numerical, and Algorithmic Differentiation

$\hookleftarrow$ Table of Contents

4.1.2 Forward Mode with borscht.forward()

$\hookleftarrow$ Table of Contents

import picograd as pg

def f(x1: pg.Tensor, x2: pg.Tensor) -> pg.Tensor:
  a = pg.exp(x1)
  b = pg.sin(x2)
  c = b**2
  d = a*c
  return d

Figure 1. Python source for the function $f:\mathbb{R},\mathbb{R}\to \mathbb{R}$ where $f:=\exp (x_1)\ast {\sin }^2{x}_2$

Here we’ve broken up the function to render the subexpressions more clearly. But this isn’t necessary — automatic differentiation will work if the function was expressed in one line. In part one, the development of picograd followed that of numpy — an array programming language similar to Matlab but embedded in the host language of Python, that could evaluate functions of the form $f(\mathbf{x})$ where Tensor objects stored their values $f(\mathbf{x})$ with the value: field and the function types that produced their values with Op. For instance, evaluating the specified function f from above with 9 and 10

if __name__ == "__main__":
  print(f(9, 10))

populates the Tensor.value fields. In part one of the book we verified this with a REPL-interface, but we can also represent the entire expression being evaluated with a graph of vertices and edges where the vertices are Tensors (along with their Ops and values) and the edges are their data dependencies:

Here you can see that even if the function was specified in one line, the graph of the expression always parses into Tensor vertices, and data dependency edges. You may have noticed the Tensor.grad fields, which supposedly store the values of derivatives $f^{\prime }(\mathbf{x})$. The question now remains in how to populate these fields.

Taking a step back to differential calculus, deriving the derivative of $f:=\exp (x_1)\ast {\sin }^2{x}_2$ involves the application of the chain rule where $f(x):=g(h(x))⟹f^{\prime }(x)=g^{\prime }(h(x))h^{\prime }(x)$. Evaluating the derivative of the function with respect to its inputs $f^{\prime }(x_1)$ and $f^{\prime }(x_2)$ results in

$$f^{\prime }(x_1)=\frac{\partial }{\partial x_1}[\exp (x_1)\ast {\sin }^2{x}_2]=\frac{\partial }{\partial x_1}[\exp (x_1)]\ast {\sin }^2{x}_2=\exp (x_1)(1)\ast ({\sin }^2{x}_2)$$ $$f^{\prime }(x_2)=\frac{\partial }{\partial x_2}[\exp (x_1)\ast {\sin }^2{x}_2]=\exp (x_1)\ast \frac{\partial }{\partial x_2}[{\sin }^2{x}_2]=\exp (x_1)\ast 2\sin x_2\cos x_2$$

symbolic and numeric differentiattion symbolic differentiation has performance issues since a large unrolled expression must be constructed in order to differentiate[^0], whereas numerical differentiation has correctness issues since evaluating finite differences requires evaluating functions to a precision point resulting in numerical instability. (trace through EXAMPLE for both. talking nets widrow)

To populate the Tensor.grad fields, the simplest idea would be to literally translate the manual derivation of the derivative into code. The translation from math to code involves a design decision: should we evaluate from outputs to inputs (symbolically outside-in, graphically right-to-left) or from inputs to outputs (symbolically inside-out, graphically left-to-right)? Although the former order seems more natural with symbolic expressions, there’s nothing illegal about the latter.

import picograd as pg

def f(x1: pg.Tensor, x2: pg.Tensor) -> pg.Tensor:
  a = pg.exp(x1)
  b = pg.sin(x2)
  c = b**2
  d = a*c
  return d

# dict[f(x), f'(x)] of local derivatives (adjoints)
dd_da, dd_dc = [c, a] # d(a,c):=a*c ==> d'(a)=c, d'(c)=a
da_dx1 = pg.exp(x1) # a(x1):=exp(x1) ==> a'(x1)=exp(x1)
dc_db = 2*b # c(b):=b^2 ==> c'(b)=2b
db_dx2 = pg.cos(x2) # b(x2):=sin(x2) ==> b'(x2)=cos(x2)

# outputs to inputs: outside-in symbolically, right-to-left graphically
dd_dd = pg.Tensor(1) # base case
dd_da, dd_dc = [dd_dd*dd_da, dd_dd*dd_dc]
dd_dx1 = dd_da*da_dx1 # DONE for the x1->d path

dd_db = dd_dc*dc_db
dd_dx1 = dd_db*db_dx2 # DONE for x2->path

# inputs to outputs: inside-out symbolically, left-to-right graphically
dx1_dx1, dx2_dx2 = [pg.Tensor(1), pg.Tensor(1)] # base case
da_dx1 = da_dx1*dx1_dx1
dd_dx1 = dd_da*da_dx1 # DONE for the x1->d path

db_dx2 = db_dx2*dx2_dx2
dc_dx2 = dc_dc*db_dx2
dd_dx2 = dd_dc*dc_dx_2 # DONE for the x2->d path

Do you notice any difference in the number of evaluations between the two orders?

4.1.3 Backward Mode with borscht.backward()

$\hookleftarrow$ Table of Contents

$$\frac{\partial f}{\partial x_i}=\sum \prod \frac{\partial v_j}{\partial v_k}$$

and as long as the operations along this path are associative — $(AB)C=A(BC)$ — then we can choose the order in how we perform these path products to minimize the number of operations. Finding the optimal ordering is an NP-hard problem because ____. For instance, if the expression graph is diamond-shaped, evaluating the derivative with forward-mode for the left-half and reverse-mode for the right-half would be more performant. In practice, we use reverse-mode as a heuristic, since most of the functions that are differentiated (so they can be optimized) in the field of machine learning are neural networks of the form $f:{\mathbb{R}}^n\to \mathbb{R}$

$\hookleftarrow$ Table of Contents

How can we generalize this into an algorithm?
All we need are 1. mappings from $f(\mathbf{x})\to f^{\prime }(\mathbf{x})$ and 2. a topological sort

For the derivative rules, the same way that optimizing compilers implement an optimization “manually” once which then gets reused many times, the authors of deep learning frameworks also implement derivatives manually which then become reused many times through automatic differentiation. In theory, we can differentiate any expression with f’(x) with only a few derivative rules for addition and multiplication, but in practice most frameworks provide sugar for complex derivatives.

For topological sort, we can simply reversed the ordering produced by a depth-first-search:

def toposort(self):
  order: list[Op] = []
  visited: set[Op] = set()

  def dfs(node: Op) -> None:
    if node in visited: return
    visited.add(node)
    for src in node.src: dfs(src)
    order.append(node)

  dfs(self)
  return order

class Tensor():
  def backward():
    for t in reversed(topo):
      t.backward()

We will now use this idea to modify the interpretation of our deep learning framework to not only evaluate $f(\mathbf{x})$, but $f^{\prime }(\mathbf{x})$ as well. This is done by dynamically overloading the operators at runtime[^0] to trace the expression graph

chain_rules = PatternMatcher([
  (Pattern(OpCode.MATMUL, name="input"), lambda output_grad, input: (_____,)),
  (Pattern(OpCode.MATVEC, name="input"), lambda output_grad, input: (_____,)),
  (Pattern(OpCode.RECIPROCAL, name="input"), lambda output_grad, input: (-output_grad * input * input,)),
  (Pattern(OpCode.SIN, name="input"), lambda output_grad, input: ((math.pi/2 - input.src[0]).sin() * output_grad,)),
  (Pattern(OpCode.LOG2, name="input"), lambda output_grad, input: (output_grad / (input.src[0] * math.log(2)),)),
  (Pattern(OpCode.EXP2, name="input"), lambda output_grad, input: (input * output_grad * math.log(2),)),
  (Pattern(OpCode.SQRT, name="input"), lambda output_grad, input: (output_grad / (input*2),)),
  (Pattern(OpCode.ADD), lambda output_grad: (1.0*output_grad, 1.0*output_grad)),
  (Pattern(OpCode.MUL, name="input"), lambda output_grad, input: (input.src[1]*output_grad, input.src[0]*output_grad)),
])

class Tensor:
  def _forward(self, f:Callable, *other:Tensor) -> Tensor: #extra_args=(), **kwargs)
    out_tensor = evaluator.eval_uop([self, other], out_uop)

  def backward(self, grad:Tensor|None=None) -> Tensor:
    """
    backward performs by collecting tensors, computing gradients with automatic differentiation, and updating said tensors.
    """
    # 1. collect all tensors that requires grad by topologically sorting the graph of uops and filter
    all_uops = self.uop.toposort()
    tensors_require_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and t.uop in all_uops and t.requires_grad]
    uops_require_grad = [t.uop for t in tensors_require_grad]
    assert grad is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
    if not (self.is_floating_point() and all(t.is_floating_point() for t in tensors_require_grad)): raise RuntimeError("only float Tensors have gradient")
    
    # 2. compute the gradient with a map of tensors to partials
    if grad is None: grad = Tensor(1.0, dtype=self.dtype, device=self.device, requires_grad=False) # base case is 1.0
    tens2grads = Tensor._automatically_differentiate(self.uop, grad.uop, set(uops_require_grad)) # skipping materializing zerod grads for now
    grads = [Tensor(g, device=t.device) for t,g in zip(tens2grads.keys, tens2grads.values)] # initialize tensor grads on device
    
    # 3. update the tensors that require grad with the gradient's partials
    for t,g in zip(tensors_require_grad, grads):
      assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
      t.grad = g if t.grad is None else (t.grad + g) # accumulate if t.grad exists
    return self

  @staticmethod
  def _automatically_differentiate(root:Op, root_grad:Op, targets:set[Op]) -> dict[Op, Op]:
    """
    _differentiate backpropagates partials on a topologically sorted expression graph with the chain rule
    and produces the gradient in the form of a map of ops to their partials (which, in turn, are ops)
    """
    tens2grads = {root: root_grad}

    # 1. topological sort
    in_target_path: dict[Op, bool] = {}
    for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
    dfs = list(root.toposort()) # lambda node: node.op not in {OpCode.DETACH, OpCode.ASSIGN} and in_target_path[node])) # don't flow through DETACH/ASSIGN or anything not in target path

    # 2. backpropagation with the chain rule
    for tensor in reversed(dfs):
      if tensor not in tens2grads: continue

      local_grads: tuple[Op|None, ...]|None = cast(tuple[Op, ...]|None, chain_rules.rewrite(tensor, ctx=tens2grads[tensor]))
      if local_grads is None: raise RuntimeError(f"failed to compute gradient for {tensor.op}\n\nin {str(tensor)[0:1000]}...")
      assert len(local_grads) == len(tensor.src), f"got {len(local_grads)} gradient, expected {len(tensor.src)}"

      for tensor,local_grad in zip(tensor.src, local_grads): # <--------------------- MOOOSE: why are we accumulating inside ad()? don't we do it in backward()??
        if local_grad is None: continue
        if tensor in tens2grads: tens2grads[tensor] = tens2grads[tensor] + local_grad # accumulate if tensor exists
        else: tens2grads[tensor] = local_grad # o/w initialize

To implement automatic differentiation with Tensor.backward(), there is a design decision to be made — the choice of implementing it dynamically or just-in-time[^3], similar to the decision of how to implement types for general programming languages[^4]. This stands in contrast to the alternative of performing a just-in-time, source-to-source transformation.

Let’s now move onto automatically differentiating the functions of neural networks, specifically the FFN language model from earlier. (johnson/ryan adams ordering) n^2 vs n^3

4.1.4 Stochastic Gradient Descent with borscht.optim.sgd()

$\hookleftarrow$ Table of Contents

4.1.5 Adam with borscht.optim.adamw()

$\hookleftarrow$ Table of Contents (§4.1 Automatic Differentiation)

4.2 Network Primitives

$\hookleftarrow$ Table of ContentsSoumith ChintalaPyTorch: An Imperative Style, High-Performance Deep Learning Library (2019)

In which we …

4.2.1 Representing Models with Graph<borscht.Tensor>

4.2.2 Implementing Forward Pass for FFN

In which we implement the forward pass for class MLP(nn.Module) with borscht.nn.Embedding() borscht.nn.Linear() borscht.nn.Tanh()

class MLP(nn.Module):
    """
    takes the previous block_size tokens, encodes them with a lookup table,
    concatenates the vectors and predicts the next token with an MLP.

    Reference:
    Bengio et al. 2003 https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdf
    """

    def __init__(self, config):
        super().__init__()
        self.block_size = config.block_size
        self.vocab_size = config.vocab_size
        self.wte = nn.Embedding(config.vocab_size + 1, config.n_embd) # token embeddings table
        # +1 in the line above for a special <BLANK> token that gets inserted if encoding a token
        # before the beginning of the input sequence
        self.mlp = nn.Sequential(
            nn.Linear(self.block_size * config.n_embd, config.n_embd2),
            nn.Tanh(),
            nn.Linear(config.n_embd2, self.vocab_size)
        )

    def get_block_size(self):
        return self.block_size

    def forward(self, idx, targets=None):

        # gather the word embeddings of the previous 3 words
        embs = []
        for k in range(self.block_size):
            tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
            idx = torch.roll(idx, 1, 1)
            idx[:, 0] = self.vocab_size # special <BLANK> token
            embs.append(tok_emb)

        # concat all of the embeddings together and pass through an MLP
        x = torch.cat(embs, -1) # (b, t, n_embd * block_size)
        logits = self.mlp(x)

        # if we are given some desired targets also calculate the loss
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

        return logits, loss

4.2.3 Implementing Forward Pass for CNN

4.2.4 Implementing Forward Pass for RNN

4.2.5 Implementing Forward Pass for GPT

import torch
import torch.nn as nn
import math

class NewGELU(nn.Module):
    """
    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
    """
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

class CausalSelfAttention(nn.Module):
    """
    A vanilla multi-head masked self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        # causal mask to ensure that attention is only applied to the left in the input sequence
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                     .view(1, 1, config.block_size, config.block_size))
        self.n_head = config.n_head
        self.n_embd = config.n_embd

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.c_proj(y)
        return y

class Block(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),
            c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),
            act     = NewGELU(),
        ))
        m = self.mlp
        self.mlpf = lambda x: m.c_proj(m.act(m.c_fc(x))) # MLP forward

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlpf(self.ln_2(x))
        return x

class Transformer(nn.Module):
    """ Transformer Language Model, exactly as seen in GPT-2 """

    def __init__(self, config):
        super().__init__()
        self.block_size = config.block_size

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # report number of parameters (note we don't count the decoder parameters in lm_head)
        n_params = sum(p.numel() for p in self.transformer.parameters())
        print("number of parameters: %.2fM" % (n_params/1e6,))

    def get_block_size(self):
        return self.block_size

    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.block_size, f"Cannot forward sequence of length {t}, block size is only {self.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)

        # forward the GPT model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)
        x = tok_emb + pos_emb
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)

        # if we are given some desired targets also calculate the loss
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

        return logits, loss

4.3 From Sequential Machines to Parallel Machines

$\hookleftarrow$ Table of ContentsIlya Sutskever, Geoffrey Hinton and Alex KrizhevskyImageNet Classification with Deep Convolutional Neural Networks (2012)

From Modal GPU Glossary (Streaming Multiprocessor), modified from a diagram in Fabien Sanglard’s blog, itself likely modified from a diagram in the CUDA C Programming Guide.

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec aliquam varius mi, pharetra scelerisque nulla vehicula ut. Cras ac faucibus dolor. Nunc tempus efficitur ultricies. Nullam non fermentum elit, et tincidunt ante. Donec hendrerit metus ac urna semper, in facilisis leo egestas. Integer sit amet quam ultrices, dapibus sem a, volutpat velit. Pellentesque eget ipsum consequat sem luctus blandit a et nibh. Mauris gravida felis ut dignissim volutpat. Quisque volutpat pulvinar enim in venenatis. Integer posuere tincidunt leo vel congue. Sed ex velit, posuere ut nibh nec, ultrices malesuada nibh. Fusce scelerisque mollis neque id fringilla. Etiam rhoncus quam a leo ultrices rutrum non vitae est. Donec bibendum, arcu sagittis consequat vestibulum, lectus elit bibendum libero, non rhoncus lorem lorem a neque. Sed ex felis, rutrum non consequat eget, tincidunt ac lacus.

Nulla quis dictum diam, a cursus nibh. Pellentesque magna ante, lobortis a aliquam at, imperdiet eget magna. Nam pharetra dolor lectus, eu hendrerit ipsum tincidunt molestie. Cras eget viverra ipsum. Cras consequat metus quis nibh fermentum, egestas vestibulum odio lacinia. Donec ut ante hendrerit, convallis sem pulvinar, pulvinar magna. Etiam consequat mattis turpis non mollis. Sed consequat et sem quis interdum. Vivamus blandit elementum tortor non convallis. Duis quis nunc nisi. Fusce scelerisque felis pellentesque ex elementum imperdiet. Nulla facilisi. Nunc molestie urna sit amet nisl ullamcorper, at convallis augue tempor.

https://nvlabs.github.io/cuda-oxide/getting-started/hello-gpu.html#install-cargo-oxide

So with that all said and done, let us extend the mixed source of Python and Rust for borscht with a third language for the devicedevice.

Recall that in §2.1.1 From Virtual to Physical Machines we installed the Rust toolchain so that rustc could translate our program into the machine language of a physical processor. That toolchain is still necessary, but it is no longer sufficient, because the physical machine we are now targeting is not the one running our program. A GPU program is really two programs — the hosthost program which runs on the CPU and orchestrates memory and launches, and the devicedevice program of kernelskernels which runs on the GPU. The two are compiled by two different compilers, for two different instruction sets, and then stitched back together.

Historically the device half of that pair was written in CUDA C++ and compiled by NVIDIA’s nvcc, which meant that a Rust deep learning framework had to reach across a foreign function interface to say anything to the GPU at all. We will not do that. Instead we will keep host and device in the same file and the same language by way of cuda-oxidecuda-oxideSee NVlabs/cuda-oxide. It is explicitly alpha software: expect bugs and API breakage., a codegen backendcodegen backend for rustc which compiles functions annotated #[kernel] down to PTXPTX, the virtual instruction set of NVIDIA GPUsPTX is virtual in exactly the sense of §2.1.1 — it is not what the silicon executes. The driver finishes the job, lowering PTX to the SASSSASS of your particular chip at load time.. Note that this is the same turtle we met in Part I: rustc never emitted x86 directly either, it emitted LLVM. Here it emits LLVM too, only the backend of that backend is NVPTX instead of X86.

So the toolchain grows a device half:

Concretely, you need five things. A Linux machine with an NVIDIA GPU and a driverdriver new enough for your toolkit. The CUDA ToolkitCUDA Toolkit (12.x or newer), which supplies nvcc, libNVVM, nvJitLink and libdevice. LLVM 21 or newer, whose llc performs the final lowering from LLVM IR to PTXThis floor is not negotiable for modern chips. The backend emits TMA, tcgen05 and WGMMA intrinsics that llc from LLVM 20 and earlier cannot lower — simple kernels may survive on an older llc, but anything Hopper or Blackwell will not.. Clang 21 with its development headers, because the host bindings are generated by bindgen, which loads libclang and wants clang’s own stddef.h. And a Rust nightly, pinned by rust-toolchain.toml, carrying the rust-src, rustc-dev and llvm-tools components — nightly because a codegen backend plugs into rustc’s own internals, which are not a stable interface.

That is a long enough list that the honest advice is to not install it by hand. The repository ships a development containerdevelopment container which pins every one of the five, so the only things your host must provide are the NVIDIA driver and the NVIDIA Container ToolkitInstalling the container toolkit is not quite the whole story: it must also be registered with the container runtime, via sudo nvidia-ctk runtime configure --runtime=docker followed by a daemon restart. Without that step --gpus=all has no runtime to bind to and the container starts with no GPU.. The CUDA Toolkit itself never touches your machine.

> root@machine ~/ $ git clone https://github.com/NVlabs/cuda-oxide.git
> root@machine ~/ $ devcontainer up --workspace-folder cuda-oxide
> root@machine ~/ $ devcontainer exec --workspace-folder cuda-oxide bash

Those with Nix may prefer nix develop, which gives the same guarantee by a different road. Those who insist on installing to the metal should consult the project’s installation chapter; the one step people trip on is LLVM, since most distributions do not package a version recent enough, and you will want llvm.sh:

> root@machine ~/ $ wget https://apt.llvm.org/llvm.sh && chmod +x llvm.sh && sudo ./llvm.sh 21
> root@machine ~/ $ llc-21 --version | grep nvptx
    nvptx      - NVIDIA PTX 32-bit
    nvptx64    - NVIDIA PTX 64-bit

However you arrive, the toolchain is driven by cargo oxide, a cargo subcommandcargo subcommand standing in the same relation to cargo that cargo stands to rustc. Inside the cuda-oxide repository it resolves through a workspace alias and works immediately; for a project of our own, such as borscht, it must be installed against the pinned nightly:

> root@machine ~/ $ cargo +nightly-2026-04-03 install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide
    Installing /usr/local/cargo/bin/cargo-oxide
     Installed package `cargo-oxide v0.2.1` (executable `cargo-oxide`)

Before writing a line of our own, ask it whether the machine is in order. cargo oxide doctor is the whole of the preceding discussion turned into a checklist, and it is worth reading line by line, because every entry is a piece of the pipeline we just described:

> root@machine ~/ $ cargo oxide doctor
cargo-oxide environment check
==============================
Rust nightly toolchain... ✓ rustc 1.96.0-nightly (55e86c996 2026-04-02)
rust-toolchain.toml... ✓ channel nightly-2026-04-03
Pinned toolchain active... ✓ nightly-2026-04-03-x86_64-unknown-linux-gnu
Required rustup components... ✓ rust-src, rustc-dev, rust-analyzer, clippy, rustfmt, llvm-tools
Codegen backend... - not built yet (run `cargo oxide setup`)
CUDA headers (cuda.h)... ✓ /usr/local/cuda/include/cuda.h
CUDA toolkit (nvcc)... ✓ Cuda compilation tools, release 13.0, V13.0.48
libNVVM (libnvvm.so)... ✓ libNVVM 2.0
nvJitLink (libnvJitLink.so)... ✓ nvJitLink 13.0
libdevice (libdevice.10.bc)... ✓ /usr/local/cuda/nvvm/libdevice/libdevice.10.bc
llc (LLVM)... ✓ Ubuntu LLVM version 21.1.8 (/usr/bin/llc-21)
clang / libclang resource dir... ✓ /usr/lib/llvm-21/lib/clang/21
NVIDIA driver / GPU... ✓ NVIDIA A100-SXM4-40GB (compute capability 8.0, driver 580.105.08)
cuda-gdb (optional)... ✓ /usr/local/cuda/bin/cuda-gdb
compute-sanitizer (optional)... ✓ Version 2025.3.0.0
✅ Environment looks good!

Two of those lines deserve a second look. The compute capabilitycompute capability 8.0 is the generation of the chip — sm_80 is Ampere, sm_90 Hopper, sm_100 Blackwell — and it is what decides which instructions are available to us in §4.4; you will meet it again as the thing that makes a kernel fast on one machine and illegal on another. And the codegen backend reported as not built yet is not an error: it is a compiler, rustc-codegen-cuda, and cargo oxide will build and cache it for you on first use.

Now we can scaffold, exactly as maturin new scaffolded the Python and Rust halves of borscht back in §2.1.2:

> root@machine ~/ $ cargo oxide new borscht-cuda
✓ Created cuda-oxide project 'borscht-cuda'

  cd borscht-cuda
  cargo oxide doctor
  cargo oxide run

> root@machine ~/ $ cd borscht-cuda && ls
Cargo.toml
README.md
rust-toolchain.toml
src

[package]
name = "borscht-cuda"
version = "0.1.0"
edition = "2024"

[workspace]

[dependencies]
cuda-device = { git = "https://github.com/NVlabs/cuda-oxide.git" }
cuda-host = { git = "https://github.com/NVlabs/cuda-oxide.git" }
cuda-core = { git = "https://github.com/NVlabs/cuda-oxide.git" }

[toolchain]
channel = "nightly-2026-04-03"
components = [
    "rust-src",
    "rustc-dev",
    "rust-analyzer",
    "clippy",
    "rustfmt",
    "llvm-tools",
]

The three dependencies partition the problem the same way the machine does. cuda-device is the vocabulary of the device — #[kernel], thread::index_1d(), the indexing types — and it is the crate whose code rustc-codegen-cuda will consume. cuda-core is the host-side runtime of contexts, streams and buffers. cuda-host supplies #[cuda_module], the macro that welds the two halves together by embedding the compiled device artifact into the host binary and generating a typed launch method for each kernel. And the pinned rust-toolchain.toml is what makes the earlier +nightly-2026-04-03 unnecessary from here on: rustup reads it and selects the toolchain on every cargo invocation, fetching llvm-tools on first use.

The generated src/main.rs is the smallest complete example of the shape all of our kernels will take, and it repays a careful reading before we write our own:

use cuda_device::{kernel, launch_bounds, launch_contract, thread, DisjointSlice};
use cuda_host::cuda_module;
use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D};

#[cuda_module]
mod kernels {
    use super::*;

    #[kernel]
    #[launch_bounds(256)]
    #[launch_contract(domain = 1, block = (256, 1, 1))]
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let idx_raw = idx.get();
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[idx_raw] + b[idx_raw];
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let ctx = CudaContext::new(0)?;
    let stream = ctx.default_stream();

    const N: usize = 1024;
    let a_host: Vec<f32> = (0..N).map(|i| i as f32).collect();
    let b_host: Vec<f32> = (0..N).map(|i| (i * 2) as f32).collect();

    let a_dev = DeviceBuffer::from_host(&stream, &a_host)?;
    let b_dev = DeviceBuffer::from_host(&stream, &b_host)?;
    let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N)?;

    // SAFETY: this package owns the embedded device bundle produced for the
    // kernels module above.
    let module = unsafe { kernels::load(&ctx)? };
    let prepared = module.prepare_vecadd(LaunchConfig1D::new((N as u32).div_ceil(256), 256, 0))?;
    module.vecadd(&stream, &prepared, &a_dev, &b_dev, &mut c_dev)?;

    let c_host = c_dev.to_host_vec(&stream)?;
    println!("{:?}", &c_host[..5]);
    Ok(())
}

Everything above the fn main runs on the GPU and everything below it runs on the CPU, in one file, in one language, built by one command. Notice that the ownership we learned in Part I has followed us across the bus: a_dev and b_dev are borrowed immutably by the launch while c_dev is borrowed mutably, and the DisjointSlice<f32> in the kernel signature is the device-side statement that no two threads will write the same element. The data race that is the traditional first lesson of GPU programming is here a borrow checkerborrow checker error.

Building it will take a moment the first time, because cargo oxide must build the codegen backend before it can build anything with it:

> root@machine ~/ $ cargo oxide run
Building rustc-codegen-cuda backend...
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 1m 06s
✓ Backend built: .../librustc_codegen_cuda.so
Detected GPU arch: sm_80 (via nvidia-smi)
    Finished `release` profile [optimized] target(s) in 22.04s
     Running `target/release/borscht-cuda`
[0.0, 3.0, 6.0, 9.0, 12.0]

Four numbers added on a hundred billion transistors. It is not much, but every matrix multiplication in §4.4 is that same skeleton with a better kernel inside it, and it is worth knowing two more subcommands before we go there, since we will lean on both: cargo oxide inspect prints the PTX that came out the far end of the pipeline, and cargo oxide pipeline prints every intermediate along the way, from Rust MIR through dialect-mir and LLVM IR to PTX. When a kernel is slow or wrong, that is where you look.

GPU MODE Lecture 3: Getting Started With CUDA for Python Programmers

GPU MODE Lecture 5: Going Further with CUDA for Python Programmers

GPU MODE Lecture 2: Ch1-3 PMPP book

GPU MODE Lecture 4: Compute and Memory Basics (Ch4-5 PMPP)

#![allow(unused)]
fn main() {
enum Device {
    Cpu,
    Gpu
}

enum Dtype {
    Float64(f64),
    Float32(f32),
    Float16(f16),
    Bfloat16(bf16),
}

struct Tensor {
    shape: Vec<u8>,
    stride: Vec<u8>,
    device: Device,
    dtype: Dtype,
    storage: Vec<Dtype>
}
}

4.4 Accelerating Matrix Multiplication on GPUs

$\hookleftarrow$ Table of Contents

In which we …

4.3.1 From Multi Core CPUs to Many Core GPUs

#![allow(unused)]
fn main() {
// gpu_host.rs
use cudarc::{driver::{self, PushKernelArg}, nvrtc};
use src_device::T; // shared type with device code
static PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/gpu_device.ptx")); // Embed the PTX code as a static string.

pub fn cudars_helloworld() -> Result<(), Box<dyn std::error::Error>> {
  // initialize device context and stream via driver api
  let process = driver::CudaContext::new(0)?; // device 0
  let queue = process.default_stream();
  
  // load ptx via nvrtc
  let dylib = process.load_module(nvrtc::Ptx::from_src(PTX))?;
  let add_kernel = dylib.load_function("add")?;

  // allocate on device
  let (a, b): ([T; _], [T; _]) = ([1.0, 2.0, 3.0, 4.0], [2.0, 3.0, 4.0, 4.0]);
  let (a_gpu, b_gpu, mut c_gpu) = (queue.clone_htod(&a)?, queue.clone_htod(&b)?, queue.alloc_zeros::<T>(a.len())?);
  let (a_len, b_len) = (a_gpu.len(), b_gpu.len());

  let cfg = driver::LaunchConfig { grid_dim: (1, 1, 1), block_dim: (4, 1, 1), shared_mem_bytes: 0, };
  unsafe {
    queue
    .launch_builder(&add_kernel).arg(&a_gpu).arg(&a_len).arg(&b_gpu).arg(&b_len).arg(&mut c_gpu)
    .launch(cfg)?;
  }
  queue.synchronize()?;

  let c = queue.clone_dtoh(&c_gpu)?;
  println!("c from cuda is = {:?}", c);
  Ok(())
}
}

#![allow(unused)]
fn main() {
// gpu_device.rs
use cuda_std::kernel;
use crate::T;

#[allow(improper_ctypes_definitions)]
#[kernel] pub unsafe fn add(a: &[T], b: &[T], c: *mut T) {
  let i = cuda_std::thread::index_1d() as usize;
  if i < a.len() {
    let elem = unsafe { &mut *c.add(i) };
    *elem = a[i] + b[i];
  }
}

#[allow(improper_ctypes_definitions)]
#[kernel] pub unsafe fn saxpy(a: &[T], b: &[T], c: *mut T) {
  let i = cuda_std::thread::index_1d() as usize;
  todo!()
}

#[allow(improper_ctypes_definitions)]
#[kernel] pub unsafe fn smul(a: &[T], b: &[T], c: *mut T) {
  let i = cuda_std::thread::index_1d() as usize;
  todo!()
}

#[allow(improper_ctypes_definitions)]
#[kernel] pub unsafe fn stanh(a: &[T], b: &[T], c: *mut T) {
  let i = cuda_std::thread::index_1d() as usize;
  todo!()
}
}

4.3.2 From CUDA Rust to PTX

4.3.3 Accelerating SGEMV on sm80 (Ampere)

4.3.4 Accelerating SGEMM on sm90 (Hopper)

4.3.5 Accelerating SGEMM on sm90 (Hopper)

4.3.6 Accelerating SGEMM on sm100 (Blackwell)

4.6 From Strides to Layouts

$\hookleftarrow$ Table of Contents

4.7 Accelerating Attention

$\hookleftarrow$ Table of ContentsTri DaoFlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022)