Deep Dive - PyTorch

You transpose a matrix, and its rows become columns without copying its values. Then the next operation gets slower. Where did the extra work happen? PyTorch lets us inspect both the shared data and the code that consumes it, so we can follow that change instead of guessing from the final output.

PyTorch is a framework for tensor computation and automatic differentiation, which calculates the gradients used in training. The running case is scaled dot-product attention (SDPA), the operation that combines queries, keys, and values in a transformer.^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

A transpose changes strides, the steps that map tensor indices into storage. A mask changes which tokens contribute. Either can change the execution path while the output shape stays the same. The previous MLflow lesson recorded the run. The question here is which operation inside that run changed.

The runnable examples use PyTorch 2.13.0 on CPU, one computation thread, and tiny tensors. They check real views, autograd, attention, compiler capture, and checkpoint reload. Run separate accelerator benchmarks and multi-rank FSDP2 tests. You should already be comfortable with a training loop and matrix multiplication; the distributed sections build on the earlier FSDP and Megatron lessons.

Eager Python, compiled kernels

Deep-learning research needs two properties that pull in opposite directions:

1. Local reasoning. A line of Python should execute now, produce a normal traceback, and allow a debugger or print statement.
2. Global optimization. The same line should reach vectorized CPU code, CUDA kernels, fused attention, graph compilers, and distributed collectives.

PyTorch's imperative tensor API, tape-based reverse-mode autodiff, and extension points made that combination practical. The 2019 paper describes a Python-first library with GPU acceleration and dynamic neural networks. The current repository still exposes that Python surface, then implements it with C++, CUDA, generated operator bindings, compiler components, and distributed packages.^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

