DL S014 · PyTorch Tensors and Autograd

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 14 / 80 · Module M03 · ~2 hours

The story — Soumith's confession, and the graph that builds itself

At NeurIPS 2016, Soumith Chintala walked on stage with a slide that had one word on it: "Autograd". PyTorch had shipped three months earlier. The room was polite. TensorFlow was already the default; every job posting said "experience with TensorFlow". Soumith's pitch was heretical: what if you didn't declare a graph at all — what if the graph built itself, one Python line at a time? The room applauded but nobody quite believed it. Six months later Karpathy tweeted a screenshot of his research code and captioned it "I have switched". By 2019, Facebook, OpenAI, Uber, and half of DeepMind were on PyTorch. By 2024, over 92% of new papers on Papers-with-Code shipped PyTorch first (source: PWC framework survey, Jan 2025). The dynamic graph won.

That's not a marketing story. It's a story about what removes friction from your brain. Static-graph frameworks (TF 1.x, Theano) asked you to first build a symbolic description of your model, compile it, and then feed it data. Debugging meant reading autogenerated C++ stack traces. Print statements didn't work. Dynamic-graph frameworks let you drop print(x) mid-forward-pass and get a number. That is the entire difference. Everything downstream — nanoGPT existing, HuggingFace being HuggingFace, you sitting here reading this — flows from that one design decision.

For thirteen sessions we've been doing something a little masochistic: we wrote a neural network with NumPy. Every matmul was ours. Every gradient — every single dL/dW1, dL/db1, dL/dW2 — we derived on paper and typed by hand. That's how you understand deep learning. But it's not how you do deep learning past about 20 lines of model code.

Here's the number that matters: the MNIST MLP from Session 011 was ~120 lines. A CNN would be ~400. A transformer with all the residual and norm plumbing, done by hand? A few thousand — and every one of them a place where the chain rule can go wrong. Nobody would ship anything.

PyTorch is what we reach for once we've paid the tuition. It gives us two superpowers: (1) tensors that live on the GPU with the same NumPy-ish API you already know, and (2) an autograd engine that records every operation into a computation graph and, on .backward(), walks the graph in reverse and fills in every gradient for you. Session 010 was you being autograd, by hand, on paper. Today, PyTorch does it in one line — and if you understood Session 010, that one line will not feel like magic. It'll feel like relief.

Prerequisites

• Session 010 — you derived backprop for an MLP by hand. Today that becomes .backward().
• Session 011 — the 120-line NumPy MLP is the thing we're rewriting.
• Session 013 — tensor-op fluency; the API is 90% the same as NumPy.

A note on which PyTorch you're running

Everything in this session works on PyTorch 2.0+, but the war stories at the end and the torch.compile discussion assume PyTorch 2.5 (Oct 2024) or 2.6 (Jan 2025) — the versions that ship TorchDynamo as default-stable, FlexAttention, and the redesigned memory profiler. Check with torch.__version__; if you're below 2.5, pip install --upgrade torch before starting. Colab defaults to 2.5+ as of early 2025.

1 · A tensor is a NumPy array that remembers

The single most important idea in PyTorch, stated as bluntly as I can:

A PyTorch tensor is a NumPy array that (a) can live on the GPU and (b) can remember every operation performed on it so we can compute derivatives later.

That's it. Everything else — nn.Module, DataLoader, torch.optim, torch.compile, the entire ecosystem — is built on top of this one primitive.

If you closed your eyes you'd think that was NumPy. Good. Now the twist:

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x * x).sum()
y.backward()
print(x.grad)                            # tensor([2., 4., 6.])

We asked for dy/dx and got it. No paper. No chain rule by hand. y = x0² + x1² + x2², so dy/dx = [2x0, 2x1, 2x2] = [2, 4, 6]. Correct.

That's autograd. The rest of this session is unpacking exactly what happened between line 2 and line 3.

A tiny bit of history: from Wengert lists to Theano to autograd

Reverse-mode automatic differentiation was invented in 1970 by Seppo Linnainmaa (Master's thesis, University of Helsinki) — years before backprop was rediscovered by Rumelhart/Hinton/Williams in 1986. Linnainmaa's key data structure was a Wengert list: a linearised record of every arithmetic operation, walked backwards to compute derivatives via the chain rule. That's exactly what PyTorch's autograd is. Every framework you've heard of — Theano (2007), Torch7 autograd (2015), Chainer (2015, the first dynamic-graph framework in Python and PyTorch's direct spiritual ancestor), TensorFlow eager (2017), JAX (2018) — is a re-implementation of Linnainmaa's Wengert list with better ergonomics. When you write y.backward(), you are running an algorithm from 1970 on a GPU chip from 2024. That's kind of beautiful.

