the tensor was only the handle
I began with a question that seemed too small for an article about an entire
framework: when I write x, what exactly do I have?
base = torch.arange(12, dtype=torch.float32).reshape(3, 4)
transposed = base.t()
flattened = transposed.reshape(-1)
print(base.stride(), transposed.stride(), flattened.stride())
print(
base.untyped_storage().data_ptr()
== flattened.untyped_storage().data_ptr()
)
The result was:
(4, 1) (1, 4) (1,)
False
I had written reshape, not clone, and still acquired new storage. The
transpose itself moved no values. It changed metadata from row-major strides
(4, 1) to (1, 4). Flattening that order into one contiguous dimension could
not be expressed by another legal stride over the original bytes, so
reshape quietly copied.
That one copy is a useful entrance into PyTorch because nearly every subsystem
had some opinion about it. The tensor object knew its sizes, strides, dtype,
device, storage offset, dispatch keys, version counter, and gradient history.
The dispatcher had to choose a reshape path. Autograd had to preserve the
relationship to base. A compiler could functionalize the operation or guard
its layout. An exporter would have to state which shapes remain valid. A
distributed wrapper might later place the resulting gradient inside a
communication bucket. The profiler would report operators, not the sentence I
had in mind when I wrote the source.
Here is the program I carried through the rest of the investigation:
class TinyBlock(torch.nn.Module):
def __init__(self):
super().__init__()
self.up = torch.nn.Linear(3, 8)
self.norm = torch.nn.LayerNorm(8)
self.dropout = torch.nn.Dropout(0.25)
self.down = torch.nn.Linear(8, 2)
self.register_buffer("temperature", torch.tensor(2.0))
def forward(self, x):
h = self.up(x)
h = torch.nn.functional.gelu(self.norm(h))
h = self.dropout(h)
return self.down(h) / self.temperature
model = TinyBlock()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)
logits = model(x)
loss = torch.nn.functional.cross_entropy(logits, target)
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)
There is no claim that this is a useful network. It is deliberately just large enough to exercise parameters, a buffer, a child-module tree, a stochastic operator, saved tensors, a reduction, gradients, and stateful optimization. The point is to identify who owns each fact while this code runs.
This chapter stays inside PyTorch. Backpropagation is not derived again, and I do not teach transformer architecture, inspect a vLLM server, or explain how an all-reduce ring moves bytes. Those have their own chapters. Here the questions are narrower and more architectural: what contract does PyTorch present at each boundary, how do its subsystems compose, and which object is responsible when the contract breaks?
The executable observations use the arm64 macOS wheel for PyTorch 2.13.0,
whose embedded git revision is
cf30153c4c131c8164ee7798e5022d810682e2cb, on CPython 3.14.6. The wheel
uses Accelerate BLAS and OpenMP, and MPS was compiled into it. MPS reported
unavailable inside the restricted repository process and available when the
verifier received normal local device authority, so I ran one bounded MPS
correctness control there. CUDA, cuDNN, NCCL, ROCm, XPU, and MPI were absent.
I separately inspected the clean upstream source tree at
12c0042f650864ce75ae2c6b83cd5f4b30c9bbc6.
That commit is newer than the release wheel. I use the wheel to establish
runtime behavior and the checkout to establish named implementation
structures. I do not pretend that reading the later source executes it.
The complete probes live in
writing-notes/probes/pytorch-framework/. One needs permission to start
PyTorch’s shared-memory manager, and the two-process distributed probe needs a
loopback socket. The CPU arithmetic, one MPS forward/backward, source
assertions, and all displayed outputs are otherwise local. CUDA behavior below
is either a PyTorch API contract or an explicitly unexecuted source path.
the Python object does not contain the numbers
At the Python level, x is a torch.Tensor. At the native boundary it is a
handle around an intrusive reference to a C++ TensorImpl. “Intrusive” means
the reference count lives with the referenced allocation rather than in a
separate control block. Copying a tensor handle normally increments that
count. Destroying the final handle can release the implementation and, through
it, the storage. This is object lifetime, not mathematical tensor equality.
The pinned
TensorImpl.h
contains four fields that explain most observable tensor behavior:
Storage storage_;
std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_;
c10::impl::SizesAndStrides sizes_and_strides_;
DispatchKeySet key_set_;
This is shortened, and TensorImpl contains substantially more state. The
important division is already visible. Storage owns access to bytes.
sizes_and_strides_ describes a multidimensional interpretation.
autograd_meta_ is optional gradient-related metadata.
key_set_ helps choose what an operator means for this tensor.
The storage has its own implementation object. In the pinned
StorageImpl.h,
the central fields are a DataPtr, a byte size, and an Allocator*:
DataPtr data_ptr_;
SymInt size_bytes_;
Allocator* allocator_;
A DataPtr combines an address with the information required to release or
otherwise manage it. The allocator decides how bytes are obtained. A CPU
allocation, a pinned host allocation, a CUDA caching allocation, a memory
mapped region, or storage supplied by an extension need not share the same
release procedure. This is why reducing a tensor to “pointer plus dimensions”
loses an ownership edge.
For an ordinary strided tensor, an index reaches the element at logical offset
where is storage_offset() and is the stride of dimension $d`.
Multiplying by the element size converts that element offset into a byte
offset. Shape limits the valid indices. Dtype tells an operator how many bytes
form an element and how to interpret them. Device says which execution and
memory domain owns those bytes. Layout says whether the ordinary strided model
even applies.
The six tensors in the first probe made these distinctions concrete:
name shape stride contiguous shares base storage
base (3, 4) (4, 1) yes yes
columns (3, 2) (4, 2) no yes
transposed (4, 3) (1, 4) no yes
flattened (12,) (1,) yes no
cloned (3, 2) (2, 1) yes no
detached (3, 4) (4, 1) yes yes
columns = base[:, ::2] skips every second element in the inner dimension.
transposed = base.t() swaps sizes and strides. Both are views: new tensor
metadata over the same storage. cloned has independent bytes and preserves
values. detached shares storage while severing the result from the autograd
history used to propagate gradients. Detach is not copy.
The distinction between view and reshape is conditional, not stylistic.
Tensor.view
requires a stride-compatible grouping of dimensions. reshape attempts such
a view and may allocate a contiguous copy when it cannot represent the
requested indexing with the existing strides. Code that depends on sharing
must not infer it from the verb.
Broadcasting introduces an even stranger view. Expanding a one-element tensor
to length four produced stride (0,). Every logical index names the same
storage location:
expanded = torch.tensor([1.0]).expand(4)
print(expanded.stride()) # (0,)
expanded.add_(1)
PyTorch rejected the in-place add:
unsupported operation: more than one element of the written-to tensor refers
to a single memory location
The problem is not that addition is unsupported. A parallel or sequential implementation would write the same address four times, and the requested elementwise semantics would be ambiguous. Cloning creates four independent locations and makes the update meaningful.
PyTorch tracks mutation with a version counter shared by relevant aliases. I
mutated the detached handle and watched both the value and counter on base
change:
detached mutation: base[0,0]=99, version 0->1
Autograd uses this counter to detect when a value saved for backward no longer matches the forward computation. Compiler transforms also need alias and mutation information. A mistake in view metadata can therefore corrupt much more than one indexing operation.
Contiguity itself is relative to a memory format. The default contiguous format
expects row-major-like strides after ignoring size-one dimensions. A
four-dimensional image tensor can instead be contiguous in channels_last
format, arranging channel values so kernels can consume them efficiently.
is_contiguous(memory_format=...) asks a specified question. A bare
“contiguous tensor” is incomplete when more than one recognized format is in
play.
Not every layout is strided dense storage. PyTorch also represents sparse COO
coordinates, compressed sparse row and column forms, block-sparse forms,
nested tensors for ragged structure, and backend-specific or quantized
layouts. A meta tensor carries shapes, strides, dtypes, and operator behavior
without real payload storage. A fake tensor goes further by imitating device
and alias metadata so tracing systems can propagate facts without executing
the real kernel. The public Tensor API is a common handle over these
different representations, not evidence that they share one memory formula.
Even a dense tensor may contain lazy metadata. Complex conjugation and negation can be represented by bits that defer work until an operator needs resolved values. Their dispatch keys let PyTorch interpose the required semantics. This is one early hint that a tensor’s type is not a single C++ class name. It is the product of dtype, layout, device, dispatch modes, wrappers, subclass behavior, gradient state, and aliasing.
The practical rule is simple but stricter than “tensors are arrays.” A tensor is a typed, indexed, dispatchable view of storage, possibly with history. Ask separately whether two handles share bytes, whether their logical elements overlap, whether they share a version counter, whether they carry the same gradient edge, and whether an operation must copy to satisfy its output contract.
one operator name opens a dispatch stack
Consider the first layer:
hidden = self.up(x)
nn.Linear eventually asks for a linear operator over x, a weight, and a
bias. No C++ function pointer is named. The operation is identified by a
schema: a namespace, name, overload, typed arguments, return types, and alias
or mutation annotations. An overload such as aten::add.Tensor is distinct
from scalar addition even though Python spells both with +.
ATen is PyTorch’s operator library and common tensor API. Operator definitions
generate much of the glue among Python bindings, C++ calls, dispatcher
registration, autograd formulas, and backend implementations. Code generation
is why one cannot reconstruct the call path by finding a single hand-written
torch.add function.
At runtime the dispatcher combines several sources of state:
- the operator handle and overload;
- dispatch keys carried by tensor arguments;
- thread-local included or excluded keys;
- active Python modes;
- registered kernels and fallbacks;
- alias keys that group related backends or functionality.
For the probe’s CPU tensor, PyTorch printed:
DispatchKeySet(CPU, ADInplaceOrView, AutogradCPU, AutocastCPU)
CPU identifies a backend family. AutogradCPU says a reverse-mode wrapper
may be needed. ADInplaceOrView participates in mutation and view handling.
AutocastCPU gives mixed precision a place to intercept eligible operators.
Not all these kernels execute on every call. Dispatch priority and mode state
decide which key is selected first; a wrapper kernel can do work and
redispatch after removing itself from consideration.
That redispatch step is the part most compact explanations omit. Autograd is
not a post-processing pass that wakes up after an arbitrary CPU function.
For differentiable operations, an autograd wrapper can prepare gradient
history, invoke a lower backend implementation, and connect the output to a
backward node. Autocast can choose input dtypes and redispatch. Functionalization
can replace a mutating operator with functional equivalents and later repair
observable input mutation. vmap can apply a batching rule. A Python tensor
subclass or TorchDispatchMode can inspect or replace the operation. Each
feature acts like an interpreter around the same operator vocabulary.
I installed a TorchDispatchMode that did nothing except record the operations
and call onward. The forward and backward of one linear, a custom affine
operator, a square, and a mean began like this:
aten.t.default
aten.addmm.default
asquare.affine.default
aten.pow.Tensor_Scalar
aten.mean.default
aten.ones_like.default
aten.expand.default
aten.div.Scalar
aten.pow.Tensor_Scalar
aten.mul.Scalar
...
The linear spelling disappeared. On this input it decomposed to transpose
plus addmm, which computes bias plus matrix multiplication. Backward emitted
more ATen operations. Twenty-three dispatched calls crossed the mode, with
fourteen distinct overloads. This trace is evidence for this run, not a
promise that every device, shape, compiler mode, or PyTorch revision decomposes
linear identically.
The pinned
OperatorEntry.cpp
builds each operator’s dispatch table. Its
computeDispatchTableEntryWithDebug logic resolves direct registrations,
composite kernels, backend fallbacks, and missing-kernel cases. Registration
is separate from selection. A schema can exist without a kernel for the
combination of backend and feature keys presented at runtime.
Several registration categories matter:
- A backend kernel implements the operation for a device or layout family,
such as CPU, CUDA, MPS, sparse CPU, or an out-of-tree
PrivateUse1accelerator. - A composite implementation expresses an operation in terms of other ATen operations. Depending on its registration category, autograd may be derived by differentiating those constituents or supplied separately.
- A meta or fake implementation computes output metadata without payload execution.
- A fallback handles many operators for one key, useful for modes or backends whose generic behavior can be expressed uniformly.
- A wrapper kernel performs feature logic and redispatches to the next key.
Eventually a CPU backend may run a vectorized loop, TensorIterator, a
specialized kernel, or a call into an external math library. A CUDA backend may
select a native kernel or call cuBLAS, cuDNN, or another library. Those choices
can depend on dtype, shape, stride, build flags, determinism settings, and
hardware. “ATen operator” names the semantic boundary while leaving the final
instruction sequence unspecified.
TensorIterator deserves one sentence because it explains a large class of
elementwise operations. It computes broadcasting, dtype promotion, output
shape, stride traversal, and iteration splitting so a kernel can focus on the
per-element calculation. It is not used for every operator, and it is not a
compiler graph. Instead, the helper constructs a legal runtime iteration
over tensor operands.
The schema’s alias annotations are equally important. If an operator returns
a fresh tensor, mutates an argument, returns a view, or writes an out=
argument, transforms need to know. A false “fresh output” declaration can let a
compiler reorder or discard a mutation. A missing fake implementation can
block shape propagation. A missing autograd registration can make forward
correct and training wrong.
I tested the full registration surface with a deliberately boring custom operator:
@torch.library.custom_op("asquare::affine", mutates_args=())
def affine(x: torch.Tensor, scale: float, shift: float) -> torch.Tensor:
return x * scale + shift
@affine.register_fake
def _(x, scale, shift):
return torch.empty_like(x)
def backward(ctx, grad_output):
return grad_output * ctx.scale, None, None
affine.register_autograd(backward, setup_context=setup_context)
The dispatch table contained CPU, Meta, and AutogradCPU entries. The recorder
saw asquare.affine.default as one opaque operator even though its eager
Python body used multiply and add. The custom operator’s schema created a
boundary; the fake and autograd registrations taught two other interpreters
how to cross it.
This is a better mental model than a ladder from Python to C++ to CUDA. PyTorch dispatch is a stack with multiple valid interpreters. Device selection is one axis. Automatic differentiation, batching, autocast, functionalization, subclasses, fake execution, and debugging modes are other axes. The selected path is the ordered composition, not merely the lowest backend kernel.
state survives longer than one operator
a Module is a registry with a call protocol
The TinyBlock looks like an ordinary Python class because it is one. Its
special behavior comes from nn.Module intercepting assignments and calls.
Assigning an nn.Parameter registers trainable state. Assigning another
Module registers a child. register_buffer registers tensor state that
should move and serialize with the module without being optimized as a
parameter.
The pinned
module.py
maintains three mappings:
_parameters
_buffers
_modules
Those dictionaries, plus non-persistent buffer names, hooks, and ordinary
attributes, are the backbone of the module object graph. self.note = "hello"
is just Python state. self.temperature = torch.tensor(2.0) would also be an
ordinary tensor attribute. Only register_buffer("temperature", ...) makes it
participate in buffer traversal, device conversion, and the default state
dictionary.
My probe reported:
children=['up', 'norm', 'dropout', 'down']
parameters=['up.weight', 'up.bias', 'norm.weight', 'norm.bias',
'down.weight', 'down.bias']
buffers=['temperature']
state_dict=['temperature', 'up.weight', 'up.bias', 'norm.weight',
'norm.bias', 'down.weight', 'down.bias']
The dotted names come from recursively walking child registries. They are paths, not object identities. Two paths can refer to the same child or parameter, as in weight tying. Traversal functions usually remove duplicate objects by default. A checkpoint must preserve the intended aliasing or the loaded model may have numerically equal but independently trainable weights.
Parameter is a Tensor subclass with registration meaning. Its multiplication
semantics remain those of a tensor. A parameter normally defaults to
requires_grad=True, and assigning it to a module tells parameters() and
optimizers where to find it. Replacing module.weight with a new Parameter
changes the registered object. An existing optimizer may still hold the old
parameter reference. This is why device moves, parameter replacement,
load_state_dict(assign=True), and optimizer construction order can interact.
Calling model(x) is not exactly calling model.forward(x).
Module.__call__ reaches _wrapped_call_impl and _call_impl, which arrange
forward pre-hooks, the forward call, forward hooks, backward-hook setup, and
exception behavior. Directly invoking forward skips that protocol. Hooks are
useful for observation and carefully designed instrumentation, but they also
become mutable program state that compilers may need to guard or ignore.
I attached a root pre-hook and post-hook. They fired in this order:
forward_pre
forward_post
That result sounds trivial until hooks modify an input, replace an output,
retain tensors, record activations, or run collectives. Global hooks, per-module
hooks, prepend, keyword handling, and always_call affect ordering. A hook
is not a comment attached to the graph. It is executable Python in the module
call protocol.
train() and eval() are another commonly misplaced boundary. They set the
recursive training flag. Modules such as Dropout and BatchNorm inspect that
flag and change behavior. They do not enable or disable autograd. After
model.eval(), I ran the module with a gradient-requiring input:
eval sets dropout.training=False; output still requires_grad=True
Conversely, torch.no_grad() can suppress gradient recording while the module
remains in training mode and Dropout remains stochastic. Inference usually
requires both the desired module mode and the desired gradient mode.
state_dict() returns a name-to-state mapping, not a self-contained clone of
the module. In the probe, state["up.weight"] had the same data pointer as the
live parameter. Mutating the module before cloning the state would mutate what
looked like a snapshot. load_state_dict() copies or assigns values into an
already constructed module according to its options and reports missing and
unexpected keys. The Python class definition that created the structure does
not arrive with it.
Buffers solve a precise ownership problem. BatchNorm running statistics, a
position index, a mask, or the temperature scalar may need to follow
module.to(device) and appear in checkpoints without receiving optimizer
updates. Setting persistent=False keeps a buffer in traversal and device
moves while omitting it from state_dict. An ordinary attribute does neither.
Module conversion methods recurse through registered parameters and buffers.
model.to("cuda") asks each relevant tensor to move or cast. Depending on the
operation and configuration, this can replace underlying tensor objects or
storage. Optimizers refer to parameters, and external code may hold aliases to
old tensors. Treat conversion as a state transition, not a decorative device
label.
The meta device separates module structure from storage allocation. A large
module can be constructed with meta tensors to obtain shapes and register the
tree without allocating payload bytes. to_empty can materialize uninitialized
storage on a target device, after which a state dictionary can populate it.
Lazy modules similarly defer dimensions until input establishes them. These
paths are valuable for models too large to instantiate twice, but they also
make “constructor completed” weaker than “all parameters contain usable
values.”
The module abstraction therefore owns names, hierarchy, call interception, mode flags, and persistent state classification. Operator implementation, autograd scheduling, device execution, optimizer state, and compilation remain elsewhere. Those systems discover and reinterpret the module through its registered tensors and executed calls.
backward is a concurrent graph execution
The earlier autograd chapter derived reverse
mode from the chain rule. Here the question is implementation ownership. What
did the framework retain between loss = ... and loss.backward()?
During forward, autograd wrappers create a graph of Node objects when at
least one input and the current gradient mode require it. Output tensors carry
an edge into that graph through grad_fn. Leaf parameters do not have a
producer node in the same sense; an AccumulateGrad node receives
contributions destined for their .grad fields.
Walking loss.grad_fn.next_functions for the tiny block found sixteen node
objects. The beginning was:
NllLossBackward0
LogSoftmaxBackward0
DivBackward0
AddmmBackward0
AccumulateGrad
MulBackward0
TBackward0
GeluBackward0
AccumulateGrad
NativeLayerNormBackward0
AddmmBackward0
AccumulateGrad
This is not the forward operator graph in reverse textual order. Some forward operators decompose, some backward formulas introduce new operations, a node may have several incoming gradient contributions, and views require special handling. Generated derivative definitions specify which values a backward rule needs and how gradients map to differentiable inputs.
Using saved_tensors_hooks, I counted fifteen saved tensors carrying 996
payload bytes for one five-example run. That is a measurement of tensors
passing through the hook for this graph, not the process’s total autograd
memory. Node objects, metadata, allocator rounding, temporary buffers,
parameters, gradients, and kernels’ workspaces are outside that sum. The hook
also changes the save/load path, so it is instrumentation with semantics, not
a free observer.
Saved tensors explain the version counter from the first section. I computed:
leaf = torch.tensor([1.0, 2.0], requires_grad=True)
saved_output = leaf.exp()
with torch.no_grad():
saved_output.add_(1)
saved_output.sum().backward()
The backward rule for exp needs its forward output. The in-place add raised
the output’s version from zero to one. Backward rejected it:
one of the variables needed for gradient computation has been modified by an
inplace operation ... is at version 1; expected version 0 instead
no_grad prevented a new history edge for the add. It did not grant permission
to invalidate an old graph. Mutation safety and gradient recording are
different concerns.
When backward() begins, the engine constructs a GraphTask describing this
execution. The pinned
graph_task.h
tracks outstanding tasks, dependencies, not-yet-ready inputs, captured
variables, errors, and completion state. The pinned
engine.cpp
queues ready nodes and calls evaluate_function as dependencies become
satisfied.
This is a dataflow scheduler. If two branches no longer depend on each other, their nodes can become ready independently. Device-specific engine threads and stream semantics complicate actual overlap. The graph describes dependency, not a promise that arbitrary backward nodes execute simultaneously.
Gradient accumulation is deliberate. If a parameter contributes to the loss
along two paths, the two vector-Jacobian products must sum. If backward() is
called twice without clearing leaf gradients, the second call adds to the
first. That behavior supports microbatch accumulation and also creates a
common stale-gradient bug.
optimizer.zero_grad(set_to_none=True) does not fill each gradient allocation
with zero. It releases the parameter’s reference by setting .grad = None.
The next backward can allocate or assign a fresh gradient. set_to_none=False
retains buffers and zeros them. Optimizers can distinguish None from a
present zero tensor, so the choice may affect whether a parameter is treated
as having participated in the step.
By default, a completed backward frees saved intermediates that are no longer
needed. retain_graph=True keeps enough state for another traversal.
create_graph=True asks backward operations themselves to be recorded so
higher-order derivatives can be taken. Those options solve different
problems. Retaining a graph is about reusing one derivative execution;
creating a graph is about differentiating the derivative computation.
torch.autograd.grad returns selected gradients without necessarily
accumulating into every leaf .grad. A non-scalar output needs an explicit
upstream tensor because reverse mode computes a vector-Jacobian product, not
an implicit full Jacobian. Forward-mode dual tensors and torch.func.jvp
propagate Jacobian-vector products in the other direction. Choosing a mode is
an arithmetic cost decision based on input and output dimensions, not a
preference for one API spelling.
Gradient mode is thread-local. torch.no_grad() suppresses recording of
ordinary operations in its region, but resulting tensors still have version
counters and can later participate in recorded work. torch.inference_mode()
removes more autograd overhead and creates tensors that are more restricted
for later autograd use. The probe could read the no-grad result’s version
counter but received:
Inference tensors do not track version counter.
detach() operates on a tensor edge. no_grad() and inference_mode() alter
ambient execution. requires_grad_(False) changes a tensor’s participation.
eval() changes module behavior. These four controls are related only because
they often appear near inference code. Substituting one for another is a
semantic change.
Custom backward logic can enter through torch.autograd.Function or through
an operator’s autograd registration. A Function supplies forward and
backward methods with a context for saved tensors and non-tensor metadata.
That route is appropriate when PyTorch cannot derive the intended derivative from
visible ATen operations. It also becomes responsible for mutation declarations,
saved-value lifetime, higher-order behavior, transform compatibility, and
correct gradients. gradcheck uses numerical perturbations to test a local
contract; it does not prove global numerical stability.
The autograd engine owns dependency execution and gradient delivery. It does not own the optimizer update. That separation becomes visible on the next line.
the optimizer has a second state graph
AdamW(model.parameters(), lr=0.01) receives an iterable of parameter objects,
not the module, its forward method, or the autograd graph. The
optimizer keeps parameter groups, each containing parameter references and
options such as learning rate, weight decay, and numerical tolerances. It also
keeps a state mapping keyed by parameter identity.
Immediately after construction, the probe’s AdamW state mapping was empty. After one backward and one step, it had six owners, one for each registered parameter. A representative state contained:
step
exp_avg
exp_avg_sq
The exponential moving average and squared moving average have the same shape as their parameter. They are allocated lazily when that parameter first participates in a step. This is why constructing an optimizer does not reveal its eventual memory footprint. Optimizer choice, dtype policy, sharding, and which parameters receive gradients decide the state payload.
For a parameter , gradient , first moment , and second moment , a simplified Adam update contains:
Bias correction and the final division follow. AdamW applies weight decay as a
decoupled parameter update rather than adding an L2 term to the gradient.
These equations describe the algorithm. They do not say whether PyTorch runs
one operator per parameter, batches tensors into foreach operations, or uses
a fused implementation. The optimizer selects among scalar-loop, foreach,
fused, capturable, differentiable, and device-specific paths under option and
backend constraints.
That distinction matters at scale. A Python loop launches many operations and holds little extra batch metadata. A foreach path groups lists of tensors and can reduce interpreter and launch overhead while increasing peak temporary memory. A fused path combines more update work inside fewer kernels when the device and dtype support it. The same AdamW equations can have different allocation, launch, and graph-capture properties.
Ordinary step() updates parameters without recording those updates into the
autograd graph. A differentiable optimizer asks for the update itself to be
differentiable, useful for meta-learning but more expensive and subject to
in-place constraints. capturable=True changes where step-related state must
live so a device graph can replay it without host decisions. These flags are
not generic speed switches. They alter the legal execution context.
Some optimizers accept a closure. The closure recomputes the loss and
gradients, allowing algorithms such as LBFGS to evaluate the objective
multiple times per optimizer step. Code that assumes one forward and backward
per step() is therefore not an optimizer-independent training loop.
A learning-rate scheduler mutates optimizer group options according to its
own state machine. Gradient clipping reads and changes gradients before the
update. Automatic mixed-precision scaling may unscale gradients, test them for
non-finite values, and skip the optimizer update. Distributed wrappers may
finish gradient communication before step() observes those fields. PyTorch
keeps these as composable objects rather than one monolithic trainer, so the
caller owns their order.
The order is observable:
with torch.autocast("cuda", dtype=torch.float16):
loss = model(x).square().mean()
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
This is a CUDA-shaped example, not one I executed on the CPU host. Clipping
before unscale would clip scaled values rather than the intended gradients.
Calling optimizer.step() directly would bypass the scaler’s non-finite check.
Zeroing before the update would discard the values. Each component has a
small contract; the training loop is the protocol that composes them.
An optimizer state dictionary does not save parameter objects. It saves state entries associated with parameter identifiers plus the parameter-group configuration. Loading matches those entries to the current optimizer’s groups. If the module’s parameters were reordered, replaced, frozen, or split into different groups, a superficially successful load can attach state to the wrong conceptual role unless names and construction are controlled.
Weight sharing creates another edge. If two module paths use one Parameter,
that object should normally appear once in an optimizer. Passing duplicate
references can produce duplicate updates or a warning, depending on how the
groups are built. The module tree supplies names. Parameter identity supplies
update ownership.
The optimizer therefore owns long-lived algorithm state and parameter mutation. Autograd produces contributions. The module exposes parameters. The caller decides when the gradient set is complete and whether this attempted step is allowed to commit.
execution is contextual
modes change meaning without changing source
The same Python line can run under several ambient modes:
y = torch.nn.functional.linear(x, weight, bias)
Gradient mode, inference mode, autocast, deterministic algorithms, active dispatch modes, current device, current stream, profiler scopes, anomaly detection, and compiler capture can all change what happens. Most are thread-local or context-managed. None appears in the function signature.
This ambient state is powerful because a library can add one context around a large body of code. It is dangerous because a helper’s semantics may depend on its caller.
gradient and inference modes
torch.enable_grad, torch.no_grad, and torch.set_grad_enabled control
reverse-mode recording. They do not set a module’s training flag. Factory
functions that accept requires_grad have documented exceptions to no-grad
behavior, which is another reason to test the resulting tensor rather than
reason from indentation alone.
Inference mode removes additional autograd bookkeeping and can enable faster
paths. Its outputs are deliberately less reusable in later recorded
computation. It is suitable when the region’s tensors will remain outside
autograd. It is not a universally stronger no_grad.
Anomaly detection asks autograd to retain traceback information and check backward outputs for invalid values. That debugging mode has substantial overhead, not a production correctness proof. A missing anomaly does not establish that gradients are numerically useful.
autocast and scaling
Autocast is a dispatch interpretation, not “convert the model to half.” It
applies operator-specific dtype policies inside a region. Matrix multiplications
may use a lower precision while reductions or numerically sensitive operations
remain or promote to float32. Operations with explicit output dtypes,
in-place variants, and out= forms have different eligibility.
The CPU probe entered:
with torch.autocast("cpu", dtype=torch.bfloat16):
output = model(x)
loss = output.float().square().mean()
It observed a bfloat16 model output and a float32 loss. The weights remained float32. Autocast selected operation behavior; it did not rewrite registered parameter storage.
Gradient scaling solves a different problem, primarily for float16’s narrow small-magnitude range. Multiplying the loss makes intermediate gradients larger, backward propagates the scaled values, and the scaler unscales before the optimizer update. It tracks whether infinities or NaNs appeared and adapts the scale. Bfloat16 has the same exponent width as float32 and usually does not need this underflow remedy, though autocast policy remains useful.
random state
torch.manual_seed changes framework generator state. A torch.Generator is
an explicit RNG state object that can be passed to supporting operations. It
is not merely a seed integer. Algorithms advance its state as they draw
values.
I created a private CPU generator, saved its state after one draw, generated four more values, restored the state, and generated again. The second sequence replayed bit for bit, while the global generator state remained unchanged:
private generator first=[0.23644638, 0.22661799, 0.80053020, 0.16918766]
restored next draw exactly=True
A seed initializes a stream. Changing operation order, device, thread
schedule, kernel, release, or worker count may consume
the same stream in the same way. Dropout’s mask is part of program state even
though it is not stored in state_dict.
DataLoader workers receive their own seeds. Distributed ranks usually need a deliberate policy: identical initialization where parameters must match, distinct data or dropout streams where samples should differ, and reproducible derivation from a run seed, rank, worker, and epoch. Reusing one seed everywhere can accidentally correlate work.
determinism and numerical policy
torch.use_deterministic_algorithms(True) asks supported operations to use
deterministic implementations and to error when only known nondeterministic
behavior is available. Backend flags, environment variables, library
versions, and workspace configuration can add requirements. Deterministic
does not mean exact real arithmetic, cross-device equality, or equality across
releases.
Floating-point reduction is order-sensitive because addition is not associative in finite precision. A batched implementation need not use the same reduction order as a loop over examples. A compiled fusion need not round at the same boundaries as eager execution. TensorFloat32, reduced-precision accumulation, denormal handling, and library algorithm selection can alter results within documented contracts.
PyTorch’s
numerical accuracy note
explicitly rejects a general bitwise-identity guarantee across platforms and
releases. The
reproducibility note
describes controls for a fixed environment. Those are different goals.
Reproducibility helps compare runs. Accuracy asks whether the computation is a
useful approximation. Determinism asks whether repeated execution chooses the
same path.
Modes are easiest to reason about as scoped interpreters over operator calls. Whenever one is enabled, ask which operators it intercepts, what state it owns, whether it redispatches, and which invariants survive after the context exits.
a device is a memory and execution domain
torch.device("cpu") is not a request to run on a vague host abstraction. It
identifies the backend and optional index participating in tensor allocation
and operator dispatch. A CPU tensor and a CUDA tensor with equal values are
different storage objects in different execution domains. to(device) is a
copy or conversion operation unless the tensor already satisfies the request
and PyTorch can return the original handle.
CPU execution
CPU operators can execute directly on the calling thread, use PyTorch’s
intra-operation thread pool, invoke a library with its own workers, or compose
these behaviors. Inter-operation parallelism schedules independent operations
in graph-execution contexts. The probe wheel reported four intra-op threads
and ten inter-op threads. Those numbers are configuration, not evidence that a
particular addmm used fourteen threads.
BLAS and OpenMP runtimes can have their own environment controls. Oversubscription
occurs when several process, operator, and library pools each assume they own
all cores. A benchmark must report PyTorch thread settings, process count, and
relevant library variables. torch.set_num_threads changes a framework
control; it does not retroactively prove native affinity or eliminate every
other pool.
The CPU allocator returns host-accessible memory. Tensor operations generally
complete their CPU work before returning, but asynchronous libraries,
background data workers, and distributed Work objects create exceptions at
higher layers. “CPU is synchronous” is too broad to be useful.
accelerator execution
CUDA operations are usually enqueued on a CUDA stream and return control to
the host before the device finishes. The current stream supplies ordering:
operations in one stream execute in issue order, while separate streams need
events or explicit waits to establish dependencies. A host read such as
.item(), a device-to-host copy without an asynchronous contract, or an
explicit synchronize can expose completion.
PyTorch tracks a current device and current stream per thread. Once a tensor
is allocated, its device travels with it. The
CUDA semantics note
documents which cross-device copies are allowed and how stream contexts affect
work. The Python source line does not tell us whether the host waited.
The CUDA caching allocator usually retains freed blocks for reuse rather than
returning each one immediately to the driver. memory_allocated tracks live
tensor allocations as PyTorch accounts them. memory_reserved includes
allocator-held segments. A process monitor may report still another number.
Deleting a tensor releases a framework ownership edge; it does not promise an
instant fall in device-visible reservation.
Allocator correctness is stream-sensitive. Memory used by an asynchronous kernel cannot be recycled for an unrelated allocation before the kernel is done. PyTorch records stream usage and events to delay reuse. Custom extensions that use tensors on non-current streams must honor the corresponding recording contract or create a use-after-reuse race invisible to Python reference counts.
Pinned host memory permits DMA-friendly transfers and enables genuinely
asynchronous copies in supported paths. non_blocking=True expresses
permission to avoid a host wait when the source, destination, backend, and
stream relationship make that safe. The keyword alone does not make pageable
memory, arbitrary backends, or a later host use asynchronous.
CUDA Graph capture records a fixed sequence of device work for replay with lower host launch overhead. PyTorch must stabilize memory addresses, control stream capture, and keep inputs or static buffers compatible with replay. Dynamic allocation, CPU decisions, synchronization, and changing shapes can break capture. The result is a replay contract layered under or beside compiler capture, not another name for an FX graph.
MPS, XPU, MTIA, HIP/ROCm, and private accelerators implement different
backends through the common tensor and dispatcher surfaces. Their operation
coverage, dtype support, asynchronous behavior, allocators, and fallbacks are
not implied by CUDA semantics. PrivateUse1 lets an out-of-tree backend
register a device family without assigning a permanent in-tree dispatch key.
It still needs allocator, device guard, generator, operator, autograd,
serialization, and extension integration to behave like a full backend.
The permitted MPS control moved one (3, 4) input, (2, 4) weight, and
two-element bias to the Apple GPU, ran affine, ReLU, mean-square loss, and
backward, synchronized explicitly, and copied output and gradients back. The
maximum absolute differences from the CPU float32 control were:
output 5.9604645e-08
input gradient 1.4901161e-08
weight gradient 1.8626451e-08
bias gradient 2.2351742e-08
All passed rtol=1e-5, atol=1e-6. MPS current allocated memory increased from
768 to 1,792 bytes while the small device tensors lived; driver-allocated
memory remained 8,880,128 bytes in that observation. Those values illustrate
the distinction between live framework payload and a backend’s larger driver
allocation. They are not a stable allocator baseline or a performance
measurement.
The meta device has no payload execution domain. Fake tensors can claim an apparent device while using fake implementations to propagate metadata. These are not toy devices. They let compilers and large-model loaders answer shape and placement questions before committing memory.
Device, dtype, and layout form a joint dispatch problem. Moving a float32 strided CPU tensor to bfloat16 channels-last CUDA changes at least three axes. An operator may be available for each axis separately yet unavailable for their combination. The error belongs to the missing composite contract, not to “GPU support” in general.
function transforms add dimensions to programs
Ordinary eager code maps input tensors to output tensors. torch.func
transforms map a function to another function:
grad(f)returns a function computing gradients off;vjpandjvpexpose reverse and forward products;jacrevandjacfwdconstruct Jacobians through those products;vmap(f)adds a batch interpretation without writing the outer loop;functional_callexecutes a module with supplied parameter and buffer mappings;functionalizeremoves intermediate mutations and, optionally, views from an operator program while preserving observable results.
These features were once developed as functorch and now live under
torch.func, with implementation pieces still named _functorch. Their
composition works because they participate in the dispatcher and operate on
ATen-level semantics.
I expressed the tiny affine-ReLU module as a pure loss over an explicit parameter dictionary:
def loss(params, buffers, x, target):
prediction = torch.func.functional_call(
model, (params, buffers), (x,)
)
return torch.nn.functional.mse_loss(prediction, target)
batch_grad = torch.func.grad(loss)(params, buffers, x, target)
Then I computed one gradient per example:
per_example = torch.func.vmap(
torch.func.grad(one_loss),
in_dims=(None, None, 0, 0),
)(params, buffers, x, target)
The weight result had shape (5, 4, 3): five examples, followed by the
parameter’s (4, 3) shape. The ordinary batch gradient had (4, 3).
For mean-squared error with mean reduction, the mean of the five per-example
gradients matched the batch gradient to the probe tolerance.
vmap did not literally promise five calls to Python. It introduced a batched
tensor interpretation and used batching rules for operations. A rule can move
the logical batch dimension, call a batched kernel, decompose an operation, or
report that the operator is unsupported. Random operations need an explicit
randomness policy because five independent draws, one shared draw, and an
error are all defensible semantics.
Functionalization exposes how mutation becomes graph data. The probe’s source mutated through a view:
def mutating_program(x):
view = x.view(-1)
view.add_(2)
return x.square()
After functionalize(..., remove="mutations_and_views"), an FX trace contained
view_copy, functional add, pow, and a final copy_ back to the input.
The intermediate in-place add disappeared. The final copy remained because
mutation of the function input is observable to the caller and must be
replayed at the boundary.
That last copy prevents a common overstatement. Functionalization does not declare mutation harmless or remove all writes. It converts internal aliasing and mutation into a representation that downstream systems can analyze, then repairs required input effects.
FakeTensor is another function-transform tool. Under FakeTensorMode, the
probe ran transpose, matrix multiplication, bias addition, and ReLU on a fake
(7, 3) input. It produced a FakeTensor of shape (7, 4) without using real
input payload. Fake implementations, symbolic shapes, and alias rules supply
the metadata.
A fake kernel cannot generally branch on tensor data because no data exists. An operator whose output shape depends on values, such as a nonzero-index query, needs a way to represent an unbacked symbolic size or a constrained operator form. A custom operator that only registers a real CUDA body may work eagerly and fail under export because the fake interpreter cannot infer its output.
Tensor subclasses and modes use related machinery but solve different
problems. A subclass attaches behavior to tensor instances. A mode applies to
all intercepted operations in a dynamic scope. DTensor is a subclass carrying
distributed placement. FakeTensor is a subclass carrying simulated metadata.
A logging mode can observe ordinary tensors without changing their class.
__torch_function__ intercepts public Python API functions at a higher level;
__torch_dispatch__ sees dispatcher operators closer to ATen. Supporting one
does not automatically make an extension correct under every transform.
Function transforms work when operators publish enough semantics: schema, aliasing, mutation, batching behavior, fake behavior, and differentiation. This is why the dispatcher is not plumbing beneath the “real” framework. It is the surface on which PyTorch’s programmable meanings compose.
graphs move decisions across time
compilation changes when decisions are made
The compiler chapter follows one expression through Dynamo, AOTAutograd, and Inductor. Here I want the system boundary: how compiled execution remains PyTorch rather than replacing it.
Default eager execution reaches the dispatcher as Python encounters each
operation. torch.compile wraps a callable and attempts to capture regions of
Python execution into graphs that a backend can optimize. Its default backend
is Inductor.
The three major owners are:
- Dynamo intercepts CPython frame evaluation, symbolically executes bytecode, records FX graphs, and emits guards for assumptions.
- AOTAutograd functionalizes relevant work, analyzes mutation and aliasing, and can trace forward and backward graphs ahead of execution.
- Inductor lowers FX/ATen graphs into scheduled loops, generated kernels, and external-library calls for supported backends.
These are not three mandatory file formats. They are compiler stages with Python and native implementation spread across the repository. Backend choice can stop after capture, hand the graph to a different compiler, or use eager ATen as a debugging backend.
I compiled:
def plain(x, weight, bias):
return torch.relu(x @ weight.t() + bias)
compiled = torch.compile(plain, fullgraph=True)
The 2.13 CPU result matched eager. torch._dynamo.explain reported one graph,
zero graph breaks, three high-level operations, and twelve guards. Guard
counts are internal and version-sensitive. The useful fact is structural:
compiled code remained conditional on assumptions about the function,
tensors, global state, and environment.
A guard can check shape, stride, dtype, device, object identity, a global value, a module property, or a Python relationship observed during capture. When a later call fails a guard, Dynamo searches another cached specialization or recompiles. Too many specializations can consume compile time and cache memory; after configured limits, execution may fall back. “The model is compiled” is not a permanent Boolean property. The callable owns guarded programs associated with code objects.
Shapes are a central example. With dynamic=None, PyTorch commonly starts
specialized and may generalize dimensions after observing changes. dynamic
symbols carry constraints and participate in guards. Some dimensions remain
specialized because operations or optimizations require them. A dynamic batch
does not imply dynamic rank, dtype, arbitrary stride, or unrestricted control
flow.
Graph breaks end one captured region and resume ordinary Python before a later
region. Data-dependent .item(), unsupported bytecode, a side effect, or an
opaque call may cause one. fullgraph=True converts a break into an error for
the requested frame, which is useful when partial capture would hide a
performance boundary. One graph can still contain many kernels.
Inside a captured graph, decompositions rewrite higher-level operators into a smaller operator set. Functionalization exposes mutations. AOTAutograd decides which forward values backward needs and can partition recomputation. Inductor’s scheduler reasons about dependencies, layouts, fusion, and buffer lifetime. Matrix multiplications and convolutions may remain external library calls while adjacent pointwise work becomes generated code. One compiled region can contain several kernels and library calls.
Inductor uses different code-generation paths by backend. CPU work can become C++ with vector and OpenMP constructs plus external ATen calls. Supported GPU work often lowers through Triton, vendor libraries, templates, or native kernels. Autotuning may benchmark candidates for a concrete shape. Generated artifacts are cached under keys influenced by code and configuration.
Compilation therefore introduces two time scales: compile time and steady execution time. A speedup after warmup can still be a loss for a short job. Dynamic inputs can trigger recompilation later. Autotuning can add startup variance. Cache reuse can hide cold cost in one measurement and expose it in another. Correct benchmarking separates cold, warm, and recompile cases.
Correctness is also layered. Eager and compiled results can differ because of
a compiler bug, a custom operator with a false schema, undefined input
aliasing, floating-point reassociation, an uninitialized read, or merely a
different documented numerical path. A useful comparison covers shapes,
dtypes, strides, aliasing, mutations, exceptional values, forward outputs,
gradients, and repeated calls. One random assert_close is a start.
Compiler diagnostics live at several levels. TORCH_LOGS categories can show
guards, recompiles, graph breaks, dynamic-shape reasoning, and generated
artifacts. torch._dynamo.explain summarizes capture. Inductor trace options
write intermediate graphs and code. TORCH_TRACE plus tlparse can package
distributed compiler events for analysis. A profiler answers runtime cost; it
does not explain why Dynamo broke a graph.
Compiled autograd can capture portions of backward that ordinary AOTAutograd could not see at forward compile time, including some hook-driven behavior. This adds another capture boundary without replacing derivative formulas. Hooks, custom Functions, distributed reducers, and mutation make backward capture substantially more stateful than compiling a pure forward function.
CUDA Graph integration may replay the device sequence selected by compiled code. Dynamo guards still decide whether the Python-level specialization is valid; CUDA Graph constraints decide whether the recorded device work is replayable. Calling both mechanisms “graphs” hides two different contracts.
Compilation succeeds because PyTorch’s lower layers remain available. Captured graphs name ATen operators. Decompositions use dispatcher semantics. Generated wrappers allocate tensors and call native kernels. Unsupported regions return to eager PyTorch. The compiler moves decisions earlier and fuses across operator boundaries, but tensor, schema, device, autograd, and allocator contracts still hold it together.
export is a contract, not a faster call
torch.compile optimizes execution inside a live Python process and may
recompile when guards fail. torch.export produces an ExportedProgram meant
to cross a stronger boundary. It contains a normalized tensor graph, graph
signature, parameters and buffers, constants, input and output structure, and
range constraints needed for sound execution.
The probe exported the affine-ReLU module with a batch dimension constrained from one through eight:
batch = torch.export.Dim("batch", min=1, max=8)
program = torch.export.export(
model.eval(),
(torch.randn(2, 3),),
dynamic_shapes=({0: batch},),
)
Its graph called:
aten.t.default
aten.matmul.default
aten.add.Tensor
aten.relu.default
Its range constraints printed one symbolic variable with [1, 8]. Batch six
ran. Batch nine did not.
I initially caught RuntimeError in the probe because that is the generic
exception I expected from a runtime guard. PyTorch raised AssertionError:
Guard failed: x.size()[0] <= 8
The exception class was not the contract. The rejected shape was. The probe now accepts the observed guard exception types and asserts the constraint. Keeping this correction avoids turning a local implementation detail into a portable API promise.
Export treats tensor data as dynamic while specializing or symbolizing
metadata according to the supplied contract. Python integers, booleans,
strings, container structure, and module attributes are commonly specialized
unless represented through supported symbolic forms. A Python branch over
dynamic tensor data cannot simply be burned in; it needs an explicit
traceable control-flow operator such as torch.cond, or export must reject it.
Parameters and buffers become graph state with signatures describing their roles. Mutations can be represented as explicit outputs that the runtime applies to state. The graph is functionalized toward a defined ATen operator set. Custom operators may remain as named boundaries if their schemas and fake behavior make the export sound.
An ExportedProgram is not a Python Module pickled with an example. It is
also not machine code. A downstream runtime or compiler consumes the graph and
state under its own supported operator, dtype, layout, and device set. Export
can succeed while a chosen target rejects an operator. Target support is a
separate lowering question.
The 2.13 API defaults and details are versioned. The current
torch.export API reference
defines the soundness guarantee around recorded assumptions. The
programming model
explains static versus dynamic values and control flow. Treating an export as
valid for untested inputs outside that contract discards the guarantee it was
designed to provide.
TorchScript historically offered scripting and tracing for graph capture and
deployment. PyTorch 2.10 deprecated TorchScript in favor of torch.export for
new work. Old TorchScript artifacts and APIs do not vanish, but they represent
a different intermediate representation and compatibility story. An article
about current PyTorch should not teach torch.jit.trace as a synonym for
export.
ONNX export, ExecuTorch packaging, AOTInductor packages, and backend-specific deployment systems begin at this boundary. They are not all PyTorch execution modes. This chapter stops after PyTorch has produced and validated its artifact.
DataLoader is a small distributed system
Training code often treats this loop as plumbing:
for x, target in loader:
...
By the time a batch reaches the loop, PyTorch may have sampled indices, serialized tasks to worker processes, read files, decoded examples, collated tensors, reordered out-of-order worker results, copied CPU storage into pinned memory, and queued a transfer to an accelerator. The model has not run, but there is already concurrency, buffering, backpressure, random state, failure, and ownership.
Dataset answers how an example is obtained. A map-style dataset implements a
key-to-example relationship, conventionally __getitem__ and __len__.
An IterableDataset produces a stream and may have no cheap random access or
known length. This distinction changes where work division belongs.
For map-style input, a Sampler chooses keys. RandomSampler,
SequentialSampler, DistributedSampler, or a user sampler defines order and
partition. A BatchSampler groups keys. Workers receive index tasks and call
the dataset. The main iterator can preserve requested order even when workers
finish at different times.
For iterable input, each worker receives a replica of the dataset object and
calls its iterator. Unless the dataset consults get_worker_info() or is
configured by worker_init_fn, every replica may emit the same stream. My
probe made the mistake visible:
unsharded iterable sorted=[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5]
sharded iterable sorted=[0, 1, 2, 3, 4, 5]
The first loader used two workers with no sharding. Each worker correctly iterated zero through five, so the combined result duplicated every element. The second began worker zero at zero and worker one at one, both advancing by the worker count. PyTorch cannot infer a safe partition for an arbitrary stream with external side effects or unknown ordering.
The map-style control produced a different result:
map dataset indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
workers=[0, 1] worker_processes=2
Both worker processes contributed, but output followed the sequential sampler. That does not mean processing completed in order. The iterator tracked task identities and emitted according to its ordering policy.
The pinned
dataloader.py
has separate single-process and multiprocessing iterators. The multiprocess
path owns per-worker index queues, a worker-result queue, task bookkeeping,
worker-status flags, and shutdown logic. With pinning enabled, another thread
can move received tensors into pinned host allocations before the main loop
observes them.
num_workers=0 runs dataset and collation code in the caller. Exceptions have
direct tracebacks, Python objects need not cross a process boundary, and no
prefetch pipeline exists. A positive worker count changes the program:
- Worker processes receive serialized dataset, function, and task state according to the platform’s process-start method.
- The iterator enqueues up to a prefetch budget of index tasks.
- Workers fetch and collate examples.
- Tensor storage crosses a multiprocessing boundary, often using shared memory rather than copying every byte through a pipe.
- The main iterator buffers or emits results and schedules replacement tasks.
- End-of-data, exceptions, timeouts, dropped consumers, and process exit must shut workers down.
The probe first failed inside the repository sandbox. Workers began, but
serializing tensor storage into the result queue tried to start
torch_shm_manager, which the sandbox denied:
torch_shm_manager ... Operation not permitted
Running the same probe with explicit permission passed. That failure is not a PyTorch numerical result. It identifies an operating-system authority the multiworker path needs. A container with a tiny shared-memory mount, a sandbox that forbids the manager, exhausted file descriptors, or a killed worker can break a DataLoader while the dataset code is correct.
On Unix, a fork-based start method initially shares parent pages through copy-on-write. Touching Python objects in workers can make those pages private, so a large parent-resident dataset may multiply memory across workers. Spawn starts a fresh interpreter and serializes needed objects, avoiding some fork hazards while adding startup and pickling constraints. macOS and Windows commonly use spawn. Dataset classes, collation functions, and worker initializers then need importable top-level definitions rather than lambdas or local closures.
persistent_workers=True keeps workers and their dataset replicas alive after
one pass through the loader. That avoids repeated startup and can also retain
file handles, caches, RNG state, and mutations into the next epoch.
prefetch_factor multiplies by the worker count, so increasing workers can
increase both concurrency and the number of resident batches. pin_memory
adds allocations and a pipeline stage. These controls spend memory to hide
latency; they do not make a slow source free.
drop_last has different consequences for map and iterable workers. With an
iterable dataset, a partial final batch can exist per worker replica. A
distributed sampler needs its epoch advanced, usually with set_epoch, so
each epoch receives a deterministic but different shuffle. Dataset length is
an estimate for many streams and should not be treated as proof of exact batch
count.
Collation is an ownership boundary too. The default collator stacks compatible tensors, converts supported NumPy arrays and numbers, and recursively preserves container structure. Variable-size examples need padding, packing, nested representation, or a list. A custom collator can accidentally copy data several times, retain large source objects, or produce non-pinnable containers.
The data path should therefore be measured independently from the model. Track queue occupancy, batch wait time, worker CPU, storage throughput, decode cost, host memory, pinned memory, transfer overlap, and consumer utilization. A GPU gap can begin in a sampler, filesystem, decoder, queue, copy, synchronization, or model. Calling all of it “DataLoader overhead” removes the owner needed to fix it.
distributed PyTorch composes processes, groups, and tensor layouts
PyTorch distributed begins below DDP. A store lets processes exchange
rendezvous information. A process group assigns ranks and a communication
backend to a participating set. Point-to-point and collective APIs return
results and, for asynchronous forms, a Work handle representing completion.
The backend is part of semantics and performance. Gloo supports CPU
communication and selected device paths. NCCL targets NVIDIA GPU collectives.
Other builds can include MPI, UCC, XCCL, or backend integrations. Availability
depends on how PyTorch was built. Backend choice also affects what wait
means for host completion, stream ordering, and errors.
My only fresh distributed execution used two local CPU processes and Gloo. It first computed a one-row linear gradient independently on each rank:
rank-local gradients=[[4.0, 2.0], [-6.0, -18.0]]
arithmetic mean=[-1.0, -8.0]
Then it wrapped an identical model in
DistributedDataParallel(gradient_as_bucket_view=True). After backward:
rank 0 DDP gradient=[-1.0, -8.0]
rank 1 DDP gradient=[-1.0, -8.0]
An explicit asynchronous all-reduce of scalars one and two returned three on
both ranks after Work.wait(). This establishes two-process CPU behavior for
the probe. NCCL, overlap between real GPU backward and bucket communication,
and network performance remain unexecuted.
DDP’s reducer
DDP is a module wrapper around replicated model state. Construction verifies or synchronizes relevant state across the process group. Dividing the input batch remains the caller’s job. Each process usually receives its own shard through a distributed sampler.
The wrapper builds gradient buckets. Autograd hooks notify a C++ Reducer as
parameter gradients become ready. When a bucket is ready, the reducer launches
communication through its process group. After the reduction, it writes or
aliases the averaged results into parameter gradients and finalizes backward.
The pinned
reducer.cpp
contains Reducer::autograd_hook and Reducer::finalize_backward. The order in
which parameters become ready can influence bucket rebuilding and overlap.
Bucket size trades earlier communication and more calls against later,
larger transfers.
With gradient_as_bucket_view=True, the probe’s parameter .grad had a
non-null _base, showing that it was a view into bucket storage. This can
avoid a gradient-to-bucket copy and means code must respect view restrictions.
Detaching the gradient in place, changing its layout, or replacing it can
break assumptions. The option changes aliasing, not the all-reduce equation.
DDP averages gradients under its ordinary reduction contract. If each rank’s local loss is already a mean over an equal local batch, averaging rank gradients equals the global mean. Unequal local batch sizes require weighting; a blind rank mean gives each rank equal influence, not each example.
no_sync() accumulates local gradients while delaying DDP synchronization.
The final synchronized backward reduces the accumulated fields. The forward
must also be inside the context when required by the wrapper’s protocol.
This is useful for microbatches and changes peak gradient values and
communication frequency.
Unused parameters complicate reducer completion. If a forward path omits a
parameter, its hook may never fire. find_unused_parameters=True traverses
the autograd graph to mark such parameters, adding overhead. static_graph
can exploit a fixed used-parameter pattern. Incorrectly declaring a dynamic
graph static risks hangs or wrong synchronization.
A communication hook can replace a bucket’s default reduction with compression, hierarchical communication, gossip, or another future-returning protocol. PyTorch cannot guarantee convergence for an arbitrary hook. The hook owns division, error feedback, dtype conversion, and completion semantics that the default reducer would otherwise supply.
DDP also assumes corresponding collective order across ranks. If control flow causes rank zero to reduce bucket A while rank one reduces bucket B, both may wait forever or corrupt protocol state. A Python exception on one rank can leave peers blocked in native communication. Timeouts and asynchronous error handling turn some hangs into errors; they do not make collective sequences transactional.
sharded state
Replicating a model makes each rank own all parameters, gradients, and ordinary
optimizer state. FSDP reduces that footprint by sharding state across a data
parallel group and gathering parameters when computation needs them. The
current composable fully_shard path, often called FSDP2, uses DTensor-based
parameter representation and hooks around module execution. Older
FullyShardedDataParallel remains a substantial API with different state and
wrapping behavior.
A sharded forward is a schedule:
- all-gather the parameter shards required by a module;
- execute its local tensor operations;
- release or reshard full parameters according to policy;
- repeat for later modules;
- during backward, gather as needed and reduce-scatter gradients.
Prefetching tries to overlap communication for the next module with current computation. Rate limiting prevents too many full parameter sets from living at once. CPU offload exchanges device capacity for transfer and host memory. Mixed precision can choose different parameter, reduction, and buffer dtypes. These policies couple module traversal, allocator lifetime, streams, and process groups.
DTensor expresses a logical tensor plus a DeviceMesh and one placement per
mesh dimension. The principal placements are:
Replicate, where each rank holds the logical value;Shard(dim), where ranks hold partitions of tensor dimensiondim;Partial, where local values need a reduction to become the logical result.
DTensor is a Tensor subclass. Its dispatch rules propagate layouts and insert redistribution collectives when an operator needs a different placement. A matrix multiplication over sharded dimensions can produce partial output; an elementwise operation over matching shards can stay local. The operator’s sharding rule is part of correctness.
DeviceMesh arranges global ranks into named dimensions. A two-dimensional
mesh can assign one dimension to data sharding and another to tensor
parallelism. Each slice corresponds to process groups. All ranks must agree on
the mesh; inconsistent construction can hang before a useful tensor error is
possible.
Tensor parallel APIs place module parameters and activations with DTensor rules. Pipeline parallel APIs divide modules into stages and schedule microbatches. Context or sequence parallel schemes shard activation dimensions. These parallelisms compose only when their mesh dimensions, layouts, collectives, autograd rules, and checkpoints agree. “Distributed” does not name one wrapper.
Distributed checkpointing must save logical state without requiring every rank to assemble it on one host. Planners assign tensor shards to files or storage objects and later reshard for a possibly different topology. The checkpoint’s key space, layout metadata, optimizer mapping, and commit protocol become durable distributed state. A successful write by one rank is not proof that a multi-rank checkpoint is complete.
Elastic launch and rendezvous can restart workers after failure, but ordinary DDP iteration is not fault tolerant in the database sense. Model parameters, optimizer state, sampler position, RNG streams, scaler state, and data-source offsets must return to a consistent checkpoint boundary. Retrying a side effecting batch without idempotence can duplicate external effects even if the weights recover.
The distributed overview
is best read as a namespace map: C10D communication, DDP replication, FSDP
sharding, DTensor layouts, device meshes, tensor and pipeline parallelism,
checkpointing, launch, and debugging. They share process groups and tensors.
They do not share one state machine.
crossing the framework boundary preserves obligations
a checkpoint is storage plus a trust decision
PyTorch recommends saving state dictionaries rather than pickling a live module object for ordinary model checkpoints. A module state dictionary contains parameters and persistent buffers. An optimizer state dictionary contains group configuration and per-parameter algorithm state. Training recovery may additionally require scheduler, gradient scaler, RNG, sampler, epoch, and application metadata.
torch.save uses a ZIP-based container for its modern format. The probe saved
the tiny module’s state and found:
.data/serialization_id
.format_version
.storage_alignment
byteorder
data.pkl
data/0
data/1
data/2
data/3
version
The exact private members are version-sensitive. Structurally, pickled metadata describes tensor objects while storage payloads occupy separate records. This lets PyTorch preserve sharing instead of writing a complete byte copy for every view.
I saved a dictionary containing a tensor and its tail view. After loading, both tensors referred to the same untyped storage and the tail retained storage offset three:
loaded aliases: same_storage=True tail_offset=3
That is correct semantics and can make a checkpoint unexpectedly large. Saving a five-element view of a one-gigabyte storage may serialize the backing storage, not merely the visible five elements. Cloning the view before saving trades alias preservation for a compact independent payload.
map_location controls where storages are restored. Loading a GPU checkpoint
directly to its recorded device can create an avoidable memory spike if the
caller intended to inspect or reshape it on CPU first. Memory-mapped loading
can defer page materialization for supported files and access patterns. These
are loading strategies, not changes to the logical values.
state_dict() is shallow. The probe cloned its values before the optimizer
step because the returned parameter tensors shared live storage. A robust
in-memory “best model” snapshot must clone tensors or serialize them before
training continues.
Serialization also has a security boundary. Python pickle can encode global
lookups and object construction with code execution. PyTorch’s
weights_only=True loader restricts construction to tensors, primitive
containers, and explicitly allowlisted types. The probe saved an ordinary
custom Python object and attempted a weights-only load:
weights-only rejected UnpicklingError: Weights only load failed.
Setting weights_only=False on an untrusted file is not a compatibility
workaround. It authorizes the pickle program. Allowlisting a class makes it
trusted for that load; it does not inspect the class for safety.
Weights-only loading narrows remote-code-execution exposure but is not a complete hostile-file sandbox. A file can request huge allocations, expensive shapes, or pathological sparse structures. Validate provenance, size, dtype, shape, key set, and application invariants before exposing loaded tensors to native kernels.
File integrity and transactional publication live above torch.save. A writer
can serialize to a temporary name, flush as required by the durability model,
and atomically publish a manifest or rename on a local filesystem. Object
storage needs its own immutable-object and manifest protocol. A .pt file
that opens is not proof that every rank wrote the intended generation.
Version compatibility is layered:
- Python class pickle compatibility matters when saving objects.
- State-dict names and shapes matter when reconstructing modules.
- Operator and layout support matters for exported programs.
- Optimizer state structure can change with options and implementation.
- Device availability matters for recorded locations.
- User metadata needs its own schema and migration.
strict=True on load_state_dict checks key agreement, not model meaning.
Equal names can still refer to a changed architecture. strict=False reports
missing and unexpected keys but can hide an accidental partial load if the
caller ignores the result.
Exported programs have their own save/load format and verifier. Inductor caches contain generated artifacts tied to compiler and environment assumptions. Distributed checkpoints describe sharded logical state. None should be silently treated as the ordinary state-dict format.
The current
serialization semantics note
documents view preservation, ZIP layout, state-dict practice, and the
weights-only threat boundary. The key design lesson is broader: a checkpoint
is not just bytes representing numbers. It combines storage aliases, a schema
connecting names to program state, and a decision about which loader is
allowed to construct what.
extensions must teach every interpreter they cross
PyTorch offers several extension levels because “custom behavior” can mean different things.
composing existing operations
The easiest extension is an ordinary Python function or nn.Module built from
supported PyTorch operators. Eager autograd can differentiate the composition,
vmap can use constituent batching rules, fake tensors can propagate through
constituent fake kernels, and compilers can see inside it. This path gives the
framework maximum information.
If Python overhead or fusion is the problem, torch.compile may optimize the
composition without inventing a new operator. If the desired semantics already
exist as ATen operations, an opaque custom boundary can make optimization
harder.
custom differentiation
torch.autograd.Function defines a forward and backward relationship. It is
useful for a mathematically custom derivative, a memory-saving rule, or native
code invisible to autograd. It must save the right tensors, mark mutation and
non-differentiable outputs, support or reject higher derivatives, and obey
device and dtype behavior.
Function is not an operator schema. Other subsystems may trace through it,
treat it as opaque, or need additional integration. Current guidance favors
registering genuine custom operators through torch.library and attaching
autograd behavior there when the operation must compose broadly.
custom operators
A custom operator has a stable qualified name such as asquare::affine and a
schema. The schema must state argument and return types plus alias and mutation
behavior. torch.library.custom_op creates such a boundary in Python;
TORCH_LIBRARY and TORCH_LIBRARY_IMPL provide C++ registration.
For the affine probe, a correct eager body was only the first registration. The operator also needed:
- a fake implementation for metadata propagation;
- an autograd formula and saved context;
- backend kernels for every claimed device;
- correct mutation and freshness declarations;
- batching support or a composition that
vmapunderstands; - autocast policy if lower-precision execution is intended;
- compiler and export visibility through its schema.
torch.library.opcheck exercises schema, autograd registration, fake behavior,
and compile interactions over supplied examples. gradcheck numerically tests
derivatives. Neither tests every shape, alias, dtype, device, exception, or
race. An operator test matrix should include noncontiguous views, empty
dimensions, broadcasting, exceptional values, repeated calls, gradients,
fake tensors, compiled calls, and each backend.
The mutation declaration is a promise to optimizers. Declaring
mutates_args=() while writing an input lets a compiler assume the input
remains unchanged. Returning an alias while promising a fresh tensor can let
storage lifetimes overlap incorrectly. A schema bug is not metadata
harmlessness; it can become silent wrong code.
tensor subclasses and modes
__torch_function__ lets Python types override high-level public torch
functions. __torch_dispatch__ operates at the dispatcher operator level and
is the foundation for deeper tensor subclasses. A subclass may wrap or extend
a physical tensor, attach distributed placement, record operations, enforce
units, or simulate storage.
A TorchDispatchMode applies interception dynamically without changing tensor
classes. The recorder probe used a mode. FakeTensorMode uses one to manage fake
execution. Debug and flop-counting tools can be modes too. Modes form a stack;
each must redispatch correctly to avoid recursion or bypassing another mode.
Subclass correctness is demanding. View creation, serialization, device conversion, autograd metadata, flatten/unflatten, wrapper aliasing, and operator coverage all matter. Falling back by unwrapping to a base tensor may erase the subclass invariant.
native extensions and backends
torch.utils.cpp_extension can compile C++ or CUDA sources against the
installed PyTorch. A C++ extension can call ATen directly, register operators,
and expose Python bindings. Build success ties it to compiler, C++ ABI,
platform, accelerator toolkit, and PyTorch binary compatibility.
The C++ frontend, often called LibTorch, provides tensors, autograd, modules, optimizers, serialization, and selected distributed functionality without Python. It shares native components with Python PyTorch but is not a guarantee that every Python feature has a stable one-to-one C++ API.
PyTorch now also documents a stable C++ ABI surface for a restricted set of
operator and registration needs. “Stable” applies to that declared surface,
not to arbitrary internal headers. Including TensorImpl.h from an extension
buys implementation access at the price of version coupling.
A full accelerator backend goes beyond a few custom kernels. It needs a device
type and guards, allocator, streams or execution queues, generators, events,
storage behavior, copy paths, serialization, autograd, autocast, distributed
integration, fake/meta kernels, and operator registrations. PrivateUse1
reserves a dispatch family for out-of-tree development, but it does not
generate these semantics automatically.
Compiler backends form another extension boundary. Dynamo passes an FX
GraphModule and example inputs to a backend callable. The backend can return
an optimized callable or reject. A production backend needs guards,
mutations, dynamic shapes, decompositions, errors, cache keys, and correctness
testing. Accepting an FX graph is much easier than preserving PyTorch for all
inputs admitted by the guards.
The extension hierarchy suggests a rule: remain at the highest layer that expresses the intended optimization. Compose operators before writing a custom one. Register a custom operator before reaching into dispatcher internals. Use public backend interfaces before private compiler modules. Move downward when evidence shows the higher layer cannot express the required semantics, not merely because native code looks more serious.
reading a running framework
observation is another program with side effects
The profiler probe wrapped one autocast forward and backward:
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU],
record_shapes=True,
profile_memory=True,
) as trace:
...
Its key averages included:
operator calls
aten::linear 4
aten::addmm 2
aten::relu 1
aten::mm 3
The two actual linear layers explain two addmm calls in forward. linear
appeared four times because the profiler can record both high-level and
redispatched activity under autocast; an event count is not a count of unique
mathematical layers. The three mm calls came from backward paths. A trace
must be interpreted as nested runtime events, not flattened into a list of
source statements.
self_cpu_time_total subtracts time attributed to nested child events from a
CPU event. It is not GPU duration and it can include framework overhead around
native calls. CUDA activity records device kernels and copies on timelines
correlated with host launches. An asynchronous launch’s CPU duration can be
microseconds while the device work lasts far longer.
Recording shapes helps attribute specialization and cost but can retain tensor references or add overhead. Stack capture, module hierarchy, FLOP estimates, memory history, and Python tracing each cost more. Profiling every iteration can change the workload enough to distort it. The scheduled profiler cycles through wait, warmup, and active windows so a long job records bounded regions.
record_function("name") adds user ranges. Modules, optimizers, DDP, FSDP, and
compiler-generated code use ranges to expose framework stages. A range says
which scope was active, not that all nested asynchronous work completed before
the host exited the context.
Memory profiling has at least four views:
- operator allocation events explain who requested or released tracked tensors;
- allocator summaries show active and reserved blocks;
- memory snapshots show segments, blocks, and allocation history;
- system tools show process or device reservation from outside PyTorch.
These accounting systems answer different questions. A tensor can be dead while its block remains reserved. External libraries can allocate memory the PyTorch allocator does not see. Unified memory and mapped files complicate resident bytes further.
torch.utils.benchmark adds warmup, repeated timing, thread control, and
comparison helpers for microbenchmarks. Device benchmarks need explicit
synchronization around timed regions when the timer otherwise measures only
enqueue. Full training throughput needs representative data, optimizer,
communication, and steady-state compilation, not a microbenchmark of one
operator.
PyTorch logging separates causes:
- dispatcher dumps show registered kernels and fallbacks;
- Dynamo logs show graph breaks, guards, and recompiles;
- Inductor traces show lowering, scheduling, generated code, and autotuning;
- distributed debug modes show collective and reducer state;
- C++ stack traces expose native failure sites;
- anomaly detection connects backward errors to forward creation;
gradcheckandopchecktest local mathematical and registration contracts;- the profiler shows where observed execution spent time.
Using the wrong tool creates confident nonsense. A profiler cannot prove an
operator schema is correct. A dispatch table cannot show which GPU kernel
dominated a real run. gradcheck cannot detect a DataLoader duplicate. An
export graph cannot show a later target’s memory pressure.
Instrumentation can also change ordering. Hooks retain tensors. Logging serializes messages. anomaly detection records extra metadata. profilers add callbacks and buffers. deterministic mode selects different kernels. Debug synchronization hides races. A clean investigation uses a diagnostic mode to locate the mechanism, then a lower-perturbation control to confirm it.
the namespaces are projections, not architectural layers
PyTorch’s public package map can look like a stack: torch at the bottom,
torch.nn above it, then optimizers, compilers, and distributed wrappers. In
practice the namespaces cut across the ownership systems already described.
torch exposes tensor construction, operators, dtypes, devices, layouts,
random state, serialization entry points, and backend namespaces. A tensor
method such as x.add(y) and a function such as torch.add(x, y) normally
reach the same operator family. Python syntax x + y uses another binding to
that family. They are API surfaces over dispatcher schemas, not three
independent implementations.
torch.nn.functional contains stateless-looking neural-network operations.
torch.nn.Module classes package those operations with registered parameters,
buffers, modes, initialization, and call hooks. nn.Linear owns weight and
bias; functional.linear consumes weight and bias explicitly. The functional
name does not guarantee mathematical purity: an operation may use random state,
mutate a supplied running statistic under its contract, or return an alias.
torch.autograd exposes differentiation control, custom Functions, graph
utilities, saved-tensor hooks, anomaly detection, and vector-Jacobian
interfaces. The actual graph nodes and engine are native. Autograd formulas
are also registered with operators. The namespace is a control and extension
surface over a cross-cutting dispatch feature.
torch.func exposes composable transforms. torch.fx exposes a Python graph
IR, Graph, Node, GraphModule, interpreters, and transformations. These
intersect but are not synonyms. Traditional torch.fx.symbolic_trace executes
Python with Proxy values and has its own control-flow limitations. Dynamo
also produces FX graphs but captures through CPython frame evaluation and
guards. AOTAutograd uses FX graphs for forward and backward. Export packages
an FX graph with a stronger signature and constraint contract. Printing
GraphModule.code does not identify which capture system produced it.
An FX node has an operation category such as placeholder, get_attr,
call_function, call_method, call_module, or output. Its target may be a
Python callable, method name, module path, or ATen overload. Passes can replace
targets, insert nodes, eliminate dead work, or partition the graph. The graph
does not intrinsically encode tensor shapes; passes attach metadata from fake
or real propagation. An arbitrary FX rewrite is responsible for preserving
aliasing, exceptions, mutations, dtypes, and output structure.
torch.optim owns optimizer algorithms and their state. Learning-rate
schedulers live nearby but are separate state machines. Distributed optimizer
and checkpoint integrations may live under torch.distributed. Fused
implementations may reach ATen operators. The namespace boundary does not
indicate execution location.
torch.amp owns autocast and scaling controls shared across supported device
types. Backend namespaces such as torch.cuda, torch.mps, and torch.xpu
own device-specific stream, event, memory, random, and diagnostic APIs.
torch.backends exposes library and policy switches. A switch in
torch.backends.cuda can affect an aten matrix multiplication reached from
nn.Linear.
torch.distributed contains communication, parallel wrappers, layout-aware
tensors, launch and elastic coordination, checkpointing, and diagnostics.
torch.multiprocessing extends Python multiprocessing with tensor-storage
sharing reducers. torch.utils.data uses multiprocessing but has its own
worker protocol. Sharing a primitive does not collapse the protocols.
torch.profiler, torch.utils.benchmark, and torch._logging observe
execution at different scales. torch.library defines operator extension.
torch.utils.cpp_extension builds native code. torch.export defines an
ahead-of-time program contract. torch.compiler exposes supported controls
over a compiler whose implementation still includes private _dynamo,
_functorch, and _inductor packages.
Private underscores matter. PyTorch exposes some public compiler controls while its implementation and debugging practice still refer to underscored modules. Importing an internal helper can be appropriate for source research or version-pinned infrastructure. It should not be advertised as a stable application contract merely because it is importable.
Several specialized tensor families reuse the same core:
- Sparse tensors carry coordinate or compressed layouts. Only a subset of operations and gradients support each sparse layout and dtype.
- Quantized tensors carry integer payload plus scale and zero-point or more elaborate quantization metadata. Current quantization workflows increasingly live in adjacent projects such as torchao, while PyTorch still owns tensor, dispatcher, export, and compiler integration points.
- Nested tensors represent ragged batches without padding every example to one rectangular size. Operator coverage and layout-specific kernels determine whether the representation avoids conversion.
- Named tensors, conjugate views, negative views, functional tensors, fake tensors, and distributed tensors attach additional semantics to the tensor interface.
The phrase “PyTorch supports operation X” therefore needs coordinates: release, device, dtype, layout, gradient mode, transform, compiler/export context, and shape. Eager dense float32 CPU success is one point in that product space.
the boundary to another array library is an ownership transfer
PyTorch can construct tensors from Python sequences, buffers, NumPy arrays, DLPack exporters, files, and device protocols. Each entrance has different copy and lifetime behavior.
torch.tensor(data) generally copies input data into newly owned tensor
storage. torch.as_tensor attempts to avoid a copy when dtype, device, and
source representation permit. torch.from_numpy creates a CPU tensor sharing
the NumPy array’s memory for supported dtypes and layouts. Writes can be
visible through both handles. A read-only NumPy source does not become safely
writable because PyTorch can manufacture a tensor view.
torch.frombuffer interprets a buffer-protocol exporter with a dtype, count,
and offset. It relies on the exporter’s lifetime and alignment without parsing
a file format or byte order for the caller. A four-byte sequence can
represent a float, integer, or two half values depending on dtype.
DLPack exchanges a tensor-like capsule containing device, dtype, shape,
strides, data pointer, and a deleter. The consumer assumes ownership according
to the protocol. Capsules are ordinarily single-consumption objects because
two consumers independently invoking one deleter would be unsafe. Modern
Python __dlpack__ also communicates stream requirements so a producer does
not hand out device memory before its writes are visible to the consumer.
Zero-copy means no payload copy at the named boundary. Metadata allocation, synchronization, later format conversion, and lifetime coupling can all remain. A noncontiguous NumPy view can become a noncontiguous tensor that a later kernel copies. A DLPack consumer can accept storage and then choose a different layout for an operation.
Tensor.numpy() on a compatible CPU tensor can return a shared NumPy view.
Gradient-requiring tensors, conjugate or negative metadata, unsupported
dtypes, and non-CPU devices require detaching, resolving, or copying according
to the API. .cpu().numpy() on a CUDA tensor includes a device-to-host copy
and completion boundary even if the final NumPy view shares with the new CPU
tensor.
Multiprocessing sharing is another interop path. share_memory_() moves CPU
storage into a shared-memory allocation where processes can map the same
bytes. Sending a tensor through torch.multiprocessing reducers can transmit
a handle to shared storage rather than serializing all values. Gradient and
version-counter semantics across arbitrary processes are not automatically
made coherent. The application still owns synchronization and object
lifetime.
The native-boundary chapter covers buffer and DLPack ownership in depth. The PyTorch-specific lesson is that interop does not bypass TensorImpl, storage, dispatch, or autograd. It creates a tensor whose storage release and visibility contract came from elsewhere.
one training step crosses all of these owners
The original program is small enough to replay chronologically.
| Source event | Immediate owner | State created or consumed |
|---|---|---|
TinyBlock() | Python and nn.Module | child registry, parameters, buffer, mode flags |
AdamW(model.parameters()) | torch.optim | parameter groups; lazy state still empty |
DataLoader yields x, target | sampler, workers, collator, queues | batch tensors and worker progress |
model(x) | module call protocol | hook scopes and child traversal |
self.up(x) | dispatcher and ATen | selected linear decomposition and backend work |
| tensor outputs appear | TensorImpl and StorageImpl | sizes, strides, storage, keys, version, history |
norm, gelu, dropout | operators plus execution modes | statistics, saved tensors, random mask |
cross_entropy | ATen and autograd | scalar loss and backward graph root |
loss.backward() | autograd engine | GraphTask, ready queues, gradient contributions |
| DDP hook, if wrapped | reducer and ProcessGroup | bucket readiness, collective Work, averaged grads |
optimizer.step() | AdamW | moments, step counters, parameter mutation |
| scheduler or scaler update | separate controller | learning rate or scale history |
zero_grad(set_to_none=True) | optimizer helper | leaf gradient references released |
| checkpoint | serialization and application protocol | named tensor state, aliases, metadata, durable publication |
| profiler and logs | observation subsystems | events, stacks, shapes, counters, traces |
This table is not a call stack. Several owners remain active at once. The autograd graph contains tensor and node references after forward. The CUDA stream can contain unfinished work after a Python call returns. The DataLoader can prefetch the next batch while the model runs. A DDP reducer can communicate one ready bucket while autograd computes another. The profiler observes nested events across them.
Object lifetime cuts through the sequence:
- Python references keep Tensor handles alive.
- Tensor handles retain TensorImpl.
- TensorImpl retains storage and optional autograd metadata.
- graph nodes retain saved tensors or packed representations.
- the optimizer retains parameters and state tensors.
- the module retains registered parameters, buffers, and children.
- queues or worker processes retain prefetched batches.
- compiled callables retain guards, graphs, constants, and cache entries.
- asynchronous device work can delay safe storage reuse after host references disappear.
- checkpoints retain a serialized storage graph after process state is gone.
Memory leaks are therefore ownership questions. Appending loss tensors to a
Python list retains their grad_fn graphs. A forward hook that stores outputs
retains activations. retain_graph=True preserves saved state. An optimizer
keeps moment tensors after a parameter stops receiving gradients unless its
state is removed. A persistent DataLoader worker can retain dataset caches.
An Inductor cache can grow on disk while all runtime tensors are freed.
Conversely, freeing too early appears as missing ownership. An extension can launch work on a stream and release scratch storage before recording its use. A DLPack producer can violate the consumer’s stream contract. A custom autograd Function can save a raw pointer instead of a tensor ownership edge. A process can unlink shared memory while peers still need it.
The source line tells us which API was called. Correctness comes from the ownership edges that outlive it.
failures sort by violated contract
PyTorch errors become less mysterious when grouped by the layer that had enough information to reject the program.
tensor and alias failures
Shape mismatch, illegal stride, overlapping in-place write, dtype promotion,
device mismatch, and modified-saved-tensor errors arise from tensor metadata
and alias contracts. Print shape, stride, storage offset, dtype, device,
layout, requires_grad, leaf status, grad_fn, and version-relevant mutation.
Do not begin with compiler flags.
dispatch failures
“No kernel for dispatch key” means the operator schema exists but the active key combination lacks an implementation or fallback. Dump the dispatch table, identify tensor keys and active modes, and reduce the call to one operator. An eager-only custom op that fails under fake tensors needs a fake rule, not a larger GPU.
autograd failures
Missing gradients can come from detached paths, disabled grad mode, a
non-differentiable operator, unused parameters, overwritten .grad, a wrong
custom backward, or optimizer timing. Inspect graph edges and compare against
finite differences or another derivative mode on a small float64 case.
requires_grad=True on a leaf does not guarantee the loss depends on it.
compiler and export failures
Graph breaks, guard failures, recompilation, unsupported symbolic shapes, and target lowering errors belong to different stages. Use fullgraph capture to turn silent breaks into an error, inspect guards, test new shapes, and examine the exported range contract. A successful Dynamo graph does not prove Inductor or a deployment target supports every operator.
device and asynchronous failures
A CUDA error may be reported at a later synchronization than the kernel that caused it. Debug synchronization can move the report closer at the price of changing timing. Check stream use, tensor device, allocator lifetime, and custom extension launch errors. A CPU stack frame that observes the error is not necessarily its origin.
data failures
Duplicate or missing examples, worker death, shared-memory exhaustion,
deadlocks, and slow batches belong to sampler, dataset-replica, queue,
collation, storage, or shutdown contracts. Run with num_workers=0 to
separate dataset logic from multiprocessing, then add workers and measure.
That control changes execution and may hide fork-specific bugs, so it diagnoses
rather than proves.
distributed failures
Hangs often mean ranks entered different collective sequences, one rank failed earlier, a process group was built inconsistently, or a backend has unfinished work. Log rank, group, sequence, tensor metadata, and timeout state. For gradients, compare local values before reduction and verify the intended weighting. Equal rank outputs can all be equally wrong.
numerical failures
NaN, infinity, overflow, underflow, cancellation, and unacceptable drift can come from the model, dtype policy, optimizer, compiler, or backend library. Start with a small high-precision reference, check finite intermediates, and separate deterministic replay from acceptable error. A tolerance should come from the computation and application, not from making a test green.
state and recovery failures
Unexpected checkpoint keys, stale optimizer moments, wrong device mapping, corrupt files, and unsafe pickle globals belong to serialization and application schema. Print missing and unexpected keys. Validate shapes and metadata. Restore all coupled state or explicitly accept a changed run.
The categories overlap because the framework composes. A false custom-op alias schema can surface as compiled wrong code. A DataLoader duplicate can look like an optimizer convergence problem. A missing DDP collective can leave autograd waiting. The point is not to force one label. It is to find the earliest owner whose contract was violated.
what PyTorch actually is
The first tensor was only a handle. The first operator was only a schema. The module was a registry, backward was a scheduled graph, AdamW was another state machine, and “the GPU call” was an asynchronous dispatch into a device and allocator domain. None of those corrections made the high-level Python less real. They explained how it could remain small.
PyTorch is not best understood as a neural-network library sitting on top of a tensor library. Its shared operator language connects a set of composable owners:
- TensorImpl and StorageImpl own identity, metadata, bytes, and lifetime.
- The dispatcher composes backend and feature meanings for operator schemas.
nn.Moduleowns named state and a call protocol.- autograd owns derivative graph construction and execution.
- optimizers own long-lived update state.
- modes and
torch.funcreinterpret operator execution. - device backends own allocation, queues, streams, and kernels.
- Dynamo, AOTAutograd, and Inductor move eager decisions into guarded compiled programs.
- export owns a constrained program artifact.
- DataLoader owns a prefetched worker protocol.
- C10D, DDP, FSDP, DTensor, and meshes own different distributed contracts.
- serialization owns a storage graph and a loader trust boundary.
- profiler and logging systems own observations with measurable perturbation.
- extension APIs let new code join only by declaring the semantics each interpreter needs.
This model also sets the boundary of the chapter. PyTorch can dispatch a matrix multiplication without defining the mathematics of a transformer. It can call a CUDA kernel without owning the GPU instruction set. It can invoke NCCL without choosing the network fabric. It can export an ATen graph without being the downstream runtime. Those systems matter, but importing them into every PyTorch explanation makes it impossible to see which layer made a decision.
The original reshape copy no longer looks like a quirk. reshape had to
preserve a tensor contract across storage and stride metadata. Later systems
were entitled to trust that contract. Autograd could reason about the copy,
functionalization could expose it, fake execution could propagate its shape,
the compiler could guard its layout, export could constrain its dimensions,
and serialization could preserve its result.
That is the framework’s central bargain. Eager Python remains permissive and dynamic because operations cross explicit semantic boundaries. The cost is that every new backend, transform, compiler, subclass, and custom operator must honor those boundaries. When they do, one small training step can pass through all of PyTorch without the user manually coordinating each system. When they do not, the shortest route to the bug is to ask who owned the missing fact.