PyTorch from the Inside: Tensors, Autograd and Precision

Everything in the series so far was written in NumPy, with backward passes by hand. That was the point: every gradient was a rule we derived. From here on the models get too big for that, and the series switches to PyTorch. This post opens it up first, so that nothing it does later is magic. A deep-learning framework is three things, and most surprises, especially performance surprises, come from mixing them up:

• A tensor library. Arrays of numbers with a shape, a data type and a layout in memory, and fast operations on them.
• An autograd tape. A record of the operations that produced a result, so the framework can run Backpropagation, by Hand’s backward rules automatically.
• A device model. Where the numbers live, CPU or GPU, and the fact that work sent to a GPU runs later, asynchronously.

Table of contents

A tensor is a view of memory

A tensor is a block of memory plus a description of how to read it: a shape, a data type, an offset, and a stride for each dimension, the number of elements to jump to move one step along it.

import torch

x = torch.arange(12.0).reshape(3, 4)
print(f"x: shape {tuple(x.shape)}, strides {x.stride()}, contiguous {x.is_contiguous()}")
t = x.T                                            # a transpose is a view: new strides, same memory
print(f"x.T: shape {tuple(t.shape)}, strides {t.stride()}, contiguous {t.is_contiguous()}, "
      f"same memory {t.data_ptr() == x.data_ptr()}")
row, col = x[1], x[:, 1]                           # slicing is a view too, with an offset
print(f"x[1]: offset {row.storage_offset()}, strides {row.stride()};  "
      f"x[:, 1]: offset {col.storage_offset()}, strides {col.stride()}")
t[0, 1] = 100.0                                    # writing through a view writes the original
print(f"after t[0, 1] = 100: x[1, 0] = {x[1, 0].item()}")
c = t.contiguous()                                 # a real copy, laid out row by row
print(f"t.contiguous(): strides {c.stride()}, same memory {c.data_ptr() == x.data_ptr()}")
try:
    t.view(12)
except RuntimeError as e:
    print(f"t.view(12) fails: {str(e).split('.')[0]}")
print(f"t.reshape(12) works (it copies when it must): {t.reshape(12)[:4].tolist()}")

Full script: tensor_views.py

x: shape (3, 4), strides (4, 1), contiguous True
x.T: shape (4, 3), strides (1, 4), contiguous False, same memory True
x[1]: offset 4, strides (1,);  x[:, 1]: offset 1, strides (4,)
after t[0, 1] = 100: x[1, 0] = 100.0
t.contiguous(): strides (3, 1), same memory False
t.view(12) fails: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces)
t.reshape(12) works (it copies when it must): [0.0, 100.0, 8.0, 1.0]

The same twelve numbers in memory, read three ways (each dot’s number is the order in which the view visits that memory cell):

• A row-major $3 \times 4$ tensor has strides $(4, 1)$. Moving down a row skips 4 elements; moving along a row skips 1.
• A transpose copies nothing. x.T is the same memory with the strides swapped to $(1, 4)$. That’s why it’s instant for any size, and why writing through it changes x.
• A slice is a view with an offset. Row 1 starts 4 elements in; column 1 starts 1 element in and strides by 4.
• Some operations need the layout to be simple. view only reinterprets memory: it works whenever the new shape can be expressed with strides, and fails on the transposed tensor, whose elements aren’t in an order view(12) can describe. contiguous() makes a real copy in row order. reshape returns a view when it can and a copy when it must, so never write through a reshape result expecting the original to change.

(PyTorch’s strides count elements; NumPy’s .strides count bytes.)

This matters for speed later. A GPU reads memory fastest in long consecutive runs (How a GPU Runs Deep Learning), so a non-contiguous tensor can make the next operation slower, or force a hidden copy.

Autograd is a tape

Backpropagation, by Hand derived backpropagation as local rules applied from the loss backward. An autograd system automates exactly that: during the forward pass, record every operation and its inputs; during the backward pass, walk the record in reverse and apply each operation’s rule. Here is the whole idea in 55 lines, enough to run Backpropagation, by Hand’s XOR network:

import numpy as np