Further reading: Baydin, Pearlmutter, Radul, Siskind — "Automatic Differentiation in Machine Learning: a Survey" (2018, still the definitive intro). Also: PyTorch autograd internals doc by Alban Desmaison — the actual engineer who maintains it.

2 · requires_grad=True — what the flag actually does

Setting requires_grad=True on a tensor tells PyTorch: "any operation that uses this tensor should be recorded, so I can later ask for d(anything)/d(this)". From that moment on, every op involving x builds a node in a hidden data structure called the computation graph.

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2                # y = x²,   also gets grad tracking
z = y + 1                 # z = y + 1
print(z)                  # tensor(10., grad_fn=<AddBackward0>)

Notice the grad_fn=<AddBackward0> in the printout. That's PyTorch saying "I remember that z was made by an Add op — if you ask me for gradients later, I know which backward function to call".

2.1 Why must the loss be a scalar?

Because a derivative ∂L/∂w is defined as "how much does one number change when I nudge one weight". If L is a vector or matrix, there's no single number to differentiate. You'd need a Jacobian, and PyTorch (rightly) refuses to build a (len(L), num_params)-shape monster unless you explicitly ask for it.

That's why every training loop ends with loss = criterion(pred, y) (a scalar) followed by loss.backward(). If you accidentally do loss.sum().backward() when loss was already a scalar, no harm; if you do loss.backward() when loss is a vector, you'll see:

RuntimeError: grad can be implicitly created only for scalar outputs

Fix: aggregate first. Usually .mean() (matches typical loss reductions).

3 · The dynamic computation graph, drawn

Let's do the smallest non-trivial example possible: L = (w * x + b - y)², a single linear neuron's squared error.

x = torch.tensor(2.0)                                # input, no grad
y = torch.tensor(5.0)                                # target, no grad
w = torch.tensor(3.0, requires_grad=True)            # leaf
b = torch.tensor(1.0, requires_grad=True)            # leaf
 
pred  = w * x + b                                    # 7.0
err   = pred - y                                     # 2.0
loss  = err ** 2                                     # 4.0
 
loss.backward()
print(w.grad, b.grad)                                # tensor(8.) tensor(4.)

Let's draw the graph PyTorch built:

The forward pass walked left→right, doing arithmetic and appending nodes. On loss.backward(), PyTorch walked right→left through the same graph, using the chain rule at each node:

• dL/derr = 2·err = 4
• derr/dpred = 1 → dL/dpred = 4
• dpred/dw = x = 2 → dL/dw = 4·2 = 8 ✅
• dpred/db = 1 → dL/db = 4·1 = 4 ✅

That's exactly the Session 010 chain rule, but PyTorch typed it for you.

The word "dynamic" in "dynamic computation graph": the graph is built on every forward pass. If your model has an if statement that takes a different branch this time, the graph is different this time. Autograd will still work. (Contrast: TensorFlow 1.x used a static graph — you had to declare it up front. Everyone hated it. Everyone.)

4 · Rewriting the MNIST MLP in PyTorch

Time for the payoff. The Session 011 forward pass was ~40 lines. Here it is in PyTorch:

import torch
import torch.nn.functional as F
 
# Parameters — leaves with requires_grad=True
W1 = torch.randn(784, 128, requires_grad=True) * 0.01
b1 = torch.zeros(128, requires_grad=True)
W2 = torch.randn(128, 10, requires_grad=True) * 0.01
b2 = torch.zeros(10, requires_grad=True)

Wait — there's a subtle bug in the code above. torch.randn(...) * 0.01 produces a non-leaf tensor (it's the result of a multiplication, so autograd tracks it as MulBackward). The requires_grad=True is on the operand, not the result you assigned to W1. Calling .backward() will fill gradients into the original torch.randn(...), which you no longer have a reference to. W1.grad will be None.

Correct idiom:

W1 = (torch.randn(784, 128) * 0.01).requires_grad_(True)
b1 = torch.zeros(128, requires_grad=True)
W2 = (torch.randn(128, 10) * 0.01).requires_grad_(True)
b2 = torch.zeros(10, requires_grad=True)

The trailing underscore in requires_grad_ means "in-place". We first do the scaling without grad tracking, then flip the flag on the final tensor, making it a leaf. This is the single most common autograd gotcha for beginners. Burn it in.

Now the forward pass:

def forward(X):
    h1  = F.relu(X @ W1 + b1)         # (batch, 128)
    logits = h1 @ W2 + b2             # (batch, 10)
    return logits

