Two engineers train the same model with the same code. One gets 3% better validation accuracy. The difference is a single missing line, model.eval(), so the other's validation pass ran with dropout on and with BatchNorm updating its statistics on the validation set. A different team's run slowly runs out of GPU memory over six hours, because total_loss += loss kept every step's computation graph alive. A third team spent two weeks tuning a model whose middle layers never trained at all, because they were stored in a plain Python list and the optimizer never saw their parameters.

None of these bugs raise an error. They are invisible unless you know what PyTorch is doing underneath. This section opens the box. It covers how nn.Module finds and saves parameters, how to write your own differentiable operations, how hooks let you watch and edit a running model, how data loading works, what mixed precision really changes, what torch.compile does, how to time and profile honestly, how to make runs reproducible, and the bugs that bite most often. Every runnable example here works on a CPU, and the exercises rebuild the key mechanisms in plain Python and NumPy so you understand them from the inside.

How nn.Module finds your parameters

When you write self.fc = nn.Linear(10, 5) inside a module's __init__, PyTorch does not just store an attribute. nn.Module overrides __setattr__, the method Python calls on every attribute assignment, and sorts the value into one of three dictionaries.

• A Parameter (a tensor marked as trainable) goes into self._parameters.
• A Module goes into self._modules, which is how modules nest into a tree.
• A tensor registered with self.register_buffer(name, tensor) goes into self._buffers. Buffers are state that is saved and moved between devices with the model but not trained, such as BatchNorm's running mean.

Everything else follows from those three dictionaries. parameters() walks the module tree and yields every parameter (each one only once, even if it is shared by two submodules, as with tied embeddings). named_parameters() gives dotted names like "encoder.layers.3.attn.q_proj.weight". state_dict() collects parameters and persistent buffers under those names, and load_state_dict() copies them back, complaining (when strict=True, the default) about missing or unexpected keys. model.to("cuda") moves both parameters and buffers. model.train() and model.eval() set a training flag on every module in the tree, which layers such as dropout and BatchNorm read.

The bug from the opening follows directly. A plain Python list is not a Module, so self.layers = [nn.Linear(10, 10) for _ in range(4)] stores the list as an ordinary attribute, and its layers' parameters never reach _modules. They are invisible to parameters(), so the optimizer never updates them, state_dict() never saves them, and .to("cuda") never moves them. nn.ModuleList (and nn.ModuleDict) exist to register containers properly.

import numpy as np

class Module:
    def __init__(self):
        raise NotImplementedError

class ModuleList(Module):
    def __init__(self, modules=()):

Custom autograd functions

Autograd knows the backward rule of every built-in operation. Sometimes you want to supply your own, by subclassing torch.autograd.Function and writing a static forward and backward. There are four good reasons.

• A better backward than autograd would derive. For example, a fused softmax cross-entropy whose backward is the clean $\mathbf{p} - \mathbf{y}$ of §0.2, instead of a chain of saved intermediates.
• Memory. Save only the inputs and recompute cheap intermediates in backward, rather than letting autograd store every intermediate tensor.
• A gradient that is not the true derivative. The straight-through estimator rounds in the forward pass (as in quantization-aware training) but passes the gradient through as if the operation were the identity, since rounding's true derivative is zero almost everywhere. You'll meet it again when models are trained for low-precision inference.
• Wrapping a custom kernel. A Triton or CUDA kernel (§11.6) becomes usable in a model once its forward and backward are wrapped in a Function.

The contract is short. forward(ctx, *inputs) computes the output and stores what backward will need with ctx.save_for_backward(...). backward(ctx, grad_output) receives the upstream gradient and returns one gradient per input (or None for inputs that need none), which is exactly a VJP. Always verify a custom backward with torch.autograd.gradcheck in float64, the finite-difference check of §0.2.

import numpy as np

def gelu_tanh(x):
    raise NotImplementedError

def gelu_tanh_backward(x, grad_out):
    raise NotImplementedError

Hooks, watching and editing a running model