class Tensor:
    """A NumPy array that remembers how it was made, so it can run backward."""
    def __init__(self, data, parents=(), backward=None):
        self.data, self.grad = np.asarray(data, dtype=float), None
        self.parents, self._backward = parents, backward   # the tape: one node per operation

    def __add__(self, other):
        out = Tensor(self.data + other.data, (self, other))
        def backward(g):                                    # sum over broadcast dimensions
            return unbroadcast(g, self.data.shape), unbroadcast(g, other.data.shape)
        out._backward = backward
        return out

    def __matmul__(self, other):
        out = Tensor(self.data @ other.data, (self, other))
        out._backward = lambda g: (g @ other.data.T, self.data.T @ g)   # "Backpropagation, by Hand"'s rules
        return out

    def relu(self):
        out = Tensor(np.maximum(self.data, 0), (self,))
        out._backward = lambda g: (g * (self.data > 0),)
        return out

    def bce_with_logits(self, y):
        """Mean log-loss of sigmoid(self) against labels y, as one operation.
        (Not stable for very large logits; PyTorch's uses log-sum-exp, as in "Cross-Entropy and the Softmax Gradient".)"""
        p = 1 / (1 + np.exp(-self.data))
        out = Tensor(-(y * np.log(p) + (1 - y) * np.log(1 - p)).mean(), (self,))
        out._backward = lambda g: (g * (p - y) / y.size,)             # p - y, "Learning Is Optimization"
        return out

    def backward(self):
        order, seen = [], set()
        def visit(t):                                        # topological order of the tape
            if id(t) not in seen:
                seen.add(id(t))
                for p in t.parents:
                    visit(p)
                order.append(t)
        visit(self)
        self.grad = np.ones_like(self.data)
        for t in reversed(order):                            # from the loss back to the inputs
            if t._backward:
                for p, g in zip(t.parents, t._backward(t.grad)):
                    p.grad = g if p.grad is None else p.grad + g   # a tensor used twice sums its gradients

def unbroadcast(g, shape):
    while g.ndim > len(shape):                               # dimensions added in front
        g = g.sum(0)
    for i, s in enumerate(shape):                            # dimensions stretched from size 1
        if s == 1:
            g = g.sum(i, keepdims=True)
    return g

Full script: tiny_autograd.py

Each operation creates a new Tensor that remembers its parents and a function mapping the gradient arriving at the output to the gradients for its inputs. backward puts every tensor in an order where each comes after its inputs, then runs the functions in reverse. Two details carry most of the weight:

• Gradients add up. A tensor used in two places gets a gradient from each, and they sum. That’s the chain rule for a quantity with several paths to the loss (Backpropagation, by Hand).
• Broadcasting has a backward too. Adding a bias vector to every row of a matrix uses the bias once per row, so its gradient is the sum over rows: unbroadcast.

check_autograd.py runs it on Backpropagation, by Hand’s worked step and compares with PyTorch, then repeats the comparison on 200 random networks. The worked step:

import numpy as np
import torch
from tiny_autograd import Tensor

# "Backpropagation, by Hand"'s worked example: one XOR input through the 2-2-1 network.
x, y = np.array([[1.0, 0.0]]), np.array([[1.0]])
W1, b1 = np.array([[0.5, -0.4], [0.3, 0.8]]), np.array([0.0, 0.1])
w2, b2 = np.array([[0.7], [-0.6]]), np.array([0.1])

ours = [Tensor(a) for a in (W1, b1, w2, b2)]
loss = ((Tensor(x) @ ours[0] + ours[1]).relu() @ ours[2] + ours[3]).bce_with_logits(y)
loss.backward()

theirs = [torch.tensor(a, requires_grad=True) for a in (W1, b1, w2, b2)]
h = torch.relu(torch.tensor(x) @ theirs[0] + theirs[1])
ref = torch.nn.functional.binary_cross_entropy_with_logits(h @ theirs[2] + theirs[3], torch.tensor(y))
ref.backward()

print(f"loss: ours {loss.data:.4f}, PyTorch {ref.item():.4f}")
for name, a, b in zip(["W1", "b1", "w2", "b2"], ours, theirs):
    print(f"dL/d{name}: {np.round(a.grad, 4).ravel().tolist()}   largest gap to PyTorch "
          f"{np.abs(a.grad - b.grad.numpy()).max():.1e}")

Full script: check_autograd.py

loss: ours 0.4932, PyTorch 0.4932
dL/dW1: [-0.2726, 0.0, 0.0, 0.0]   largest gap to PyTorch 0.0e+00
dL/db1: [-0.2726, 0.0]   largest gap to PyTorch 0.0e+00
dL/dw2: [-0.1947, 0.0]   largest gap to PyTorch 0.0e+00
dL/db2: [-0.3894]   largest gap to PyTorch 0.0e+00
200 random cases, shared weight and (1, 5) bias: largest gap to PyTorch 1.1e-16