The same tensor API can select different accelerator backends. On Apple silicon, the Metal Performance Shaders (MPS) backend executes supported operations on the local GPU, often within a unified-memory architecture.^{[4]Reference 4MPS backend.https://docs.pytorch.org/docs/stable/notes/mps[5]Reference 5Accelerated PyTorch training on Mac.https://developer.apple.com/metal/pytorch/}

A discrete NVIDIA device instead commonly requires an explicit host-to-device transfer before its CUDA kernels can consume a CPU-created tensor. Operator coverage, numerical behavior, and memory capacity vary by backend; test the exact operations and device you plan to use.^{[6]Reference 6CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Before reading the stack as a component list, predict where one Python call can change path: storage layout, dispatch mode, gradient recording, compiler capture, kernel eligibility, or communication. torch.Tensor is only the front of that stack:

There's no single "PyTorch performance number." Eager execution, compiled execution, fused kernels, memory layout, device, sequence length, and communication topology all change the path. "It's PyTorch" isn't a diagnosis. Name the path before changing a flag.

The tensor is metadata plus storage

The transpose is the first place that extra work can hide, before any kernel runs. For an ordinary dense, strided tensor, torch.Tensor is a Python handle to an internal object whose TensorImpl carries metadata and refers to storage. Several tensors can view the same allocation. This model doesn't describe every tensor kind: sparse tensors use different representations, and meta tensors have no data allocation.

• Storage owns an allocation and byte-sized region on a device. TensorImpl carries dtype metadata that tells kernels how to interpret those bytes.
• Sizes describe the length of each logical dimension.
• Strides describe how an index moves through storage.
• Offset points the logical zero element into storage.
• Device and dtype select where values live and how kernels interpret bytes.
• Autograd metadata links the tensor to history, leaves, and version tracking.

Before touching the shape, predict what a transpose should change: coordinates and strides, not the six storage cells. Take six values in storage, [0, 1, 2, 3, 4, 5]. A contiguous tensor with shape [2, 3] and element strides [3, 1] places element (i, j) at offset + 3i + j. A transpose exposes shape [3, 2] with strides [1, 3] without moving bytes. Value 3 is x[1, 0] and y[0, 1] at once.

View larger

Check the actual PyTorch storage and strides. Strides and storage_offset() count elements, not bytes: element 3 of an int64 allocation starts 24 bytes from its beginning. Mutation demonstrates aliasing; reshape demonstrates a case where a new allocation is necessary.

import torch

torch.set_num_threads(1)
x = torch.arange(6, dtype=torch.int64).reshape(2, 3)
y = x.t()
assert x.stride() == (3, 1) and y.stride() == (1, 3)
assert x.untyped_storage().data_ptr() == y.untyped_storage().data_ptr()
assert x[1, 0].item() == y[0, 1].item() == 3
assert x.untyped_storage().nbytes() == 48
print("x:", x.tolist(), "strides:", x.stride())
print("y:", y.tolist(), "strides:", y.stride())
print("storage bytes:", x.untyped_storage().nbytes())
try:
    y.view(-1)
except RuntimeError:
    print("view(-1): incompatible strides")
else:
    raise AssertionError("This transpose cannot be flattened with view")
flat = y.reshape(-1)
assert flat.tolist() == [0, 3, 1, 4, 2, 5]
assert flat.untyped_storage().data_ptr() != y.untyped_storage().data_ptr()
y[0, 1] = 30
assert x[1, 0].item() == 30 and flat[1].item() == 3
print("alias changed x; copied reshape stayed unchanged")

x: [[0, 1, 2], [3, 4, 5]] strides: (3, 1)
y: [[0, 3], [1, 4], [2, 5]] strides: (1, 3)
storage bytes: 48
view(-1): incompatible strides
alias changed x; copied reshape stayed unchanged

A training step that mutates one view can affect the other because both names reach the same storage. A compiler guard may also specialize on the stride pattern and recompile when a later batch arrives with a different layout.

Follow that layout into attention. Queries (Q), keys (K), and values (V) often begin as [batch, sequence, heads, head_dim], then become [batch, heads, sequence, head_dim] with a permutation.

The shape changes instantly, but the next matrix multiplication may still pay for a contiguous conversion if its kernel can't consume the resulting strides.

view requires a compatible stride pattern. reshape can return a view or allocate a copy. When a kernel slows down, inspect shape, stride(), storage_offset(), is_contiguous(), and memory format at that boundary. Don't use private attributes such as _base as application logic.^{[7]Reference 7PyTorch 2.14: Tensor Viewshttps://docs.pytorch.org/docs/2.14/tensor_view.html}

Operator schemas and the dispatcher

After Python creates an operator call, PyTorch needs a stable contract for its inputs, outputs, aliases, and mutation. Operators are registered with schemas that describe names, arguments, returns, overloads, defaults, and aliasing behavior. Many generated Python and C++ bindings lead back to an ATen operator, so torch.mm(x, w) is a typed operator invocation rather than an opaque Python function.

The dispatcher uses a DispatchKeySet assembled from tensor backends and active modes. Keys cover CUDA, CPU, autograd, functionalization, and batching, among other behaviors. It chooses the highest-priority applicable implementation, then may redispatch to a lower key. Python fake-tensor and tracing modes intercept through Python dispatch; the pinned revision also has a Fake key for C++ fake-tensor mode. Those transformation keys don't identify a physical device. This layered selection lets one operator expose one public contract while several systems observe or implement it.^{[8]Reference 8PyTorch Source Snapshot: 4ee8feahttps://github.com/pytorch/pytorch/tree/4ee8fea1f8cbb27acf0dac1c4f2e8b8bb6aef7f7}

For the SDPA call in the diagram, predict the first labels a profiler might expose: schema, CUDA dispatch, and Autograd mode. The chosen branch then depends on backend eligibility:

View larger

Generated code reduces handwritten binding work, but it doesn't remove semantic obligations. A new backend must implement the operators it advertises, preserve dtype and device behavior, and respect aliasing rules.

A custom operator that returns a wrong stride can pass a toy example and fail later when a view or compiler guard depends on it.

The dispatch stack also explains why a Python context can change results without changing model code. torch.no_grad() removes autograd recording. Inference mode can remove additional metadata work under stricter rules.

Functionalization rewrites in-place mutations into out-of-place operations for transformations. Fake tensors carry shape, dtype, and device information without allocating real storage. Each mode participates through dispatch keys or interpreter layers rather than through a second tensor API.

Autograd records a tape, not a static model

Before loss.backward(), predict what graph exists: only operations actually executed in this iteration, plus saved values needed for their reverse functions.

In a grad-enabled region, when a tensor with requires_grad=True participates in a differentiable operation, PyTorch creates a dynamic graph of backward nodes. Each node saves the values or metadata needed for its derivative. Calling loss.backward() walks that graph in reverse, accumulates gradients into leaves, and releases saved intermediates when they're no longer needed unless the graph is retained.

That graph follows the actual branch taken by Python. A loop can stop early, a condition can select a different module, or an attention mask can change shape. Ordinary eager mode has no separate graph-building phase.

Views and in-place mutation

Autograd must detect when a saved value changes before backward uses it. Ordinary autograd-tracked tensors use version counters that increment on tracked in-place updates. Tensors created in inference mode are an exception: they don't track a version counter.

When a backward node saves a tensor, it saves that version too. Accessing the saved tensor later raises if the version has moved.^{[9]Reference 9Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html}

Views share storage and often share version tracking with their base. Mutating a view can therefore invalidate a backward node that saved the base, producing errors such as "one of the variables needed for gradient computation has been modified in-place."

Some writes fail sooner: mutating a leaf tensor that requires gradients while grad mode is enabled can raise at the update itself. The saved-version check explains delayed failures; it isn't the only protection against unsafe writes.

The error is useful: silently accepting the mutation could produce an incorrect gradient. Clone before mutation when the value is needed for backward, or rewrite the operation out-of-place. Don't "fix" it by disabling gradient tracking around a required training mutation.

detach() creates a tensor that shares storage but stops gradient history. Use it when a value mustn't contribute to a later gradient. In-place writes through detached aliases can still mutate the original storage, though.

clone().detach() separates both history and bytes when isolation is required.

For a scalar loss $L=x^2$ at $x=3$, the derivative is $2x=6$. Backward needs the original value of $x$. This CPU check first obtains that gradient, then shows why writing through a detached alias is still unsafe. no_grad or detach stops graph recording, not shared-storage mutation.

import torch

torch.set_num_threads(1)
x = torch.tensor([3.0], dtype=torch.float64, requires_grad=True)
x.square().sum().backward()
assert x.grad.item() == 6.0
print("gradient at x=3:", x.grad.item())

x = torch.tensor([3.0], dtype=torch.float64, requires_grad=True)
loss = x.square().sum()
x.detach().add_(1.0)  # Shares storage with the saved x.
try:
    loss.backward()
except RuntimeError as error:
    assert "modified by an inplace operation" in str(error)
    print("detached alias mutation: rejected during backward")
else:
    raise AssertionError("Expected a saved-value version mismatch")

x = torch.tensor([3.0], dtype=torch.float64, requires_grad=True)
loss = x.square().sum()
independent = x.detach().clone()
independent.add_(1.0)
loss.backward()
assert x.item() == 3.0 and independent.item() == 4.0
assert x.grad.item() == 6.0
assert torch.autograd.gradcheck(lambda t: t.square().sum(), (x,))
print("independent clone: original gradient preserved; gradcheck passed")

gradient at x=3: 6.0
detached alias mutation: rejected during backward
independent clone: original gradient preserved; gradcheck passed

torch.no_grad() tells autograd not to record operations in a dynamic region. Inference mode goes further: tensors created there aren't meant to re-enter a grad-tracked computation after you leave the context.^{[9]Reference 9Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html}

torch.autograd.gradcheck compares analytical and numerical derivatives for small double-precision inputs. Test distributed and low-precision kernels under their actual execution modes as well.

Hooks, custom autograd.Function, saved tensor hooks, and checkpointing extend this machinery. Each extension adds a contract: a custom backward must match the forward's aliasing and dtype semantics, and activation checkpointing must replay a compatible forward when backward asks for saved values.

Attention is a kernel-selection problem

Attention makes those execution choices visible. Before naming a kernel, count the score cells. Given Q, K, and V with shape [B, H, S, D], scaled dot-product attention computes:

Here $M$ is an additive mask: zero for allowed positions and negative infinity for blocked ones. Softmax runs along the key dimension. This expression omits dropout and assumes equal query/key lengths and head counts; the API also supports rectangular attention and a different value-head dimension.

For B=1, H=2, S=4, and D=8, the score tensor has 1 × 2 × 4 × 4 = 32 elements before dtype accounting. At S=8192, it has 1 × H × 8192 × 8192 scores. That's why materializing the full score matrix becomes a memory problem. Fused kernels tile the computation so intermediate scores stay in on-chip memory or are recomputed instead of written as a giant matrix.

batch, heads, sequence = 1, 2, 8192
score_cells = batch * heads * sequence * sequence
print(score_cells)
assert score_cells == 134_217_728

134217728

torch.nn.functional.scaled_dot_product_attention is documented as beta. Its API overview lists FlashAttention-2, memory-efficient attention, and the math implementation, while the backend enum also exposes cuDNN attention. Don't treat the overview's three-item list as an exhaustive inventory of implementations in an installed build.^{[10]Reference 10PyTorch 2.14: SDPBackendhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.attention.SDPBackend.html[11]Reference 11torch.backends.cuda: SDPA backend eligibilityhttps://docs.pytorch.org/docs/stable/backends.html}

Backend selection depends on device, dtype, shape, mask, causal mode, build support, and enabled backends. Its heuristics can select a slower implementation for a particular workload. Prefer torch.nn.attention.sdpa_kernel when you need to restrict the candidate backends during diagnosis.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html[11]Reference 11torch.backends.cuda: SDPA backend eligibilityhttps://docs.pytorch.org/docs/stable/backends.html[12]Reference 12FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

Same function, different inputs, different legal paths. Keep these contracts beside the call:

• Dropout isn't implicit in eval(). The function always applies dropout_p. Pass 0.0 when the module isn't training.
• Boolean masks mean "keep." True participates in attention. That's the inverse of MultiheadAttention's key_padding_mask, where True means masked out.
• Grouped-query attention is experimental in this API. It lets several query heads share each key/value head. The 2.14 documentation lists FlashAttention, cuDNN, and math support on CUDA, plus memory-efficient attention on NVIDIA CUDA. Query-head count must be divisible by key/value-head count, key and value head counts must match, and nested tensors aren't supported. An older wheel may have narrower support; check its backend eligibility instead of transferring this list across versions.
• Fused kernels change numerics. The math backend keeps intermediates in float32 when inputs are half or bfloat16, and it supports float64 when you need a higher-precision check.

For non-square attention, is_causal=True uses an upper-left causal alignment. A one-token query against a longer cached key prefix therefore doesn't automatically mean “this is the final token; allow the whole prefix.” Construct the mask or causal bias for the actual query/key positions. Also ensure every query has permitted keys rather than assuming all-masked rows behave identically across implementations.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

At the call site, inspect these signals:

A fused attention kernel isn't automatically faster for every shape. Very short sequences can be launch-bound, an unsupported mask can force a fallback, and a layout conversion can cost more than the fused kernel saves.

Record the selected backend and shape when comparing runs. Strides alone reveal layout, not the history of allocations.

Start with a numerical check before measuring a fused kernel. With two queries, two keys, and a mask allowing only the first key, both output rows must equal the first value row. The example forces CPU math, compares the formula and gradients, then checks that mask. Finally, it keeps only the first query: upper-left causal alignment permits only key 0, while a query representing the final cached position can use both keys. It doesn't exercise CUDA fused attention.

import math
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

torch.set_num_threads(1)
q = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
k = q.detach().clone().requires_grad_()
v = torch.tensor([[[[10.0, 0.0], [0.0, 20.0]]]], dtype=torch.float64, requires_grad=True)
with sdpa_kernel(SDPBackend.MATH):
    actual = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
reference = (q @ k.transpose(-2, -1) / math.sqrt(2)).softmax(dim=-1) @ v
torch.testing.assert_close(actual, reference)
actual_grads = torch.autograd.grad(actual.square().sum(), (q, k, v))
reference_grads = torch.autograd.grad(reference.square().sum(), (q, k, v))
for actual_grad, expected_grad in zip(actual_grads, reference_grads):
    torch.testing.assert_close(actual_grad, expected_grad)

keep_first = torch.tensor([[True, False], [True, False]])
with sdpa_kernel(SDPBackend.MATH):
    masked = F.scaled_dot_product_attention(q, k, v, attn_mask=keep_first, dropout_p=0.0)
expected = torch.tensor([[[[10.0, 0.0], [10.0, 0.0]]]], dtype=torch.float64)
torch.testing.assert_close(masked, expected)
print("CPU math: forward and all Q/K/V gradients match the formula")
print("True means keep:", masked.detach().tolist())

with sdpa_kernel(SDPBackend.MATH):
    upper_left = F.scaled_dot_product_attention(q[:, :, :1], k, v, is_causal=True)
    cached_last = F.scaled_dot_product_attention(q[:, :, :1], k, v, is_causal=False)
torch.testing.assert_close(upper_left, expected[:, :, :1])
torch.testing.assert_close(cached_last, reference[:, :, :1])
assert not torch.allclose(upper_left, cached_last)
print("one query, two keys: upper-left causal sees key 0; cached-last sees both")

CPU math: forward and all Q/K/V gradients match the formula
True means keep: [[[[10.0, 0.0], [10.0, 0.0]]]]
one query, two keys: upper-left causal sees key 0; cached-last sees both

FlexAttention is a separate API for custom attention: a score modifier changes attention scores, while a block mask describes sparsity that the kernel can exploit. It isn't another backend selected by scaled_dot_product_attention.^{[13]Reference 13PyTorch 2.14: torch.nn.attention.flex_attentionhttps://docs.pytorch.org/docs/2.14/nn.attention.flex_attention.html}

torch.compile: capture, transform, lower

torch.compile is a compiler entry point, not one kernel. In the common PyTorch 2 path, TorchDynamo observes Python bytecode and extracts graph regions.

Ahead-of-time autograd (AOTAutograd) can stage forward and backward graphs together. TorchInductor lowers those graphs to device-specific code, often including Triton kernels for GPUs and generated C++ or vectorized code for CPUs.^{[14]Reference 14PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilationhttps://doi.org/10.1145/3620665.3640366[15]Reference 15torch.compilehttps://docs.pytorch.org/docs/2.14/generated/torch.compile.html}

Predict what survives compilation: a guarded region can preserve eager semantics, but changing Python values or tensor metadata can send execution back to eager or trigger a new variant. Four contracts make that boundary visible:

1. Capture. Dynamo specializes on Python values, tensor metadata, and control-flow assumptions. It emits guards that must hold for a cached graph.
2. Functionalization and decomposition. In-place or composite operations may be rewritten into a form easier to transform. Operator decompositions can replace a high-level op with lower-level ATen ops.
3. AOTAutograd. Forward and backward graphs are staged so compiler passes can see saved values and gradient dependencies.
4. Lowering. Inductor fuses compatible operations, selects or generates kernels, schedules loops, and caches compiled artifacts.

View larger

Official guidance is to compile the highest-level function that doesn't cause excessive trouble, typically the train or eval step without the data loop, or a top-level nn.Module.

The general compiler guide recommends compiling the inner module when a distributed wrapper causes capture trouble, and using model.compile() for a top-level module. Its Distributed Data Parallel (DDP) example compiles before wrapping. Don't turn that example into a universal ordering rule for composable Fully Sharded Data Parallel (FSDP2): validate the block boundaries, hooks, and collectives for the specific sharding/compilation integration.^{[16]Reference 16Where to apply torch.compile?https://docs.pytorch.org/docs/main/user_guide/torch_compiler/compile/programming_model.where_to_apply_compile.html}

Graph breaks and recompilation

Python operations that Dynamo can't represent cause graph breaks. Common triggers include data-dependent control flow, unsupported library calls, mutation patterns, and side effects such as printing a tensor's value.

A graph break returns to Python, then compilation resumes around a later region. The program can remain correct while losing fusion and adding synchronization.

Guards are another failure boundary. A graph specialized for one dtype, stride, device, or static shape can recompile when a later batch changes.

Shape polymorphism reduces needless recompiles but can produce more general kernels. A compile cache can fill with variants when a serving workload has highly varied sequence lengths.

Use the compiler's explanation and recompilation reports instead of guessing. For the example below, run TORCH_LOGS="graph_breaks,guards,recompiles" python compile_boundaries.py to see captured assumptions and which changed value invalidates them. The official troubleshooting guidance recommends reducing the program to a small reproducer and inspecting graph breaks and guards.^{[17]Reference 17Troubleshooting torch.compilehttps://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/torch.compiler_troubleshooting.html}

Then decide whether to rewrite code, allow a fallback, or compile only stable regions. Use fullgraph=True to expose graph breaks during diagnosis; production settings depend on workload behavior and compiler support, since partial capture may run valid programs that full capture rejects.

The distinction is observable without benchmarking kernels. A custom backend below counts captured graphs and executes each graph eagerly. Repeating the same Python scale reuses its graph; changing that guarded scalar creates another specialization. A separate aot_eager check stages autograd and compares gradients, but doesn't run Inductor code generation.

import torch
import torch._functorch.config

torch.set_num_threads(1)
torch.compiler.reset()
captures = []

def counting_backend(graph, inputs):
    captures.append(graph)
    return graph.forward

def scaled_square(x, scale):
    return x.square() * scale

compiled = torch.compile(scaled_square, backend=counting_backend, fullgraph=True, dynamic=False)
x = torch.tensor([1.0, 2.0], dtype=torch.float64)
torch.testing.assert_close(compiled(x, 2.0), torch.tensor([2.0, 8.0], dtype=torch.float64))
compiled(x, 2.0)
assert len(captures) == 1
torch.testing.assert_close(compiled(x, 3.0), torch.tensor([3.0, 12.0], dtype=torch.float64))
assert len(captures) == 2

def loss_fn(value):
    return value.sin().square().sum()

a = x.clone().requires_grad_()
b = x.clone().requires_grad_()
staged = torch.compile(loss_fn, backend="aot_eager", fullgraph=True)
eager_loss, staged_loss = loss_fn(a), staged(b)
eager_loss.backward()
staged_loss.backward()
torch.testing.assert_close(eager_loss, staged_loss)
torch.testing.assert_close(a.grad, b.grad)

# Match the compiler's backward assumption to backward outside autocast.
torch.manual_seed(0)
inputs = torch.randn(4, 3)
weight = torch.randn(3, 2, requires_grad=True)
staged_weight = weight.detach().clone().requires_grad_()
def amp_loss(value):
    return (inputs @ value).square().mean()
with torch.autocast("cpu", dtype=torch.bfloat16):
    eager_amp_loss = amp_loss(weight)
eager_amp_loss.backward()
with torch._functorch.config.patch(backward_pass_autocast="off"):
    compiled_amp_loss = torch.compile(amp_loss, backend="aot_eager", fullgraph=True)
    with torch.autocast("cpu", dtype=torch.bfloat16):
        staged_amp_loss = compiled_amp_loss(staged_weight)
    staged_amp_loss.backward()
torch.testing.assert_close(eager_amp_loss, staged_amp_loss, rtol=0, atol=0)
torch.testing.assert_close(weight.grad, staged_weight.grad, rtol=0, atol=0)
print("guarded Python scales: 2 captured graphs for 3 calls")
print("aot_eager: forward and gradient agreement; no optimized-kernel benchmark")
print("CPU BF16 autocast: backward outside autocast matches with assumption off")

guarded Python scales: 2 captured graphs for 3 calls
aot_eager: forward and gradient agreement; no optimized-kernel benchmark
CPU BF16 autocast: backward outside autocast matches with assumption off

Compilation can hurt when startup or generated code costs more than it saves. A tiny model may spend more time compiling than executing, and dynamic shapes can trigger repeated compilation.

Compiled backward also captures assumptions that eager backward can read dynamically. The compiler's autograd guide specifically documents the ambient autocast assumption: a mismatch between that assumption and where backward actually runs can cause silent incorrectness. Its default, "same_as_forward", assumes backward runs under the forward's autocast context. If forward and loss use autocast but backward runs outside it, the guide specifies backward_pass_autocast="off". The example scopes that assumption with torch._functorch.config.patch and keeps both compilation and backward inside the patch.^{[18]Reference 18torch.compile Has Different Autograd Semanticshttps://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_backward.html}

This is an internal configuration knob documented for this compiler contract, so check the installed version's guidance before using it. The added CPU BF16 check exercises a matrix operation and verifies loss and gradients with aot_eager; it doesn't validate CUDA AMP, GradScaler, Inductor-generated kernels, or a distributed training recipe. A precision change needs those checks separately.

A custom CUDA extension can be opaque to the compiler. A fused kernel can increase register pressure or change numerical error. Benchmark warm runs after compilation, include realistic shape distributions, and compare outputs, memory, and average tokens per second together.

Distributed PyTorch: process groups first

torch.distributed starts with a world of processes and one or more process groups. A rank is an endpoint in a group; a collective defines which ranks exchange tensors and how results are reduced or gathered. NCCL commonly carries CUDA collectives, while Gloo or other backends serve CPU or testing paths.^{[19]Reference 19PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

Make the failure concrete: rank 0 enters all_reduce while rank 1 enters all_gather. Each call can be healthy in isolation, yet the group can hang because collective order is part of the contract.

The primitive operations are familiar:

One rank taking a different Python branch is enough to desynchronize a supposedly identical training step. Set process-group timeouts, enable distributed debug logs for diagnosis, and keep a minimal rank-symmetric reproducer.

DeviceMesh and DTensor make placement explicit

Nested process groups become difficult when one model uses data, tensor, sequence, and pipeline axes. DeviceMesh names an n-dimensional arrangement of devices and derives submeshes and process groups from it. A mesh lets code say "the tensor-parallel axis" instead of rebuilding rank lists in every layer.^{[19]Reference 19PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

Read each placement as the declared logical layout of a tensor across mesh dimensions. DTensor adds a global tensor view plus a placement on each mesh dimension; measure execution to determine communication and performance costs. The three placements are:^{[20]Reference 20PyTorch 2.14: torch.distributed.tensorhttps://docs.pytorch.org/docs/2.14/distributed.tensor.html}

• Replicate: ranks along this mesh dimension hold the same value. Another mesh dimension can still shard it.
• Shard(dim): ranks along this mesh dimension store slices of tensor dimension dim, following torch.chunk(dim) semantics.
• Partial: ranks along this mesh dimension hold contributions that need a specified reduction, commonly a sum, before the value is complete.

Operations propagate placements and insert collectives when needed. A sharded matrix multiply can yield a Partial output. Redistributing Shard(dim) to Replicate() is an all-gather; Partial() to Replicate() is an all-reduce; Partial() to Shard(dim) is a reduce-scatter. These rules define shape and placement; the required collectives can still move data.^{[20]Reference 20PyTorch 2.14: torch.distributed.tensorhttps://docs.pytorch.org/docs/2.14/distributed.tensor.html}

For a sum-Partial vector on two ranks, local values [1, 2] and [10, 20] represent the complete vector [11, 22]. All-reduce leaves [11, 22] on both ranks. Reduce-scatter instead leaves [11] on rank 0 and [22] on rank 1. Simply concatenating the Partial inputs would produce four elements and the wrong logical tensor.

For a mesh mesh = (dp=2, tp=4), a weight with global shape [4096, 16384] can be Shard(1) over tp, giving each TP rank [4096, 4096], while DP replicas hold the same shard. If a downstream operation needs a full output, DTensor may insert an all-gather or reduce. Estimate communication before selecting a placement.

Each edge below redistributes one mesh dimension; other dimensions keep their placement.

View larger

FSDP2 uses this placement machinery for data-parallel state sharding. Its default is Shard(0) on a one-dimensional mesh, or hybrid sharding (Replicate(), Shard(0)) on a two-dimensional mesh. The shard_placement_fn option can override the parameter dimension or mesh. Placement dimensions describe storage axes; they don't by themselves identify a parallelism strategy.^{[21]Reference 21torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/2.14/distributed.fsdp.fully_shard.html}

Tensor and pipeline parallelism split different work

Before comparing parallel axes, predict their traffic: TP moves partial layer results, PP moves activations between stages, and FSDP2 moves parameter and gradient shards around module use.

Tensor parallelism (TP) divides one operator across ranks. parallelize_module takes a one-dimensional DeviceMesh slice plus a plan such as ColwiseParallel or RowwiseParallel. A transformer multilayer perceptron (MLP) can split its first projection's output features, keep that activation partitioned, then split the second projection's input features and reduce partial outputs. The rank-local matrix multiplications are smaller, but their layouts must agree on communication.^{[22]Reference 22PyTorch 2.14: Tensor Parallelismhttps://docs.pytorch.org/docs/2.14/distributed.tensor.parallel.html}

Watch the matrix convention: nn.Linear.weight is stored as [out_features, in_features]. Splitting output features with ColwiseParallel shards that stored weight on dimension 0, even though a mathematical right-hand matrix in $XW$ would be described as column-sharded. FSDP2 can also use dimension 0, but it reconstructs weights for data-parallel compute rather than leaving the operator split across tensor-parallel ranks.^{[8]Reference 8PyTorch Source Snapshot: 4ee8feahttps://github.com/pytorch/pytorch/tree/4ee8fea1f8cbb27acf0dac1c4f2e8b8bb6aef7f7}

Pipeline parallelism (PP) divides model execution into stages. Stage 0 can run early layers for microbatch 2 while stage 1 runs later layers for microbatch 1.

PyTorch's torch.distributed.pipelining package provides stage construction and schedules including GPipe, one-forward-one-backward (1F1B), interleaved 1F1B, and looped breadth-first search (BFS). Current documentation still labels the package alpha and notes it was migrated from PiPPy, so pin the PyTorch version and test schedule changes before relying on its API.^{[23]Reference 23PyTorch 2.14: Pipeline Parallelismhttps://docs.pytorch.org/docs/2.14/distributed.pipelining.html}

These axes solve different limits and compose through different communication:

Pipeline throughput depends on microbatch count and schedule. Too few microbatches leave stages idle during fill and drain. Too many reduce bubble fraction but can raise activation memory and scheduling overhead.

For 1F1B, forward and backward work alternate after warmup. With equally sized microbatches and mean-reduced losses, the schedule's scale_grads=True divides accumulated gradients by the microbatch count. For sum-reduced losses, use scale_grads=False to preserve the full-batch sum. Dividing a mean loss yourself and leaving schedule scaling on divides twice. Unequal microbatch sizes need explicit weighting to reproduce a full-batch mean.^{[23]Reference 23PyTorch 2.14: Pipeline Parallelismhttps://docs.pytorch.org/docs/2.14/distributed.pipelining.html}

FSDP2 shards parameters around the forward

With the default FSDP2 policy, managed parameters are DTensor shards at rest. Its composable API, fully_shard, defaults to dimension-0 sharding. Applying it bottom-up, first to transformer blocks and then the root, creates smaller communication groups of parameters than sharding only the root.^{[21]Reference 21torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/2.14/distributed.fsdp.fully_shard.html}

Around module execution, pre-forward and pre-backward hooks all-gather each managed parameter group and register unsharded parameters for compute. For an ordinary tensor parameter that's a plain torch.Tensor; when composing with TP, the compute parameter can remain a DTensor with TP placement. After the module runs, the unsharded copy can be freed and the data-parallel shard registered again. Gradients are reduce-scattered so each rank keeps its optimizer-owned slice.^{[21]Reference 21torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/2.14/distributed.fsdp.fully_shard.html}

This hook boundary changes a common debugging call: use model(input), not model.forward(input). The hooks that all-gather hang off the module __call__ path.

If you must invoke forward directly, call model.unshard() first. For a different entry method, such as encode, use register_fsdp_forward_method(model, "encode") so that method runs the FSDP pre/post hooks. Manual unsharding isn't recursive and doesn't itself reproduce every hook's lifecycle behavior.

reshard_after_forward trades memory for a second all-gather in backward:

• True: free the full parameter after forward; all-gather again in backward.
• False: keep the unsharded parameter until backward. The usual choice for the root module, which is needed as soon as backward starts.
• None: True for non-root modules, False for the root.

An integer can instead request resharding to a smaller shard-group size, subject to the documented divisibility constraints. These options change parameter residency; “FSDP2 always frees full weights immediately after forward” is false.

View larger

For a parameter with P elements and a data-parallel shard degree d, the steady-state model copy is roughly P/d elements per rank. Peak memory still includes all-gathered parameters, activations, gradients, optimizer state, communication buckets, and temporary kernels.

Applying fully_shard only to the root groups all its managed parameters for gathering. Block-level application allows prefetch and compute to overlap, but peak memory can include several prefetched or retained blocks. Size the actual live set rather than assuming that exactly one full layer is ever resident.^{[21]Reference 21torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/2.14/distributed.fsdp.fully_shard.html}

FSDP2 and tensor parallelism solve different dimensions. FSDP2 reconstructs a dim-0-sharded parameter around use. TP keeps a layer's computation split across ranks and communicates inside or around the operation. Combining them requires matching axis names, parameter layouts, and collective order. A model can fit and still underperform if all-gathers contend with TP collectives or if the network can't hide resharding latency.

Checkpoint and restart form one recovery contract

Sharding changes checkpoint ownership. Saving one ordinary file from each rank can tie files to a particular world size and rank assignment. PyTorch Distributed Checkpoint (DCP) coordinates sharded writes and reads.

DCP coordinates writes across ranks and loads tensor state in place into destination storage. Its planner and the constructed model's target sharding enable supported changes in world size. File layout is an implementation detail, not a rule for assigning ownership to restarted ranks. State-dict helpers expose canonical parameter names across supported distributed wrappers.^{[24]Reference 24PyTorch Distributed Checkpoint (DCP) Recipehttps://docs.pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html[25]Reference 25Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html}

First construct the logical model, optimizer, and process groups, then call matching save or load code on every rank. DCP isn't torch.load(path): it fills the preallocated shards of the model you constructed.

This single-process CPU example uses DCP's real save/load and state-dict APIs. It saves a tiny linear model after one Adam update, then compares an uninterrupted second update with a fresh model and optimizer restored from disk. Fixed input removes data-order and RNG differences from this narrow test. The example checks optimizer-state recovery, not multi-rank resharding or worker restart.

import tempfile
import torch
from torch.distributed.checkpoint import load, save
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

torch.set_num_threads(1)
torch.manual_seed(7)
inputs = torch.tensor([[1.0, 2.0], [2.0, -1.0]], dtype=torch.float64)
targets = torch.tensor([[0.5], [-0.5]], dtype=torch.float64)

def make_model():
    model = torch.nn.Linear(2, 1).double()
    return model, torch.optim.Adam(model.parameters(), lr=0.01)

def update(model, optimizer):
    optimizer.zero_grad(set_to_none=True)
    loss = (model(inputs) - targets).square().mean()
    loss.backward()
    optimizer.step()
    return loss.detach().clone()

model, optimizer = make_model()
update(model, optimizer)
model_state, optim_state = get_state_dict(model, optimizer)
state = {"model": model_state, "optimizer": optim_state, "step": 1}
with tempfile.TemporaryDirectory() as directory:
    save(state, checkpoint_id=directory, no_dist=True)
    expected_loss = update(model, optimizer)
    expected_params = [p.detach().clone() for p in model.parameters()]
    expected_optim = get_state_dict(model, optimizer)[1]

    resumed, resumed_optimizer = make_model()
    model_state, optim_state = get_state_dict(resumed, resumed_optimizer)
    restored = {"model": model_state, "optimizer": optim_state, "step": 0}
    load(restored, checkpoint_id=directory, no_dist=True)
    set_state_dict(resumed, resumed_optimizer,
                   model_state_dict=restored["model"],
                   optim_state_dict=restored["optimizer"])
    assert restored["step"] == 1
    actual_loss = update(resumed, resumed_optimizer)
    torch.testing.assert_close(actual_loss, expected_loss, rtol=0, atol=0)
    for actual, expected in zip(resumed.parameters(), expected_params):
        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
    actual_optim = get_state_dict(resumed, resumed_optimizer)[1]
    for name, expected_state in expected_optim["state"].items():
        for key, expected in expected_state.items():
            torch.testing.assert_close(actual_optim["state"][name][key], expected, rtol=0, atol=0)
print("restored step: 1")
print("second update: loss, parameters, Adam moments and counters match exactly")

restored step: 1
second update: loss, parameters, Adam moments and counters match exactly

Model and optimizer tensors aren't the whole training state. A reliable checkpoint also records scheduler position, gradient-scaler state, global step, random-number-generator state, and enough sampler or data-cursor state to define which examples run next.

Save configuration and code revision beside that state. Otherwise a restart can load valid weights while silently changing learning rate, data order, or update count.

torchrun supplies worker lifecycle, not saved progress. When one worker fails, its elastic agent stops and restarts the worker group up to the configured restart limit.

Rank assignments aren't stable across restarts, and elastic membership can change world size. Recovery code must derive identity from the new environment, rebuild process groups, then load the latest complete checkpoint instead of opening a file named after an old rank.^{[26]Reference 26PyTorch 2.14: torchrun (Elastic Launch)https://docs.pytorch.org/docs/2.14/elastic/run.html}

View larger

Test recovery as a normal distributed feature, not as a final deployment ritual. Save at step $N$, kill one worker, and verify every worker restarts from checkpoint $N$.

Compare model and optimizer state, scheduler position, next data item, and first resumed loss within tolerance for the chosen deterministic settings. The latest complete, durable checkpoint determines recoverable progress; failed or delayed saves can make the replay window longer than the configured interval. Launcher retries can't recover work that was never saved.

A PyTorch LLM training path

Return to the same SDPA call. Its forward and backward route crosses several ownership boundaries. FSDP hooks surround module execution, so their parameter gathers happen before the operators below consume those parameters:

1. The data loader creates token IDs with a known device, dtype, shape, and stride.
2. The embedding lookup returns hidden states, possibly sharded across a DTensor mesh.
3. Attention calls SDPA, which selects a backend from layout, dtype, mask, and device.
4. MLP and normalization operators dispatch through ATen and may be fused by Inductor.
5. Autograd saves only what backward needs, then executes reverse nodes.
6. FSDP2's backward hooks gather any needed parameters and reduce-scatter completed gradients.
7. After the required gradient collectives finish, the optimizer updates the parameter and optimizer state owned by this rank.
8. DCP saves sharded model and optimizer state plus application progress needed for restart.

For inference, steps 5 through 8 change, while MLP and normalization operators in step 4 still dispatch through ATen and may be compiled or fused. There's no gradient tape, and parameters can stay replicated or use a serving-specific sharding plan.

Attention may use a paged KV-cache engine outside core PyTorch. torch.inference_mode() can remove metadata work in a PyTorch-native path, but a production engine such as vLLM or SGLang may own batching and cache policy instead.

Use ownership to choose the next receipt. A slow token stream might come from an engine scheduler instead of an ATen kernel. CUDA out-of-memory during FSDP2's all-gather points to a sharding peak rather than model-parameter total.

Compile graph breaks around environment calls belong to Python capture, not autograd. A logged MLflow metric doesn't tell you which of those owners failed; it only records that the step produced a number.

Strengths and weaknesses

Use PyTorch when model code, kernels, compiler choices, and distributed training need one composable substrate. Measure torch.compile for each model, inspect the costs of its stride patterns, and choose a separate production serving layer when needed. When the output is right and the step is not, which layer do you inspect first?

Failure-oriented debugging loop

When the output is correct but slow, or it changes only after sharding, start on one device and hold the input and expected output fixed. Then move one boundary at a time:

1. Reproduce on one device. Fix seed, dtype, shape, layout, and a small batch. Save input and expected output.
2. Inspect tensor contracts. Print device, dtype, shape, strides, storage offset, requires_grad, and memory format at the failing boundary.
3. Separate eager from compile. Run eager, compile only the stable inner module, then inspect graph breaks and guards.
4. Name the selected kernel. For SDPA, vary sdpa_kernel and record which path succeeds and its peak memory.
5. Check autograd versions. Remove in-place writes, compare clone and view behavior, and use gradcheck on a tiny double-precision case.
6. Reduce the distributed world. Run one rank, then two ranks with one collective. Confirm every rank enters calls in the same order.
7. Measure communication and memory. Profile all-gathers, reduce-scatters, allocator peaks, compile time, and kernel launches separately.
8. Drill restart. Kill one worker after a known checkpoint and verify group restart, resharded load, data cursor, scheduler, and resumed loss.

Use the observations to locate the failing layer:

Project identity

PyTorch began in 2016 as a collective effort from the Torch community, with substantial Meta funding and contributors from NVIDIA, Twitter, and other organizations.^{[28]Reference 28PyTorch Strengthens Its Governance by Joining the Linux Foundationhttps://pytorch.org/blog/PyTorchfoundation/} The 2019 paper by Adam Paszke, Sam Gross, Soumith Chintala, and collaborators documents the imperative, Python-first design that made dynamic model experimentation practical.^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703}

The repository contains operator code, compiler internals, distributed APIs, tests, and release work.^{[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch} The source walkthrough pins the August 2, 2026 snapshot below. Local checks use the PyTorch 2.13.0 wheel, a different build. Test accelerator behavior on the exact installed wheel. API review uses 2.14 documentation checked September 23, 2026, plus the explicitly cited moving compiler and distributed guides. These three versions serve different jobs: pinned source for reading, a pinned wheel for the CPU checks, and newer documentation for current API boundaries.

Important research roots map to present code:

The paper gives design intent. A source checkout gives behavior at its pinned revision. Read both, and record accelerator and dependency versions when measuring a path. Where, in one commit, does each of those boundaries live?

A source-reading path

The official PyTorch repository at commit 4ee8fea1f8cbb27acf0dac1c4f2e8b8bb6aef7f7 is the map.^{[8]Reference 8PyTorch Source Snapshot: 4ee8feahttps://github.com/pytorch/pytorch/tree/4ee8fea1f8cbb27acf0dac1c4f2e8b8bb6aef7f7} Start with a small, traceable operation:

1. Read torch/_tensor.py for the Python tensor façade, c10/core/TensorImpl.h for tensor metadata, and torch/overrides.py for interception hooks.
2. Open aten/src/ATen/core/dispatch/Dispatcher.h and c10/core/DispatchKey.h to see schema dispatch and key ordering.
3. Start at aten/src/ATen/native/native_functions.yaml, then follow torchgen/ to connect operator schemas to generated registrations and backend implementations.
4. Follow torch/autograd/graph.py, torch/csrc/autograd/engine.cpp, and torch/csrc/autograd/saved_variable.cpp to see backward scheduling, saved tensors, and version checks.
5. Read torch/nn/functional.py around scaled dot-product attention, then inspect backend capability checks and profiler traces.
6. Trace torch/_dynamo/eval_frame.py, torch/_functorch/aot_autograd.py, and torch/_inductor/compile_fx.py for capture, forward/backward staging, and lowering.
7. Read torch/distributed/device_mesh.py, torch/distributed/tensor/parallel/, and torch/distributed/fsdp/ for mesh, TP plans, and sharding contracts.
8. Follow torch/distributed/pipelining/ from stage construction into GPipe and 1F1B schedules.
9. Inspect torch/distributed/checkpoint/ and torch/distributed/run.py for sharded state, resharded load, worker restart, and unstable rank assignments.
10. As a separate accelerator exercise, run a tiny two-rank collective and a one-layer fully_shard model, then kill one worker after a checkpoint. The CPU examples here don't perform this distributed recovery test.

Use source paths as a map; internal names and file locations move between revisions. The invariant to carry forward is the ownership chain: tensor metadata, operator contract, dispatch path, gradient history, compiler region, and distributed placement.

What to remember

Separate a view from the copy a later kernel may still make, a dispatch key from the kernel that ran, and a restored weight from a training step that actually resumed.

• An ordinary dense, strided tensor views storage through sizes, strides, and an offset; it also carries device, dtype, and autograd metadata.
• Operator schemas and dispatch keys let one API route through backend, tracing, functionalization, and autograd behavior.
• Reverse-mode autograd records the branch actually executed; views and in-place writes are guarded by version tracking.
• SDPA backend eligibility depends on shapes, masks, strides, dtype, device, and the installed build. Inspect profiler output to verify which kernel the backend selected.
• torch.compile captures guarded regions and can stage forward/backward for lowering. Validate the intended module and sharding boundaries. Graph breaks, guard failures, and kernel compilation are different events.
• c10d collectives and DeviceMesh define communication ownership. DTensor placements make replicate, shard, and partial states explicit.
• FSDP2 defaults to dim-0 sharding, gathers around use, and follows its configured resharding policy. Peak memory includes more than the active full parameter. Call model(input) so the hooks run.
• TP splits operators, while PP splits model execution into scheduled microbatches. They move different tensors and expose different idle time.
• DCP stores sharded training state in place; torchrun restarts workers. Recovery works only when training code reloads complete state without assuming stable ranks.
• PyTorch owns tensor and training semantics. Serving schedulers, model-definition packs, rollout orchestration, and experiment tracking may belong to vLLM, Transformers, SGLang, Ray, or another layer.

Mastery check

Evaluation rubric

• Demonstrate the shared storage and changed strides of a transpose, identify when reshaping copies, and reproduce the saved-tensor mutation failure without suppressing it.
• Compare SDPA outputs and Q/K/V gradients with a reference formula. Distinguish mask semantics, backend eligibility, graph capture, guard specialization, and optimized-kernel execution.
• Explain parameter residency for the chosen FSDP2 policy and verify a restored optimizer's next update. Separate the demonstrated single-process checks from untested distributed restart and accelerator behavior.

Follow-up questions