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 (i0,,in1)(i_0,\ldots,i_{n-1}) reaches the element at logical offset

o+d=0n1idsd,o + \sum_{d=0}^{n-1} i_d s_d,

where oo is storage_offset() and sds_d 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 PrivateUse1 accelerator.
  • 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 θ\theta, gradient gg, first moment mm, and second moment vv, a simplified Adam update contains:

mt=β1mt1+(1β1)gt,m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, vt=β2vt1+(1β2)gt2.v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2.

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 of f;
  • vjp and jvp expose reverse and forward products;
  • jacrev and jacfwd construct Jacobians through those products;
  • vmap(f) adds a batch interpretation without writing the outer loop;
  • functional_call executes a module with supplied parameter and buffer mappings;
  • functionalize removes 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:

  1. Dynamo intercepts CPython frame evaluation, symbolically executes bytecode, records FX graphs, and emits guards for assumptions.
  2. AOTAutograd functionalizes relevant work, analyzes mutation and aliasing, and can trace forward and backward graphs ahead of execution.
  3. 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:

  1. Worker processes receive serialized dataset, function, and task state according to the platform’s process-start method.
  2. The iterator enqueues up to a prefetch budget of index tasks.
  3. Workers fetch and collate examples.
  4. Tensor storage crosses a multiprocessing boundary, often using shared memory rather than copying every byte through a pipe.
  5. The main iterator buffers or emits results and schedules replacement tasks.
  6. 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:

  1. all-gather the parameter shards required by a module;
  2. execute its local tensor operations;
  3. release or reshard full parameters according to policy;
  4. repeat for later modules;
  5. 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 dimension dim;
  • 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 vmap understands;
  • 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;
  • gradcheck and opcheck test 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 eventImmediate ownerState created or consumed
TinyBlock()Python and nn.Modulechild registry, parameters, buffer, mode flags
AdamW(model.parameters())torch.optimparameter groups; lazy state still empty
DataLoader yields x, targetsampler, workers, collator, queuesbatch tensors and worker progress
model(x)module call protocolhook scopes and child traversal
self.up(x)dispatcher and ATenselected linear decomposition and backend work
tensor outputs appearTensorImpl and StorageImplsizes, strides, storage, keys, version, history
norm, gelu, dropoutoperators plus execution modesstatistics, saved tensors, random mask
cross_entropyATen and autogradscalar loss and backward graph root
loss.backward()autograd engineGraphTask, ready queues, gradient contributions
DDP hook, if wrappedreducer and ProcessGroupbucket readiness, collective Work, averaged grads
optimizer.step()AdamWmoments, step counters, parameter mutation
scheduler or scaler updateseparate controllerlearning rate or scale history
zero_grad(set_to_none=True)optimizer helperleaf gradient references released
checkpointserialization and application protocolnamed tensor state, aliases, metadata, durable publication
profiler and logsobservation subsystemsevents, 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.Module owns named state and a call protocol.
  • autograd owns derivative graph construction and execution.
  • optimizers own long-lived update state.
  • modes and torch.func reinterpret 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.