The same numbers Backpropagation, by Hand derived by hand ($-0.2726$, $-0.1947$, $-0.3894$), and identical to PyTorch’s here. On 200 random networks, with a weight used twice and a bias broadcast from shape $(1, 5)$, the two agree to about $10^{-16}$, the last bit or two: the same arithmetic done in a slightly different order. (The one thing this tape doesn’t do that PyTorch’s does is free itself: PyTorch discards the graph after backward unless it is called with retain_graph=True, and walks it with a work queue instead of recursion, which a very deep graph would overflow.)

The graph PyTorch records

PyTorch keeps the same kind of record, one node per operation, and the record can be walked:

import torch

x = torch.tensor([[1.0, 0.0]])
W1 = torch.tensor([[0.5, -0.4], [0.3, 0.8]], requires_grad=True)
w2 = torch.tensor([[0.7], [-0.6]], requires_grad=True)
z = torch.relu(x @ W1) @ w2

# The graph PyTorch recorded while computing z, walked from the output back.
def walk(fn, depth=0):
    if fn is None:
        return
    print("  " * depth + type(fn).__name__)
    for nxt, _ in fn.next_functions:
        walk(nxt, depth + 1)
walk(z.grad_fn)

# The graph is rebuilt on every forward pass, so ordinary Python control flow
# decides its shape: here, how many times a layer is applied.
def count(fn, seen=None):
    seen = set() if seen is None else seen
    if fn is None or fn in seen:
        return 0
    seen.add(fn)
    return 1 + sum(count(nxt, seen) for nxt, _ in fn.next_functions)

W = torch.eye(2, requires_grad=True)
for repeats in (1, 3):
    h = x
    for _ in range(repeats):
        h = torch.relu(h @ W)
    print(f"layer applied {repeats}x: {count(h.grad_fn)} nodes in the graph")

# Under no_grad nothing is recorded: no graph, and no saved activations.
with torch.no_grad():
    print(f"under no_grad: grad_fn = {torch.relu(x @ W1).grad_fn}")

Full script: graph.py

MmBackward0
  ReluBackward0
    MmBackward0
      AccumulateGrad
  AccumulateGrad
layer applied 1x: 3 nodes in the graph
layer applied 3x: 7 nodes in the graph
under no_grad: grad_fn = None

• The graph mirrors the forward pass, backward. From the output: a matrix multiply, a ReLU, another matrix multiply, and an AccumulateGrad node for each weight, which is where .grad gets written.
• It’s rebuilt on every forward pass. Applying a layer three times in a Python loop gives a graph with three copies of the layer’s nodes (7 instead of 3; the shared weight has one accumulator). Ordinary Python control flow, if and for, decides the graph’s shape. That’s what “define-by-run” or eager execution means.
• torch.no_grad() turns recording off. No graph is built, and nothing is saved for a backward pass. For inference, torch.inference_mode() is a stricter, faster version. Neither switches dropout or batch statistics to their test-time behavior; that’s model.eval() (Dropout and Normalization).

Two practical consequences follow from the tape.

• Gradients accumulate. .grad is added to, not overwritten, which is what the summing rule needs inside one backward pass. Between steps it’s usually a bug, which is why every training loop calls zero_grad() (which by default sets the gradients to None). Sometimes it’s the point: gradient accumulation runs several small batches, calls backward on each, and steps once, to train with a batch too big for memory.
• The tape holds memory. To run Backpropagation, by Hand’s rules, the backward pass needs what each operation saved: its inputs for a matrix multiply, its output for a ReLU or a sigmoid. The graph keeps them alive until backward runs. That’s the activation memory Backpropagation, by Hand described. no_grad avoids it for inference, and activation checkpointing, torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False), trades it for recomputation during training.

Devices, and why GPU timings lie

Every timing below was measured on the machine this post was written on, an Apple M4 laptop, using PyTorch 2.11 on its CPU and on its GPU through Apple’s Metal backend (mps). The data-center numbers in AI Engineering, from the Inside are from an H100’s datasheet; the lessons transfer, the absolute numbers don’t.

A tensor lives on a device, and operations run where their inputs are. Moving data between devices is an explicit copy (.to("mps"), .cpu()), and it’s slow relative to computing, so the rule is to move once and stay.

The less obvious part is timing. A GPU call returns as soon as the work is queued, not when it’s done. That’s what keeps the GPU busy while Python prepares the next operation, and it means a naive timer measures the wrong thing. gpu_sync below waits for whichever GPU the script finds, torch.mps.synchronize on this laptop and torch.cuda.synchronize on an NVIDIA card. The excerpt is the timer and the first two measurements; the script’s last part times a small multiply waiting after every call, and again with 200 calls queued.

DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
gpu_sync = {"cuda": torch.cuda.synchronize, "mps": torch.mps.synchronize}.get(DEVICE, lambda: None)   # the post ran on "mps"