Hooks are callbacks that PyTorch runs at specific moments. They let you observe or modify a model without editing its code, which makes them essential for debugging, logging, and interpretability research.

• module.register_forward_hook(fn) calls fn(module, inputs, output) after the module's forward. Return a value to replace the output. This is how you record activations, the starting point of the probing and steering methods of interpretability research.
• module.register_forward_pre_hook(fn) runs before forward and can modify the inputs.
• tensor.register_hook(fn) calls fn(grad) when that tensor's gradient is computed, and can return a modified gradient. It is useful for inspecting gradients of intermediate tensors, or for per-tensor gradient manipulation.
• module.register_full_backward_hook(fn) sees the gradients with respect to a module's inputs and outputs.

Every registration returns a handle, and handle.remove() removes the hook. Forgetting to remove hooks that store tensors is another quiet memory leak.

Datasets and data loaders

A map-style Dataset is anything with __len__ and __getitem__(i). An iterable-style IterableDataset yields examples from __iter__, for streams too large to index, which is the usual choice for pretraining corpora (§6.1). A DataLoader wraps either and handles the rest.

• A sampler decides the order of indices (sequential, or shuffled with shuffle=True), and batch_size groups them. drop_last=True discards a final short batch, which keeps shapes constant (useful for compiled models and for BatchNorm).
• A collate_fn turns a list of examples into a batch. The default stacks tensors, which fails for variable-length sequences. Language models need a collate function that pads to the longest sequence in the batch and returns a mask, which the exercise below builds.
• num_workers > 0 loads batches in background processes, so the GPU is not left waiting on Python. Each worker gets its own copy of the dataset object. pin_memory=True puts batches in page-locked host memory, which makes tensor.to("cuda", non_blocking=True) faster and lets the copy overlap with compute.

import numpy as np

def batch_indices(n, batch_size, shuffle, drop_last, seed=0):
    raise NotImplementedError

def pad_collate(seqs, pad_id=0):
    raise NotImplementedError

Train mode, eval mode, and no_grad

Three switches are easy to confuse, and they do completely different things.

So a correct evaluation loop needs both model.eval() and torch.no_grad() (or inference_mode), and the training loop must switch back with model.train(). Forgetting eval() makes validation numbers noisy and wrong. Forgetting no_grad() makes evaluation use far more memory than it should.

Mixed precision and number formats

A floating-point number stores a sign bit, an exponent, and a mantissa (the significant digits). The exponent's width sets the range (how large and how small a number can be), and the mantissa's width sets the precision (how many significant bits).

Here machine epsilon is the gap between 1 and the next representable number. Halving the bits roughly halves memory traffic and, on tensor cores, multiplies matmul throughput, which is why every large model trains in reduced precision. The two 16-bit formats make opposite trade-offs, and the difference is a favorite interview question. fp16 keeps more precision but has a tiny range. Below about $6 \times 10^{-5}$ its values become subnormal and lose precision, below about $6 \times 10^{-8}$ small gradients underflow to zero, and above 65,504 large activations overflow to infinity. bf16 (Kalamkar et al., 2019) keeps fp32's full exponent range and gives up precision instead, so it almost never overflows or underflows, while keeping only about 2 to 3 significant decimal digits.

Explore the formats below. Slide a value across the range, and watch where each format rounds it, turns it into a subnormal, flushes it to zero, or overflows. The lower panel shows the loss-scaling trick that makes fp16 training work.

What autocast and GradScaler do

Automatic mixed precision keeps the weights themselves in fp32 (the master weights) and casts them to 16 bits on the fly inside selected operations. Inside torch.autocast(device_type, dtype), operations that are safe and fast in low precision (matmuls, convolutions, linear layers) run in fp16 or bf16, while operations that need range or accumulate many terms (on CUDA, softmax, log-softmax, layer norm, losses, and many reductions) run in fp32. The optimizer updates the fp32 master weights, because adding a tiny update to a large weight in bf16 would often round away entirely ($1 + 2^{-8}$ rounds to 1 in bf16, as the cell shows).

