PyTorch From First Principles

Introduction

PyTorch is remarkably easy to start using before it is easy to understand.

You can create a tensor, copy a model architecture, call loss.backward(), construct an optimizer, and have a training loop running in a few minutes.

That is one of PyTorch’s strengths.

It is also where the trouble begins.

A program can run while broadcasting the wrong dimensions. A parameter can have a gradient and still never be updated. A model can train successfully while consuming inputs whose numerical meaning is wrong. A GPU can appear underutilized even though adding workers makes the program slower. A checkpoint can load without error while the model behaves differently because its preprocessing changed. torch.compile can appear to work while quietly producing more graphs than you expected.

None of those problems are solved by knowing one more API call.

They are solved by knowing what computation PyTorch is actually performing.

That is what this book is about.

Make hidden structure visible before guessing.

We will build PyTorch upward from its smallest mechanisms and repeatedly ask the same kind of question:

What does PyTorch think is happening here, what actually happened, and what evidence would distinguish the two?

The goal is not to memorize PyTorch.

The goal is to make unfamiliar PyTorch code inspectable.

The system we are going to build

The book begins with almost nothing:

one value
    ↓
one parameter
    ↓
one prediction
    ↓
one loss
    ↓
one gradient
    ↓
one update

By the end, those same mechanisms will have expanded into a small GPT-style language model:

raw data
    ↓
Dataset / transforms / DataLoader
    ↓
token tensors
    ↓
embeddings
    ↓
attention
    ↓
residual blocks
    ↓
logits
    ↓
cross entropy
    ↓
autograd
    ↓
optimizer
    ↓
updated parameters
    ↓
checkpoint / generation / evaluation

Around that model sits another system:

CPU work
GPU work
memory
profiling
compilation
reproducibility
measurement
regression testing

Those are not separate subjects.

They are different views of the same running program.

The book builds enough of that program from first principles that when something goes wrong, you can identify which part of the system to interrogate.

The recurring debugging move

A lot of PyTorch debugging begins too late.

A tensor reaches an operation with the wrong shape, so we reshape it until the exception disappears.

A parameter has grad=None, so we add requires_grad=True.

GPU utilization is low, so we increase num_workers.

Training is unstable, so we lower the learning rate.

Compilation is slow, so we disable torch.compile.

Sometimes those changes work.

But a working change is not automatically a diagnosis.

This book uses a different pattern:

predict
    ↓
observe
    ↓
find the first divergence
    ↓
identify the mechanism
    ↓
change one thing
    ↓
verify the repair

The chapters gradually turn that pattern into a collection of concrete diagnostic techniques.

For tensors:

Find where the tensor first became wrong, not merely where the wrong tensor finally became illegal.

For autograd:

Find the first point where the gradient path stops existing.

For modules and parameters:

Inspect the structure PyTorch sees rather than inferring it from the source code.

For input pipelines:

Find where useful work stops flowing between producer and consumer.

For transforms:

Find the first boundary where the sample stops satisfying the representation the model expects.

The particular mechanism changes.

The investigation does not.

PyTorch contains several overlapping structures

One reason PyTorch can become confusing is that several different systems occupy the same Python program.

Consider a model with a few tensors and submodules.

There is the ordinary Python object structure:

model
 ├── encoder
 ├── blocks
 └── head

There is the structure PyTorch has registered:

_modules
_parameters
_buffers

There is the computation graph created by the operations executed during the current forward pass:

parameter
    ↓
operation
    ↓
intermediate
    ↓
loss

There are optimizer parameter groups:

optimizer
    ↓
specific parameter objects

There is the input pipeline:

source
  ↓
Dataset
  ↓
transforms
  ↓
collation
  ↓
DataLoader workers
  ↓
batch

And there is the execution system underneath all of it:

Python
CPU threads
worker processes
host memory
CUDA
GPU memory
compiled graphs

A tensor can belong to one of these systems without belonging to another.

A parameter can participate in autograd without being registered in the model.

A registered parameter can exist without being owned by the optimizer you created earlier.

A transform can return a tensor with the correct shape and dtype while changing what its values mean.

A batch can arrive quickly and still be wrong.

A GPU can be busy while the program as a whole makes little progress.

Much of PyTorch debugging becomes easier once you stop treating all of these structures as one thing.

What this book is designed to teach

The book progresses from small mechanisms to complete systems.

You will learn to reason about:

• tensor geometry — shape, dtype, device, stride, broadcasting, contiguity, layout, and the meaning attached to each axis;
• autograd — what PyTorch records, leaves and intermediates, grad_fn, gradient accumulation, broken paths, retained gradients, and graph lifetime;
• neural networks from raw tensors — parameters, matrix multiplication, nonlinearities, logits, loss, gradients, and manual updates before nn.Module hides the bookkeeping;
• registered model structure — modules, parameters, buffers, recursive composition, state_dict(), device movement, train/eval behavior, optimizer membership, and the places where those structures can disagree;
• input delivery — datasets, collation, multiprocessing, workers, prefetching, persistent state, host-to-device transfer, and the difference between a pipeline that waits and one that merely moves the bottleneck;
• input representation — transforms, scaling, normalization, augmentation, target semantics, pretrained preprocessing contracts, and the difference between a legal tensor and a meaningful one;
• spatial models — channels, height, width, convolution, receptive geometry, and deriving CNN shapes rather than guessing them;
• higher-dimensional representations — features, embeddings, attention, heads, masks, logits, and transformer-style tensor geometry;
• training failures — data, targets, losses, gradients, optimizers, parameter movement, numerical instability, and models that run while refusing to learn;
• execution and performance — CUDA timing, memory, profiling, bottlenecks, compilation, graph breaks, guards, recompilation, and dynamic behavior;
• experiments — baselines, reproducibility, comparable runs, regressions, benchmark validity, and turning discovered failures into tests; and
• integration — assembling the mechanisms into a small GPT-style language model whose behavior can be inspected from input through generation.

The important part is not the size of that list.

It is that the topics accumulate.

Later chapters should feel like larger applications of structures you already understand.

One system, built progressively

The sequence is deliberate.

We begin with the training loop because almost everything else in PyTorch eventually has to attach to it.

Then we expose the tensor flowing through that loop.

Then we expose the graph connecting the tensor operations.

Then we build a neural network without the machinery that normally organizes it.

Only after those pieces are visible do we introduce nn.Module and ask what PyTorch thinks belongs to the model.

Then we move outside the model.

We ask how the next batch reaches it.

Then we ask what that batch actually means.

Only after the input contract is trustworthy do we start exploiting spatial structure with convolution and increasingly complex model architectures.

Later, once all the pieces exist, the book becomes more forensic.

We ask questions like:

Why does this model run but not learn?

Where is GPU memory actually going?

Why is the GPU waiting?

Which part of this operation is asynchronous?

Why did this compiled function produce another graph?

Which input property created that guard?

Did this optimization improve throughput,
or did the measurement merely move?

Is this new result genuinely worse,
or are the two experiments not comparable?

And finally we assemble the pieces.

By then the capstone should not feel like twenty unrelated PyTorch tricks combined into a transformer.

Each mechanism should already have somewhere to attach.

Legal is not the same as correct

This distinction appears again and again throughout the book.

PyTorch is very good at rejecting computations that are impossible.

It cannot reject every computation that is meaningless.

A broadcast may be legal while pairing the wrong axes.

A floating-point image may be perfectly valid while still representing values in 0..255 when the model expects 0..1.

A horizontal flip may execute perfectly while changing the truth of the label.

A parameter may receive gradients while remaining outside the optimizer.

A checkpoint may load with every key matched while inference preprocessing no longer matches training.

A benchmark may run correctly while comparing two different workloads.

So we will repeatedly separate:

CAN PYTORCH EXECUTE THIS?

from

IS THIS THE COMPUTATION WE INTENDED?

Error messages help with the first question.

This book is primarily about learning to answer the second.

Observation before explanation

Another rule follows naturally:

Inspect what happened before explaining why it happened.

Before diagnosing a tensor problem, print the tensor.

Before diagnosing a gradient problem, inspect the graph and the gradients.

Before diagnosing an optimizer problem, check whether the parameter object is actually in the optimizer and whether it moved.

Before diagnosing an input problem, trace one known sample through the transforms.

Before diagnosing DataLoader performance, measure batch wait and consumer work separately.

Before diagnosing CUDA performance, distinguish host time from device time.

Before diagnosing compilation, inspect the actual graph breaks and recompilation reasons.

Before diagnosing a regression, establish that the two runs are comparable.

An explanation that arrives before the evidence is a hypothesis.

Treat it as one.

Using AI without outsourcing understanding

Modern PyTorch development increasingly involves AI-generated code.