def timed(fn, device, trials=10):
    """Median seconds per call. GPU work is queued, not finished, when a call returns,
    so wait for the device before reading the clock."""
    sync = gpu_sync if device == DEVICE else (lambda: None)
    fn(); sync()                                        # warm up
    times = []
    for _ in range(trials):
        t0 = time.perf_counter(); fn(); sync()
        times.append(time.perf_counter() - t0)
    return sorted(times)[trials // 2]

cases = [("cpu", torch.float32), (DEVICE, torch.float32), (DEVICE, torch.float16), (DEVICE, torch.bfloat16)]
for n in (256, 1024, 4096):
    row = [f"n = {n:4d}:"]
    for device, dt in cases:
        a = torch.randn(n, n, device=device, dtype=dt)
        b = torch.randn(n, n, device=device, dtype=dt)
        s = timed(lambda: a @ b, device)
        row.append(f"{device} {str(dt).split('.')[1]:8s} {2 * n ** 3 / s / 1e12:5.2f}")
    print("  ".join(row) + "  TFLOP/s", flush=True)

# The same GPU matmul, timed without waiting for it: the clock only sees the launch.
a = torch.randn(4096, 4096, device=DEVICE)
a @ a; gpu_sync()
t0 = time.perf_counter(); a @ a; t_launch = time.perf_counter() - t0
gpu_sync()
print(f"4096 matmul on {DEVICE}, no synchronize: {t_launch * 1e3:.2f} ms 'measured', "
      f"{timed(lambda: a @ a, DEVICE) * 1e3:.2f} ms real")

Full script: timing.py

n =  256:  cpu float32   1.36  mps float32   0.12  mps float16   0.14  mps bfloat16  0.13  TFLOP/s
n = 1024:  cpu float32   1.88  mps float32   1.52  mps float16   2.38  mps bfloat16  2.46  TFLOP/s
n = 4096:  cpu float32   1.45  mps float32   2.99  mps float16   3.43  mps bfloat16  3.42  TFLOP/s
4096 matmul on mps, no synchronize: 0.05 ms 'measured', 47.58 ms real
256 matmul on mps: 230 us per call waiting each time, 30 us per call when 200 are queued (1.12 TFLOP/s)

• An unsynchronized timer is wrong by orders of magnitude. The 4096 matrix multiply “took” 0.05 ms when the clock stopped at the call’s return, and 47.58 ms when it waited for the result, about 1,000 times longer here. The first number is the time to queue the work. A benchmark needs a synchronize before each clock reading: torch.cuda.synchronize() on NVIDIA, torch.mps.synchronize() here. (torch.utils.benchmark.Timer handles this and the warm-up.)
• But don’t synchronize inside the work. Timed one call at a time, a $256 \times 256$ multiply takes 230 µs on this GPU, 0.15 TFLOP/s, because every measurement includes a full round trip to the GPU and back. Queue 200 of them and synchronize once, and each takes 30 µs, 1.1 TFLOP/s: the GPU overlaps the calls. Real training queues work the same way, which is why anything that forces a wait in the loop, .item(), .cpu(), printing a tensor, branching on a tensor’s value, costs more than it looks.
• Small work still doesn’t pay for the GPU. Even queued, the $256 \times 256$ multiply is slower on this GPU than on this laptop’s CPU (1.4 TFLOP/s). At $4096 \times 4096$ the GPU is ahead, 3.4 TFLOP/s in BF16 against the CPU’s 1.5 in FP32. This CPU is unusually fast at matrix math because Apple’s chips have matrix units that PyTorch reaches through Apple’s Accelerate library; a typical server CPU is much further behind a data-center GPU.
• The habits transfer, the speeds don’t. The H100 in AI Engineering, from the Inside does 989 TFLOP/s of dense BF16, about 300 times this laptop’s GPU, and its BF16 tensor cores are about 15 times its own FP32 rate (989 against 67). Here BF16 is only about 15% faster than FP32. (On an H100, PyTorch’s torch.set_float32_matmul_precision("high") runs FP32 matrix multiplies in TF32 on the tensor cores, 494 TFLOP/s dense, which narrows that gap to 2×.) The habits transfer, move data once, batch work into big operations, synchronize before timing, but absolute speeds and precision trade-offs don’t.

Numbers, bit by bit

A floating-point number stores a sign, an exponent, which sets its range, and a mantissa, which sets its precision. The three formats deep learning uses split their bits differently:

import torch

bits = {torch.float32: (8, 23), torch.float16: (5, 10), torch.bfloat16: (8, 7)}
print(f"{'':9s} {'exponent':>8s} {'mantissa':>8s} {'largest':>10s} {'smallest normal':>16s} "
      f"{'smallest subnormal':>19s} {'step at 1.0':>12s}")
for dt, (e, m) in bits.items():
    f = torch.finfo(dt)                                   # a subnormal is the normal minimum's last bits
    print(f"{str(dt).split('.')[1]:9s} {e:8d} {m:8d} {f.max:10.3g} {f.tiny:16.3g} "
          f"{f.tiny * f.eps:19.3g} {f.eps:12.3g}")

show = lambda v, dt: torch.tensor(v, dtype=torch.float32).to(dt).item()
print()
for v in (1e-8, 70_000.0, 1.001):
    print(f"{v:>9g} stored as  float16: {show(v, torch.float16):<12.6g} bfloat16: {show(v, torch.bfloat16):.6g}")

# Adding many small updates to one number, as an optimizer does to a weight.
for dt in (torch.float32, torch.bfloat16):
    w = torch.tensor(1.0, dtype=dt)
    for _ in range(10_000):
        w = w + torch.tensor(1e-4, dtype=dt)
    print(f"1 + 10,000 steps of 0.0001 in {str(dt).split('.')[1]:8s} = {w.item():.4f} (should be 2.0000)")

# FP16's narrow range: small gradients vanish unless the loss is scaled up first.
g = torch.tensor([1e-8, 2e-7, 1e-4])
fmt = lambda t: "[" + ", ".join(f"{v:.3g}" for v in t.tolist()) + "]"
print(f"gradients {fmt(g)} in float16: {fmt(g.half().float())}")
print(f"scaled by 1024, stored, unscaled: {fmt((g * 1024).half().float() / 1024)}")

Full script: float_formats.py

          exponent mantissa    largest  smallest normal  smallest subnormal  step at 1.0
float32          8       23    3.4e+38         1.18e-38             1.4e-45     1.19e-07
float16          5       10   6.55e+04          6.1e-05            5.96e-08     0.000977
bfloat16         8        7   3.39e+38         1.18e-38            9.18e-41      0.00781

    1e-08 stored as  float16: 0            bfloat16: 1.00117e-08
    70000 stored as  float16: inf          bfloat16: 70144
    1.001 stored as  float16: 1.00098      bfloat16: 1
1 + 10,000 steps of 0.0001 in float32  = 2.0002 (should be 2.0000)
1 + 10,000 steps of 0.0001 in bfloat16 = 1.0000 (should be 2.0000)
gradients [1e-08, 2e-07, 0.0001] in float16: [0, 1.79e-07, 0.0001]
scaled by 1024, stored, unscaled: [1e-08, 2e-07, 0.0001]

• FP32 has 8 exponent bits and 23 mantissa bits: a range up to $3.4 \times 10^{38}$ and a step of about $10^{-7}$ near 1.
• FP16 has 5 exponent bits and 10 mantissa bits, three more than BF16. Its range stops at 65,504, so 70,000 becomes infinity. Below its smallest normal number, $6.1 \times 10^{-5}$, it keeps going with subnormal numbers, which give up precision bit by bit, down to $6 \times 10^{-8}$. So $2 \times 10^{-7}$ comes back 10% off, and $10^{-8}$ is past the end: it becomes 0.
• BF16 (“brain float”) keeps FP32’s 8 exponent bits and cuts the mantissa to 7. Its range is FP32’s, so $10^{-8}$ is fine and 70,000 is in range, stored as 70,144. Its precision is poor: near 1, the next number is 1.0078, so 1.001 is stored as 1.

Why BF16 won for training. Gradients span a huge range, and many are tiny. In FP16 the small ones fall into the subnormals or underflow to zero, as the last lines show. The fix is loss scaling: multiply the loss by a large factor, so by the chain rule every gradient is that much larger and out of the danger zone, then divide it back out in FP32 before the update (Micikevicius et al., 2017). The demo scales the gradients directly, which is the same thing. PyTorch’s torch.amp.GradScaler does it adaptively: it starts the scale at $2^{16}$, skips any step whose gradients overflow and halves the scale, and doubles it again after 2,000 clean steps. (Unscale before clipping gradients, with scaler.unscale_(opt).) FP16’s 65,504 ceiling can also overflow in the forward pass, in large activations or attention scores, which loss scaling can’t fix. BF16 has FP32’s range, so it needs no loss scaling at all. And its low precision matters little for matrix multiplications because the hardware multiplies BF16 inputs but adds up the products in FP32, and autocast keeps numerically sensitive operations, softmax, norms, the loss, in FP32 (Kalamkar et al., 2019). BF16 spread with the hardware that supports it, first Google’s TPUs, then NVIDIA’s GPUs from the A100 on. FP16 is still common for inference, and newer GPUs train with 8-bit formats for the matrix math, still with FP32 weights and higher-precision accumulation (Quantization from Scratch covers FP8).

Why the weights stay in FP32 anyway. The accumulation line is the reason. Adding $10^{-4}$ to 1, ten thousand times, should give 2. In BF16 it gives 1: each addition is smaller than half the gap between 1 and the next representable number, so it rounds away, every time. (Even FP32 drifts a little, to 2.0002.) An optimizer does exactly this. Each update is typically $10^{-3}$ to $10^{-5}$ of the weight it changes, below BF16’s half-step of about 0.4%, so in BF16 most updates would round away. So the weights and the optimizer’s state are kept in FP32, and only the heavy math runs in BF16. That’s mixed precision. The same reason is why the matrix hardware accumulates in FP32: a long sum in BF16 would stall the same way.

The router, in PyTorch

Here is Cross-Entropy and the Softmax Gradient’s four-way router trained the way Initialization and Optimizers trained it, with PyTorch doing the work, and Initialization and Optimizers’s hand-written optim.py running alongside on the same mini-batches. router_torch.py loads the data with those posts’ helpers, copied into its folder, then trains both, and finally prints the PyTorch model’s validation accuracy:

# The four-way router as a PyTorch module, trained the way "Initialization and Optimizers" trained it.
model = torch.nn.Linear(X.shape[1], 4, dtype=torch.float64)
torch.nn.init.zeros_(model.weight); torch.nn.init.zeros_(model.bias)
opt = torch.optim.AdamW([{"params": [model.weight], "weight_decay": 0.3},
                         {"params": [model.bias], "weight_decay": 0.0}],   # no decay on biases
                        lr=0.01)

# The same run with "Initialization and Optimizers"'s hand-written AdamW, on the same mini-batches.
W, b, ours = np.zeros((Xtr.shape[1], 4)), np.zeros(4), Adam(0.01, decay=0.3)

rng = np.random.default_rng(0)
for step in range(1, 3001):
    i = rng.choice(len(ytr), 16, replace=False)
    loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])   # takes logits ("Cross-Entropy and the Softmax Gradient")
    opt.zero_grad()                                               # gradients accumulate otherwise
    loss.backward()
    opt.step()
    G = (softmax(Xtr[i] @ W + b) - np.eye(4)[ytr[i]]) / 16
    ours.step(W, Xtr[i].T @ G, "W"); ours.step(b, G.sum(0), "b")
    if step in (1, 10, 3000):
        gap = np.abs(model.weight.detach().numpy().T - W).max()
        print(f"step {step:4d}: loss {loss.item():.4f}, largest gap between PyTorch and optim.py weights {gap:.1e}")