Three lines. Compare to the Session 011 version, which was the same computation but with np.dot, manual ReLU, and hand-typed weight init. Same math. Cleaner code.

Loss and backward:

X = torch.randn(32, 784)
y = torch.randint(0, 10, (32,))
 
logits = forward(X)
loss   = F.cross_entropy(logits, y)   # scalar
loss.backward()
 
print(W1.grad.shape)                  # torch.Size([784, 128])  ✅
print(b1.grad.shape)                  # torch.Size([128])       ✅

We wrote zero backprop code. F.cross_entropy combined softmax + NLL into a numerically stable scalar, and .backward() filled all four .grad fields. Every derivative you painstakingly typed in Session 011? Gone. This is why we're all still employed and not going insane.

5 · The SGD update — and why you need .zero_grad()

We have .grad. Now do the update:

lr = 0.1
with torch.no_grad():                 # don't track THIS math
    W1 -= lr * W1.grad
    b1 -= lr * b1.grad
    W2 -= lr * W2.grad
    b2 -= lr * b2.grad
 
# CRITICAL — zero the gradients before the next step
W1.grad.zero_()
b1.grad.zero_()
W2.grad.zero_()
b2.grad.zero_()

Two things worth staring at.

5.1 torch.no_grad() — turn off recording

The parameter update is not something we want to differentiate through. If we didn't wrap it in no_grad(), PyTorch would build graph nodes for W1 -= lr * W1.grad, and the next forward pass would build on top of them, and by step 100 you'd have a computation graph the length of a novella and OOM the moment you called .backward(). no_grad() says "for this block, don't track". Absolutely required.

5.2 .grad accumulates — you must zero it

This is not a bug. It's a design choice, and once you know why, it makes sense. When .backward() is called, PyTorch adds the new gradients to whatever's already in .grad. This lets you do things like:

• Split a huge batch into 4 mini-batches, .backward() each one, and the gradients accumulate as if you'd done the whole batch (this is gradient accumulation, and we cover it in Session 024).
• Compute gradients w.r.t. multiple losses without redoing the forward pass.

But it means: if you don't .zero_() between training steps, gradients from step 1 will still be sitting in .grad when step 2 backprops. Your effective batch size becomes step_count × real_batch_size, and your training will look weird and slowly diverge.

Every PyTorch training loop ends with optimizer.zero_grad() (which loops over parameters and does exactly this). Session 015 will show you torch.optim.SGD, which packages this up nicely.

6 · The full mini-training-loop (raw, no nn.Module yet)

Putting it all together:

lr = 0.1
for step in range(200):
    # 1. FORWARD
    logits = forward(X)
    loss   = F.cross_entropy(logits, y)
 
    # 2. BACKWARD
    loss.backward()
 
    # 3. STEP
    with torch.no_grad():
        for p in [W1, b1, W2, b2]:
            p -= lr * p.grad
            p.grad.zero_()
 
    if step % 20 == 0:
        print(f"step {step:3d}  loss {loss.item():.4f}")

That is a complete supervised-learning training loop. Nine functional lines. In Session 011 the equivalent code was pushing 100. This is the moment the tuition pays off. You wrote every gradient by hand in Session 010 — now you know what's inside .backward(), so it isn't magic. It's just a very efficient version of what you already did.

6.1 What .backward() is doing under the hood

When you call loss.backward(), PyTorch's C++ autograd engine does roughly this:

1. Walk the graph starting from loss, in reverse topological order.
2. At each node, call its grad_fn.backward(grad_output) — a hand-written C++ function (one per op: MmBackward, AddBackward, ReluBackward, etc.) that knows the local Jacobian-vector product.
3. Accumulate the result into each input tensor's .grad (leaves) or pass it further backward (intermediates).
4. Free saved activations as soon as they're no longer needed (this is why you can only call .backward() once by default — pass retain_graph=True if you truly need to call it again).

The engine is multithreaded by default since PyTorch 1.6: independent branches of the graph run in parallel on CPU. On GPU, all backward ops are queued to the same CUDA stream as the forward, so they overlap naturally with the next forward pass if you're careful. The parallelism is invisible until you're profiling — Session 019 shows how to see it in the Kineto trace.

6.2 Vector-Jacobian products — the actual math autograd computes

Autograd never materialises a full Jacobian. If y = f(x) maps ℝⁿ → ℝᵐ, the Jacobian ∂y/∂x is an m×n matrix — potentially gigantic (imagine m = n = 1B). What autograd computes is a vector-Jacobian product (VJP): given ∂L/∂y (a length-m row vector, called the upstream gradient), it computes ∂L/∂x = (∂L/∂y) · (∂y/∂x) in one pass, giving you a length-n vector.