With fp16, gradients also need loss scaling (Micikevicius et al., 2018). Multiply the loss by a large factor $S$ before backward(), so every gradient is $S$ times larger and stays above fp16's underflow threshold. Then divide the gradients by $S$ in fp32 before the optimizer step. A dynamic scaler starts with a large $S$, and whenever a gradient overflows to infinity it skips that step and halves $S$. After a run of good steps it doubles $S$ again. That is exactly what torch.amp.GradScaler does. bf16 has fp32's range, so it needs no loss scaling, which is a big part of why it became the default for training. The precision story for frontier training, including fp8, continues in §8.5.

import numpy as np

class DynamicLossScaler:
    def __init__(self, init_scale=2.0 ** 16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000):
        raise NotImplementedError

    def step(self, scaled_grads):
        raise NotImplementedError

Gradient accumulation

§0.6 showed why accumulation works. The average of $k$ equal-sized micro-batch averages is the full batch's average, so $k$ small forward and backward passes followed by one optimizer step take exactly the step one big batch would. PyTorch makes the sum free, because .grad accumulates across backward() calls (§0.3). If every micro-batch has the same number of examples, calling (loss / k).backward() on each micro-batch's mean loss gives the full batch's mean gradient.

When micro-batches have different sizes, as with language-model batches whose sequences contain different numbers of real (non-padding) tokens, dividing by $k$ is wrong. Averaging $k$ per-micro-batch means gives each micro-batch equal weight, so tokens in short micro-batches count for more than tokens in long ones. The correct gradient weights each micro-batch by its share of the tokens,

where $n_i$ is the number of counted tokens in micro-batch $i$ and $\bar{\mathcal{L}}_i$ is its mean loss. In plain words, undo each micro-batch's mean by multiplying by its count, add, and divide by the total count. In code it is simpler still. Count the target tokens of the whole accumulated batch before the first backward pass, then for each micro-batch call backward() on its summed per-token loss divided by that global count. This is not hypothetical. Until October 2024, the Hugging Face Trainer averaged per-micro-batch means for causal language models, so runs with accumulation reached a worse loss than the same runs in one big batch, and the library switched to dividing by the total number of counted tokens (Hugging Face, 2024).

One optimizer step, in order

A real training step has six parts, and their order matters.

1. Count. Count the target tokens over all $k$ micro-batches (and, with data parallelism, sum the count across devices).
2. Accumulate. For each micro-batch, run the forward pass under autocast, divide the summed token loss by the global count, and call backward(). With data parallelism, skip the gradient synchronization on all but the last micro-batch (model.no_sync() in DDP), since one all-reduce per step is enough (§10.2).
3. Unscale. With fp16 loss scaling only, divide the gradients by the scale.
4. Clip. Clip the global norm of the accumulated gradient, and log the norm before clipping.
5. Step. Call optimizer.step(), then scheduler.step(), once per accumulated batch, never once per micro-batch.
6. Zero. Call optimizer.zero_grad(set_to_none=True).

The equality with the big batch is exact up to rounding, with two exceptions worth knowing. Layers that look across the batch, such as BatchNorm, see only the micro-batch. And dropout draws different masks, which changes the gradient but not its distribution.

import numpy as np

def combine_microbatch_grads(grads, counts):
    raise NotImplementedError

Here is the whole step in PyTorch, on a tiny model so it runs anywhere. It checks the claim numerically. Token-weighted accumulation over four uneven micro-batches reproduces the big batch's gradient, and the naive average of per-micro-batch means does not.

torch.compile

PyTorch runs eagerly. Each operation launches its own kernel as Python reaches it, which is flexible and easy to debug but leaves performance on the table. A chain like x * torch.sigmoid(x) + bias launches three kernels, each reading its input from GPU memory and writing its output back. For such memory-bound elementwise operations (§11.2), the memory traffic, not the arithmetic, is the cost.