Full script: router_torch.py

step    1: loss 1.3863, largest gap between PyTorch and optim.py weights 1.7e-18
step   10: loss 1.1757, largest gap between PyTorch and optim.py weights 2.8e-17
step 3000: loss 0.0428, largest gap between PyTorch and optim.py weights 1.1e-15
validation accuracy 0.750

The standard training loop is four lines: compute the loss, zero_grad, backward, step. Two details are worth noticing.

• Weight decay by parameter group. AdamW takes groups of parameters with their own settings; here the bias gets no decay, as Initialization and Optimizers recommended.
• The hand-written version matches. After 3,000 steps the two sets of weights differ by $10^{-15}$, which is floating-point rounding. Initialization and Optimizers’s optim.py is PyTorch’s AdamW, and PyTorch’s AdamW is Initialization and Optimizers’s equations.

What a training step keeps in memory

Initialization and Optimizers counted 16 bytes per parameter for mixed-precision Adam, using the ZeRO paper’s accounting: 16-bit weights and gradients, plus an FP32 master copy and Adam’s two FP32 averages. PyTorch’s usual mixed precision, torch.autocast, arrives at the same total differently. Measured:

import torch

# Count the bytes a training step actually keeps, for a small MLP trained
# with AdamW, the way mixed precision usually runs in PyTorch: FP32 weights,
# BF16 math inside autocast.
model = torch.nn.Sequential(torch.nn.Linear(1024, 4096), torch.nn.ReLU(), torch.nn.Linear(4096, 1024))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
x = torch.randn(64, 1024)
saved = []                                              # what autograd keeps for backward
weights = {p.shape for p in model.parameters() if p.dim() == 2}
with torch.autograd.graph.saved_tensors_hooks(lambda t: saved.append(t) or t, lambda t: t):
    with torch.autocast("cpu", dtype=torch.bfloat16):
        loss = model(x).float().pow(2).mean()