That is not a problem this book tries to avoid.

It is one of the reasons the book matters.

An assistant can often produce a working training loop, a custom dataset, an attention block, a profiler script or a torch.compile configuration much faster than you would write one from scratch.

The question is what happens when the generated program behaves strangely.

If the only thing you know how to do is ask for another implementation, then every failure becomes another generation problem.

If you understand the underlying mechanisms, you can ask much better questions:

Show me the tensor shapes after every operation.

Which exact edge disconnects this parameter from the loss?

Which registered parameters are absent from the optimizer?

What does this transform assume its input range is?

Where is the DataLoader consumer actually waiting?

Which graph break caused this function to recompile?

What observation would distinguish these two explanations?

AI can generate the implementation.

You still need enough understanding to know what evidence to request when the implementation stops making sense.

That is the relationship with AI this book assumes.

What you should be able to do after reading

The goal is not that you remember every PyTorch API used in these chapters.

The goal is that an unfamiliar PyTorch program becomes legible.

You should be able to open one and ask:

What does this tensor represent?

What should its shape be?

What does each dimension mean?

Where was this tensor created?

What operations produced this value?

Does the loss depend on this parameter?

Where does that dependency first disappear?

Which tensors does PyTorch consider parameters?

Which state is registered?

Which exact parameter objects does the optimizer own?

Did those objects move after the step?

What happened to this sample before the model received it?

Does the transformed sample still mean what its target says it means?

Where does the next batch spend its time?

Is the CPU waiting, computing, or contending?

Is this CUDA timing measuring submission or completed work?

Where is memory allocated?

Why did compilation create another graph?

What assumption produced that guard?

Did this intervention actually improve throughput or accuracy?

Are these two experiments comparable?

Those questions are more useful than knowing where a particular method sits in the documentation.

They tell you how to investigate.

You will also learn when not to use more PyTorch

Not every problem needs another abstraction.

You do not always need:

• a more complicated model;
• another layer;
• mixed precision;
• additional workers;
• pinned memory;
• compilation;
• a custom CUDA kernel;
• a more elaborate optimizer;
• another augmentation;
• another framework around PyTorch.

Sometimes the tensor is simply wrong.

Sometimes the parameter is not in the optimizer.

Sometimes the dataset is duplicated.

Sometimes the validation pipeline differs from the training contract.

Sometimes the GPU is not the bottleneck.

Sometimes a plain eager implementation is already fast enough.

The book therefore tries to earn complexity.

We add machinery when the simpler system exposes a limitation that the new mechanism actually solves.

Who this book is for

You should be comfortable with basic Python:

variables
functions
loops
classes
lists
dictionaries
running scripts

You do not need previous PyTorch experience.

You do not need to know transformer internals before beginning.

You do not need to arrive with a complete mathematical treatment of neural networks.

You do need to be willing to run small experiments.

Many chapters deliberately break working code.

That is not an interruption to the lesson.

It is the lesson.

The book is written for programmers who want to be able to investigate PyTorch rather than merely operate it when everything goes according to plan.

What this book does not try to cover

This is not an encyclopedia of every PyTorch API.

It does not attempt to make every specialization a prerequisite for competence.

Topics such as:

• large-scale distributed training;
• multi-node infrastructure;
• custom CUDA kernels;
• quantization;
• production serving;
• LoRA and other parameter-efficient fine-tuning systems;
• mixture-of-experts architectures;
• every current transformer variant;
• every vision architecture;
• every compiler backend;

are important.

But they are not the foundation.

The purpose of this book is to give those subjects somewhere to attach.

If tensors, computation graphs, registered state, input contracts, optimization, execution, measurement and debugging are clear, unfamiliar features stop arriving as isolated incantations.

They become extensions of a system you already understand.

The promise

By the end of PyTorch From First Principles, you should be able to encounter PyTorch code you have never seen before and work out:

what tensors mean, how they were produced, what PyTorch has recorded about them, what state belongs to the model, what the optimizer can actually change, what representation reaches the model, where execution time and memory are going, and what evidence distinguishes a real repair from a plausible-looking change.

More compactly:

Make the hidden structure visible. Find the first divergence. Repair the mechanism that caused it. Verify that the system actually improved.

That is the standard the rest of the book is built around.

The final GPT-style model matters.

But it is not the real destination.

The real destination is reaching the point where, when a model does something you did not expect, you know how to find out why.