torch.compile(model) (Ansel et al., 2024) captures your Python code as a graph and compiles it. It has three stages. TorchDynamo analyzes the Python bytecode as it runs and extracts sequences of PyTorch operations into graphs, falling back to ordinary Python (a graph break) for anything it cannot capture. AOTAutograd traces the backward pass ahead of time, so it can be optimized too. TorchInductor generates fused kernels, Triton on GPUs and C++ on CPUs, so that the three operations above become one kernel that reads x once and writes the result once. Kernel fusion and profiling are the subject of §11.7.

Practical points.

• The first call is slow, since it compiles. Later calls reuse the compiled code, as long as input shapes and other guarded properties do not change. A new shape triggers a recompilation, after which Dynamo by default marks that dimension as dynamic and compiles shape-generic code, so it does not recompile for every new size. dynamic=True asks for shape-generic code from the start.
• Graph breaks (from print, .item(), data-dependent Python control flow, unsupported libraries) split the graph and limit fusion. TORCH_LOGS="graph_breaks" lists them.
• mode="reduce-overhead" uses CUDA graphs to cut kernel-launch overhead for small batches, and mode="max-autotune" searches for faster kernel configurations at a longer compile time.
• The usual speedups for training Transformers are tens of percent, which is why most large training codebases compile their model.

Timing and profiling honestly

GPU operations are asynchronous. A call like y = model(x) returns as soon as the work is queued, long before the GPU finishes. So time.time() around it measures how fast Python can enqueue work, not how fast the GPU does it. Honest timing needs three things. Run a few warmup iterations first (to trigger compilation, autotuning, and memory allocation). Call torch.cuda.synchronize() before reading the clock, every time. And report a median over many repetitions. torch.utils.benchmark.Timer does all of this for you.

To find out where the time goes, use torch.profiler. It records every operator and kernel with its time, memory, and input shapes, and record_function lets you label your own regions. The output can be viewed as a sorted table or exported as a timeline (a Chrome trace) that shows gaps where the GPU sits idle waiting for the CPU, often the most revealing view.

Reproducibility

Making two runs produce the same numbers takes more than one seed.

• Seed every generator, torch.manual_seed, np.random.seed (or use explicit NumPy generators), and random.seed.
• Seed the data order, with a torch.Generator passed to the DataLoader, and seed workers as described above.
• Ask for deterministic kernels. Some GPU operations (atomic additions in scatter and index operations, some convolution algorithms) are nondeterministic because the order of floating-point additions varies, and floating-point addition is not associative. torch.use_deterministic_algorithms(True) makes PyTorch use deterministic versions or raise an error. On CUDA it also requires setting CUBLAS_WORKSPACE_CONFIG=:4096:8, and torch.backends.cudnn.benchmark = False stops cuDNN from choosing algorithms by timing.

Even then, bitwise reproducibility holds only on the same hardware, software versions, and configuration. A different GPU, a different number of GPUs, or a different batch split changes the order of additions and therefore the low bits of the results. For research, aim to make results reproducible statistically (the same conclusion across seeds) and report variation across seeds, not just a single run.

The most common PyTorch bugs

Read more

• Paszke et al. (2019). PyTorch: An Imperative Style, High-Performance Deep Learning Library. The design principles behind modules, eager execution, and autograd.
• PyTorch documentation. Extending PyTorch. Custom autograd functions, done properly.
• Micikevicius et al. (2018). Mixed Precision Training. Master weights, loss scaling, and fp32 accumulation.
• Kalamkar et al. (2019). A Study of BFLOAT16 for Deep Learning Training. Why bf16's range matters more than its precision.
• Hugging Face (2024). Fixing Gradient Accumulation. The token-normalization bug in a widely used trainer, and its fix.
• Ansel et al. (2024). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. How TorchDynamo and TorchInductor work.
• PyTorch documentation. Automatic Mixed Precision package - torch.amp, Reproducibility, and the Performance Tuning Guide. The official guidance on autocast, determinism, and speed.
• Edward Z. Yang (2019). PyTorch internals. Tensors, dispatch, and autograd from the inside.

Watch next

▶

▶

▶