loss.backward()
opt.step()

size = lambda ts: sum(t.numel() * t.element_size() for t in ts)
dtype = lambda ts: str(next(iter(ts)).dtype).split(".")[1]
n = sum(p.numel() for p in model.parameters())
state = [t for s in opt.state.values() for t in s.values() if t.dim() > 0]
parts = {"weights": list(model.parameters()),
         "gradients": [p.grad for p in model.parameters()],
         "Adam m and v": state}
for name, ts in parts.items():
    print(f"{name:13s} {size(ts) / n:4.1f} bytes/param  ({dtype(ts)})")
print(f"{'total':13s} {sum(size(ts) for ts in parts.values()) / n:4.1f} bytes/param for {n:,} parameters")
copies = [t for t in saved if t.dtype == torch.bfloat16 and t.shape in weights]
print(f"during the step: {len(copies)} BF16 weight copies saved for backward, "
      f"{size(copies) / n:.1f} more bytes/param")

Full script: step_memory.py

weights        4.0 bytes/param  (float32)
gradients      4.0 bytes/param  (float32)
Adam m and v   8.0 bytes/param  (float32)
total         16.0 bytes/param for 8,393,728 parameters
during the step: 1 BF16 weight copies saved for backward, 1.0 more bytes/param

Under autocast the FP32 weights are the master copy. Autocast casts each one to BF16 when a matrix multiply first uses it, reuses that cast for the rest of the forward pass, and the multiply keeps the BF16 copy for backward. The gradients flow back through the cast and land in FP32. What persists between steps is $4 + 4 + 8 = 16$ bytes per parameter, the same total as ZeRO’s accounting (which uses FP16, the same size as BF16), split differently. During the step, the saved BF16 copies add up to 2 more bytes per parameter in a deep network. Here it’s 1, because the first layer’s input needs no gradient, so its weight’s copy isn’t saved. The split matters at scale: FP32 gradients mean a data-parallel all-reduce moves 4 bytes per parameter instead of 2, which is why sharded setups keep 16-bit gradients.