For most ops this VJP is cheaper than materialising J. Example: matmul Y = X @ W. The Jacobian ∂Y/∂W has shape (mn, np) = O(mn²) entries. The VJP is just dW = X.T @ dY, one matmul, same cost as the forward. This is the whole efficiency trick of reverse-mode AD: backward costs ~2× forward, regardless of parameter count. Forward-mode AD (which JAX also supports via jax.jvp) computes Jacobian-vector products instead — cheaper when you have few inputs and many outputs (rare in DL, common in physics simulation).

7 · .detach(), .item(), and .data — three ways to leave the graph

Three utilities you will use every hour.

7.1 .item() — pull a scalar out to a Python number

loss_val = loss.item()          # <class 'float'>, no tensor, no graph

Only works on scalar tensors. Use it for logging, printing, tqdm postfix. Never use tensors in print(loss) inside a training loop — you'll build up references and leak memory.

7.2 .detach() — new tensor, same data, no graph

logits_for_metrics = logits.detach()   # can compute accuracy, plot, etc.

Detach whenever you're going to use a tensor's values but never backprop through them. E.g., computing accuracy on training predictions — you want the numbers, not the gradient path.

7.3 .data — legacy escape hatch

tensor.data returns a raw tensor pointing at the same memory, with no graph. Don't use this. It's still in the API for historical reasons. Modern code uses .detach() (safer — errors out if you try to backprop through it later).

7.4 .clone() vs .detach() — the confusion table

The four combinations you'll meet in the wild:

Common bug: saved = x.detach() then later saved += noise — you just corrupted x. Fix: saved = x.detach().clone().

7.5 · torch.compile — the 2023–2025 story you cannot skip

The biggest change to PyTorch since autograd itself is torch.compile, which shipped in PyTorch 2.0 (March 2023) and became default-stable in 2.5 (Oct 2024). One line:

model = torch.compile(model)          # or: @torch.compile decorator

What it does, in one sentence: it traces your Python forward pass with TorchDynamo (a Python bytecode-level tracer), fuses the resulting graph with TorchInductor (the backend, which lowers to Triton kernels on GPU and OpenMP-vectorised C++ on CPU), and hands you a compiled function that's typically 1.3×–2.5× faster than eager mode on modern GPUs. On the HuggingFace transformers benchmark suite (BERT, GPT-2, T5, ViT), median speedup on A100 is ~1.7× (PyTorch Foundation 2024 benchmarks).

The magic vs. TF 1.x: it's still dynamic. TorchDynamo watches Python bytecode, extracts pure-tensor subgraphs, and falls back to eager mode whenever it hits Python it can't trace (a print, a data-dependent if, a call into a C extension). You never lose eager's debuggability. This design — called "graph breaks" — is what took seven years to figure out.

The rule of thumb for M03–M12: develop and debug in eager, then wrap model = torch.compile(model) for real runs. If you see graph breaks (torch._dynamo.explain(model)(x)), that's where your speedup is leaking.

2025 twist — FlexAttention: In PyTorch 2.5, FlexAttention lets you write custom attention patterns (causal, sliding-window, ALiBi, document masks) as tiny Python functions and get FlashAttention-2-level speed via torch.compile. In 2.6, FlexAttention supports arbitrary block-sparse patterns. We use this heavily in Module M07.

8 · War stories — the autograd bugs that cost me a weekend

W = torch.randn(10, 5) * 0.01
W.requires_grad = True
loss = (W @ x - y).pow(2).sum()
loss.backward()
print(W.grad)             # None. Nothing. Silent failure.

W = torch.empty(10, 5).normal_(std=0.01).requires_grad_(True)
# or
W = torch.randn(10, 5, requires_grad=True)
# then scale INSIDE no_grad if you must:
with torch.no_grad():
    W *= 0.01

losses.append(loss.item())

9 · The autograd cheat sheet

10 · Diagram — forward builds, backward walks

The forward pass is the graph construction. There is no separate "compile" step. Every op you run just... appends a node. Then backward() unwinds it.

11 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §3 (the graph drawing) + §6.2 (VJP math) in 24 hours. Revisit the full session on day 7. Redo the §8 war-story diagnoses from memory on day 14.

Next session (S015): nn.Module — the object that packages parameters + forward together, gives you .parameters() iteration for optimizers, .state_dict() for saving, and .train() / .eval() mode switches. Your MLP will shrink from 40 lines to 8 and start looking like a real project.

Sticky note (keep on your desk): zero → forward → backward → step — one word. Everything else is a variation.

Previous: ← DL S013 · Vectorization · Next: DL S015 · nn.Module →