Deep Dive - PyTorch

PyTorch is a tensor-computing and automatic-differentiation framework. Its Python expressions select implementations through several layers: tensor layout, operator dispatch, gradient recording, optional compilation, and distributed placement. Understanding those layers explains why two calls with the same output shape can take very different execution paths.

We'll trace scaled dot-product attention (SDPA). A transpose can preserve its values while changing strides; a mask can preserve its output shape while changing which tokens contribute. A training step may still produce a plausible loss in either case. The previous MLflow lesson showed how to record a run. Here we'll inspect the PyTorch operations that produced it.^{[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/stable/generated/torch.nn.functional.scaled_dot_product_attention.html}

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. They don't benchmark accelerators or validate multi-rank FSDP2. 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. Backend availability doesn't promise identical operator coverage, numerical behavior, or memory capacity.^{[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.

What is the first useful question when an LLM operation is slower than expected?

Answer

Identify its execution path: eager or compiled, which dispatch backend, which dtype and layout, which attention kernel, and whether a collective or synchronization is on the critical path.

The tensor is metadata plus storage

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.

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.

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 and reshape aren't interchangeable promises. 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.13: Tensor Viewshttps://docs.pytorch.org/docs/2.13/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. The selected key can include CUDA, CPU, autograd, functionalization, batching, fake tensors, or a tracing layer. It chooses the highest-priority applicable implementation, then may redispatch to a lower key. This layered selection lets one operator expose one public contract while several systems observe or implement it.

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:

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.

Why can a custom operator that computes correct values still be wrong?

Answer

It can violate dtype, device, stride, aliasing, mutation, autograd, or dispatch contracts. Downstream views, compiler guards, and distributed wrappers depend on those details as well as numerical values.

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.

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.^{[8]Reference 8Autograd 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."

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.

A view and its base share storage. Why does view.add_(1) sometimes fail only during backward?

Answer

The in-place update increments shared version tracking. If a backward node saved the old value, its version check detects the mutation only when that node executes. The failure can appear far from the original update.

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.

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.^{[8]Reference 8Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html}

torch.autograd.gradcheck can compare analytical and numerical derivatives for small double-precision inputs, but it doesn't prove a distributed or low-precision kernel correct.

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:

$A = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{D}} + M\right)V.$

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.

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.^{[9]Reference 9PyTorch 2.13: SDPBackendhttps://docs.pytorch.org/docs/2.13/generated/torch.nn.attention.SDPBackend.html[10]Reference 10torch.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 don't guarantee the fastest implementation for every 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/stable/generated/torch.nn.functional.scaled_dot_product_attention.html[10]Reference 10torch.backends.cuda: SDPA backend eligibilityhttps://docs.pytorch.org/docs/stable/backends.html[11]Reference 11FlashAttention-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. The documented support is Flash and math kernels on CUDA, with query-head count divisible by key/value-head count, equal key and value head counts, and no nested tensors. Treat those as documented support limits, not a proof that every other combination must raise.
• 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/stable/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. It doesn't exercise CUDA fused attention.

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.^{[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

For an attention input with S=8192, why is avoiding the materialized score matrix more important than shaving one Python function call?

Answer

The score matrix grows as S². A tiled fused kernel reduces peak intermediate memory and memory traffic, while Python dispatch overhead is nearly constant with sequence length.

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.^{[12]Reference 12PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilationhttps://doi.org/10.1145/3620665.3640366[13]Reference 13torch.compilehttps://docs.pytorch.org/docs/main/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.

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.^{[14]Reference 14Where 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. The official troubleshooting guidance recommends reducing the program to a small reproducer and inspecting graph breaks and guards.^{[15]Reference 15Troubleshooting torch.compilehttps://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_troubleshooting.html}

Then decide whether to rewrite code, allow a fallback, or compile only stable regions. fullgraph=True is a diagnostic constraint, not a universal production setting: it makes breaks visible but can reject valid programs that would run with partial capture.

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.

What does a graph break mean for correctness, and what does it usually cost?

Answer

It usually preserves correctness by returning to eager Python for an unsupported region. It costs optimization opportunity, extra boundary overhead, and sometimes synchronization. It isn't proof that model math is wrong.

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 affect correctness. Follow the version-specific guidance for mixed precision and compare gradients, not only forward outputs.^{[16]Reference 16torch.compile Has Different Autograd Semanticshttps://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_backward.html} Our double-precision CPU example avoids autocast; it doesn't validate an automatic mixed precision (AMP) training configuration.

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.^{[17]Reference 17PyTorch 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.^{[17]Reference 17PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

Read each placement as a promise about where a logical tensor lives, not as an automatic performance guarantee. DTensor adds a global tensor view plus a placement on mesh dimensions. The three placements are:^{[17]Reference 17PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

• Replicate: every rank stores the global value.
• Shard(dim): each rank stores a slice of one global dimension, following torch.chunk(dim) semantics.
• Partial: each rank stores a contribution that must be reduced before treating it as 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. This is a shape-and-placement contract, not a guarantee that every operation avoids communication.

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.

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.^{[18]Reference 18torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/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.^{[17]Reference 17PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

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.^{[19]Reference 19PyTorch 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, 1F1B, interleaved 1F1B, and looped 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.^{[17]Reference 17PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

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. Configured loss scaling must match whether each microbatch loss is averaged or summed.

Why can't adding pipeline stages fix a tensor-parallel collective bottleneck by itself?

Answer

Pipeline parallelism moves activations and gradients between consecutive model stages. Tensor parallelism still communicates inside each sharded layer, so its collective cost remains unless the tensor layout, mesh, or operator plan changes.

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.^{[18]Reference 18torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html}

Around module execution, pre-forward and pre-backward hooks all-gather a parameter and register a plain torch.Tensor for compute. After the module runs, the unsharded copy can be freed and the DTensor shard registered again. Gradients are reduce-scattered so each rank keeps its optimizer-owned slice.^{[18]Reference 18torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/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 or register that method with register_fsdp_forward_method.

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.

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.^{[18]Reference 18torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/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.

Why does FSDP2 not mean every parameter is permanently absent from every rank?

Answer

The owning module needs a full parameter view for its local computation, so FSDP2 all-gathers around the forward or backward and can reshard afterward. The memory win comes from limiting how much is materialized at one time and sharding optimizer and gradient state.

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.^{[20]Reference 20PyTorch Distributed Checkpoint (DCP) Recipehttps://docs.pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html[21]Reference 21Getting 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.

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.^{[3]Reference 3PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

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 worker dies after step 1,240, and the latest complete checkpoint is step 1,200. What should torchrun restore automatically?

Answer

torchrun restarts the worker group, but training code must load checkpoint 1,200. Up to 40 completed steps may be replayed or lost, depending on data-cursor handling. Rank IDs from the failed group must not determine restored ownership.

A PyTorch LLM training path

Return to the same SDPA call. Its forward and backward route crosses several ownership boundaries:

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 all-gathers parameters, runs a module, and reduce-scatters gradients.
7. The optimizer updates local state, while c10d synchronizes the replicas required by the chosen parallelism.
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

PyTorch is a strong default when model code, kernels, compiler choices, and distributed training need one composable substrate. It doesn't promise that torch.compile improves every model, that every stride is cheap, or that a high-level module owns production serving. Pick the layer that matches the bottleneck.

Failure-oriented debugging loop

When output is correct but slow, or changes only after sharding, hold input and expected output fixed. 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, incubated and heavily funded by Meta's Facebook AI Research group, with important contributions from NVIDIA, Twitter, and other organizations. 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, which is a different build; source inspection isn't proof of that wheel's accelerator behavior. API references distinguish versioned 2.13 pages from the explicitly cited moving compiler guides.

Important research roots map to present code:

The paper gives design intent. The source checkout gives current behavior. Read both, pin the repository commit, and record accelerator and dependency versions when measuring a path.

A source-reading path

This walkthrough uses the official PyTorch repository at commit 4ee8fea1f8cbb27acf0dac1c4f2e8b8bb6aef7f7.^{[19]Reference 19PyTorch 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, not as a stable public API guarantee. Internal names and file locations move. 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

• 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. A backend preference isn't proof of the selected kernel.
• 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

A transpose has the right values and shape, but attention slows down. What observation distinguishes a view from a subsequent layout copy?

Answer

Shared storage pointers and changed strides demonstrate the view. They don't show whether the next operator copied. Inspect allocation and profiler events around that operator, then compare a supported contiguous input without changing the attention mask or values. Time the conversion as part of the path, not outside the measurement.

The CPU checkpoint example passes. What remains before claiming recovery after losing a distributed worker?

Answer

Run the actual multi-rank sharding configuration, save a complete durable checkpoint, fail a worker, and rebuild the group. Verify model and optimizer shards, scheduler and scaler state, RNG, data position, and the first resumed update. If world size changes, test the supported resharding path and data-assignment semantics explicitly. A single-process reload doesn't establish those properties.

Can the two captured graphs in the compiler example establish a speedup?

Answer

No. The counting backend executes its graphs eagerly, and aot_eager checks staged autograd without Inductor optimization. Benchmark the chosen optimized backend separately, including startup and warm execution, and compare outputs and gradients under realistic shapes, strides, and autocast settings.