Graphs, eager, and compilation

Eager execution, running each operation as Python reaches it, is easy to debug and lets Python control flow shape the model. It has a cost: each operation is launched separately, reads its inputs from memory and writes its output back, and the framework can’t see the next operation coming. Dropout and Normalization showed that simple elementwise work is bound by memory traffic, so a chain of four such operations pays for four round trips.

TensorFlow started from the other end (Abadi et al., 2016). In its first version a program built a static graph first and ran it afterwards, which let TensorFlow optimize and plan the whole graph (and later compile it with XLA), but made debugging and dynamic models awkward. TensorFlow 2 made eager the default and kept graphs behind tf.function. PyTorch started eager (Paszke et al., 2019), following Chainer, which coined “define-by-run”, and added compilation later: torch.compile (Ansel et al., 2024) watches the Python code run, captures the tensor operations into a graph, compiles the backward pass too, and generates fused kernels, Triton on NVIDIA GPUs and C++ on CPUs. Code it can’t capture falls back to eager at a “graph break”, and a new input shape can trigger a recompile, so the compilation cost is paid once only if shapes stay stable. mode="reduce-overhead" also replays whole sequences of GPU calls at once (CUDA graphs), which attacks the per-call cost measured above. The two frameworks converged on the same design from opposite ends.

What fusion buys, on the CPU:

import time
import torch

def gelu_scale_shift(x, g, b):
    """Four elementwise steps. Eager PyTorch runs each as its own pass over memory."""
    return torch.nn.functional.gelu(x) * g + b + x

def scale_shift(x, g, b):
    """The same without GELU: three steps, almost no arithmetic."""
    return x * g + b + x

x = torch.randn(4096, 4096)
g, b = torch.randn(4096), torch.randn(4096)

def timed(fn, trials=20):
    fn(); times = []
    for _ in range(trials):
        t0 = time.perf_counter(); fn(); times.append(time.perf_counter() - t0)
    return sorted(times)[trials // 2]

for name, fn in [("GELU, scale, shift, add", gelu_scale_shift), ("scale, shift, add", scale_shift)]:
    t0 = time.perf_counter()
    fast = torch.compile(fn)
    fast(x, g, b)
    compile_s = time.perf_counter() - t0
    eager_s, fused_s = timed(lambda: fn(x, g, b)), timed(lambda: fast(x, g, b))
    same = torch.allclose(fast(x, g, b), fn(x, g, b), atol=1e-5)
    print(f"{name:24s} eager {eager_s * 1e3:5.1f} ms   compiled {fused_s * 1e3:4.1f} ms "
          f"({eager_s / fused_s:.1f}x)   first call {compile_s:.1f} s   same result: {same}")
t = timed(lambda: x.clone())
print(f"{'one read + one write':24s} {t * 1e3:5.1f} ms  (x.clone(), the floor for any one-pass kernel)")

Full script: fusion.py

GELU, scale, shift, add  eager   9.7 ms   compiled  4.0 ms (2.4x)   first call 1.6 s   same result: True
scale, shift, add        eager   6.7 ms   compiled  1.8 ms (3.7x)   first call 0.0 s   same result: True
one read + one write       1.6 ms  (x.clone(), the floor for any one-pass kernel)

In this run the two compiled functions are 2.4 and 3.7 times faster than eager, with the same results. Count the memory traffic to see why, and why the two chains differ.

• Eager runs each step as its own pass. For the GELU chain, GELU reads x and writes a temporary, the multiply reads that and writes another, and so on: 9 full-array reads and writes. The compiled kernel reads the inputs once and writes the result once, 2 transfers.
• Without GELU, fusion reaches the floor. The compiled three-step chain takes 1.8 ms, close to the 1.6 ms of x.clone(), which does nothing but read the array once and write it once. It’s purely memory-bound, and fusion removed every pass it could.
• With GELU, the math becomes the limit. The compiled GELU chain takes about two and a half times that floor. Once the passes are fused, the kernel is no longer waiting on memory: GELU’s arithmetic, an erf per element, is now the bottleneck on this CPU.

That’s Dropout and Normalization’s lesson from the other side: for memory-bound work, fewer passes is the whole game, until the work stops being memory-bound. The price is the first call, which spent 1.6 seconds compiling. Most of that appears to be the compiler starting up, a one-time cost: the second function’s first call, compiled right after, took under 0.05 seconds (printed as 0.0 s). For a training run of thousands of steps with stable shapes that’s nothing; for a function called once it’s a loss.

Where this leads

• How a GPU Runs Deep Learning goes inside the GPU: what a kernel is, why reading memory in long consecutive runs matters, the roofline that decides whether an operation is bound by math or by memory, and a first kernel written by hand.
• Build a Tiny GPT, End to End, and Train It on This Blog builds a Transformer in PyTorch, with views to reshape the activations into heads and autograd for training.
• Quantization from Scratch comes back to precision from the other side: storing a trained model in 8 or 4 bits for serving.

Takeaways

• A framework is a tensor library, an autograd tape and a device model.
• A tensor is memory plus shape, strides and an offset. Transposes and slices are views that copy nothing; contiguous() copies.
• Autograd records each operation during the forward pass and applies its backward rule in reverse order. Gradients from different uses of a tensor add up, and broadcast inputs get summed gradients. A 55-line version matches PyTorch to about $10^{-16}$.
• PyTorch rebuilds the graph on every forward pass, so Python control flow shapes it. no_grad skips it; zero_grad is needed because gradients accumulate.
• GPU calls return before the work is done. Synchronize before reading a clock, but not inside the work: queued calls overlap, and every hidden synchronize (.item(), .cpu()) costs a round trip.
• BF16 has FP32’s range and little precision, so it needs no loss scaling. FP16 needs loss scaling to keep small gradients. Weights and optimizer state stay in FP32 because small updates vanish in BF16.
• Mixed-precision training with Adam keeps 16 bytes per parameter, whether counted ZeRO’s way or measured under torch.autocast.
• torch.compile captures eager code into a graph and fuses it, which pays off for chains of memory-bound operations.

References

Start with the PyTorch paper (Paszke et al.) for the design, then Ansel et al. for torch.compile. Micikevicius et al. is the reference for mixed precision, and CS336’s lecture 2 for accounting memory and FLOPs.

• Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library (NeurIPS, 2019): eager execution and PyTorch’s autograd.
• Abadi et al., TensorFlow: A System for Large-Scale Machine Learning (OSDI, 2016): the static-graph design PyTorch is contrasted with.
• Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation (ASPLOS, 2024): how torch.compile captures and fuses graphs.
• Micikevicius et al., Mixed Precision Training (ICLR, 2018): FP16 training, FP32 master weights and loss scaling.
• Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training (2019): BF16 as a drop-in for FP32 training.
• Stanford CS336, Language Modeling from Scratch, lecture 2 (2025): resource accounting, tensors, FLOPs and memory for training.