Deep Dive - PyTorch

PyTorch looks like a Python library because model authors start there. A transformer call quickly crosses a stack of tensor metadata, generated operators, dispatch keys, autograd nodes, compiler graphs, and collectives. Ask which contract owns each value and which path will execute it, rather than reducing the operation to "Python or C++."^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

This lesson follows one LLM-shaped operation through that stack. A query, key, and value tensor enter scaled dot-product attention. Eager mode resolves an operator schema and a backend kernel. Autograd records a backward edge when gradients are enabled. torch.compile may capture the Python frame, partition forward and backward graphs, then lower them to generated kernels. FSDP2 and DTensor may change where the tensor's storage lives without changing the model author's shape-level code.

Why PyTorch became the substrate

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 original paper describes a Python-first library with GPU acceleration and dynamic neural networks, while the current repository includes 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[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

The project is more than torch.Tensor:

There is no single "PyTorch performance number." Eager execution, compiled execution, fused kernels, memory layout, device, sequence length, and communication topology all change the path.

The tensor is metadata plus storage

At the Python boundary, torch.Tensor is a handle to an internal tensor object. In the C++ core, TensorImpl carries the tensor's metadata and points to storage. Treat a tensor as a view over storage, not a boxed multidimensional array copied on every operation.

• 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.

For a contiguous tensor with shape [2, 3] and element strides [3, 1], element (i, j) lives at offset + 3i + j. A transpose can expose shape [3, 2] with strides [1, 3] without moving bytes. A kernel that requires contiguous input may materialize a copy; a stride-aware kernel can read the view directly.

That distinction matters for attention. Q, K, and V often begin as [batch, sequence, heads, head_dim], then become [batch, heads, sequence, head_dim] with a permutation. The shape changes instantly. The next matrix multiplication may still pay for a contiguous conversion if its kernel can't consume the resulting strides.

The output is a small contract check, not a benchmark. x and y share storage, but they present different logical coordinates. A training step that mutates one can affect the other. A compiler guard may also specialize on the stride pattern and recompile when a later batch arrives with a different layout.

Views are why reshape and view are not interchangeable promises. view requires a compatible stride pattern. reshape can return a view or allocate a copy. Use tensor._base only for debugging, not as application logic. In production, inspect shape, stride, is_contiguous(), and memory format at the boundary where a kernel slows down.

Operator schemas and the dispatcher

PyTorch operators are registered with schemas that describe names, arguments, returns, overloads, defaults, and mutation or aliasing behavior. Many generated Python and C++ bindings lead back to an ATen operator. A call such as torch.mm(x, w) is therefore a typed operator invocation, not an arbitrary Python function with opaque side effects.

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.

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

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 no longer needed unless the graph is retained.

The 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. There is no separate graph-building phase in ordinary eager mode.

Views and in-place mutation

Autograd must detect when a saved value changes before backward uses it. Tensors carry version counters. An in-place operation increments the counter; backward checks the expected version and raises when the saved value is stale. Views share storage and often share version tracking with their base, so mutating a view can invalidate a backward node that saved the base.

This is the source of 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.

detach() creates a tensor that shares storage but stops gradient history. Use it when a value mustn't contribute to a later gradient, but remember that in-place writes through detached aliases can still mutate the original storage. clone().detach() separates both history and bytes when isolation is required.

torch.no_grad() tells autograd not to record operations in a dynamic region. Inference mode goes further by using inference tensors with stricter aliasing and metadata behavior; tensors created there shouldn't be moved back into a grad-tracked computation casually.^{[3]Reference 3Autograd 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

LLM attention makes PyTorch's layers visible. 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.$

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, which is 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 exposes one API while selecting among math, memory-efficient, and FlashAttention-style implementations based on device, dtype, shape, mask, causal mode, and backend availability.^{[4]Reference 4torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html[5]Reference 5FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691} Supported SDPA backends provide autograd paths, but backward availability, memory use, and numerical behavior depend on backend, device, dtype, shape, and mask. A backend control context can force or disable implementations during diagnosis.

The LLM owner still needs to check:

A fused attention kernel is not automatically faster for every shape. Very short sequences can be launch-bound. An unsupported mask can force a fallback. A layout conversion can cost more than the fused kernel saves. Record the selected backend and shape when comparing runs.

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

The sequence is easier to reason about as four contracts:

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.

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, inspecting graph breaks and guards, and deciding whether to rewrite code, allow a fallback, or compile only stable regions.^{[8]Reference 8Troubleshooting torch.compilehttps://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_troubleshooting.html} 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.

Compile can hurt. A tiny model may spend more time compiling than executing. Dynamic shapes can trigger repeated compilation. 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.^{[9]Reference 9PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

The primitive operations are familiar:

Collectives have ordering contracts. If rank 0 enters an all-reduce while rank 1 enters an all-gather on the same group, the job can hang or report a transport error. 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.^{[9]Reference 9PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html}

DTensor adds a global tensor view plus a placement on mesh dimensions. Common placements are:

• Replicate: every rank stores the global value.
• Shard(dim): each rank stores a slice of one global dimension.
• 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. A later redistribution to replicated or another shard placement makes the required communication visible in the abstraction. 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.

Tensor and pipeline parallelism split different work

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 MLP can shard its first projection by output columns, keep the expanded activation split, then shard its second projection by input rows. The rank-local matrix multiplications are smaller, but their layouts must agree on where an all-reduce, reduce-scatter, or all-gather belongs.^{[9]Reference 9PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

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 and one-forward-one-backward (1F1B). Current documentation labels the package alpha, so pin the PyTorch version and test schedule changes before relying on its API.^{[9]Reference 9PyTorch Distributed Communication Packagehttps://docs.pytorch.org/docs/main/distributed.html[2]Reference 2PyTorch 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.

FSDP2 shards parameters around the forward

FSDP2's composable API, exposed through fully_shard, uses per-parameter sharding and a device mesh to avoid keeping full model state on every rank. A parameter is sharded when idle and all-gathered before its owning module computes, then released or resharded according to reshard_after_forward. Non-root modules default to resharding, while a root module defaults to keeping parameters unsharded after forward. Gradients are reduce-scattered so each rank keeps its optimizer-owned slice.^{[10]Reference 10torch.distributed.fsdp.fully_shardhttps://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html}

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, but peak memory includes all-gathered parameters, activations, gradients, optimizer state, communication buckets, and temporary kernels. A layer-by-layer schedule changes the requirement from "the full model must fit" to "one active module plus its working set must fit."

FSDP2 and tensor parallelism solve different dimensions. FSDP2 shards a parameter across a mesh axis and reconstructs it 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, while its state-dict helpers expose canonical parameter names across FSDP2, distributed data parallelism, and tensor parallelism. A later load can read only shards needed by each rank and reshard state for a different trainer count or parallel layout.^{[11]Reference 11Getting Started with Fully Sharded Data Parallel (FSDP2)https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch}

The core save and load shape appears below after model, optimizer, and process groups exist. Every rank must call this distributed skeleton with matching keys.

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

Test recovery as a normal distributed feature. 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. Checkpoint interval sets the maximum recomputation or replay window; launcher retries can't recover work that was never saved.

A PyTorch LLM training path

An LLM training step 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 is no gradient tape, parameters can stay replicated or use a serving-specific sharding plan, and 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.

The boundary matters when debugging. 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.

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

Use a narrow loop instead of toggling flags randomly:

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 module, then inspect graph breaks and guards.
4. Name the selected kernel. For SDPA, vary backend controls 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.

Common symptoms map to different owners:

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 current repository remains the executable source of truth for operator code, compiler internals, distributed APIs, tests, and release work.^{[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch} Read its governance and license files before describing ownership or redistribution rights.

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.^{[2]Reference 2PyTorch Source Repositoryhttps://github.com/pytorch/pytorch} Start with a small, traceable operation:

1. Read torch/_tensor.py for the Python tensor façade 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. Search generated operator registrations under aten/src/ATen and torchgen/ to connect schemas to backend implementations.
4. Follow torch/autograd/graph.py and the C++ engine to see nodes, saved tensors, version checks, and backward scheduling.
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. Run a tiny two-rank collective and a one-layer fully_shard model, then kill one worker after a checkpoint before attempting a full transformer.

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

• A tensor is a view over storage with sizes, strides, offset, 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 can choose math or fused kernels. Shapes, masks, strides, dtype, and device decide whether the fast path is legal.
• torch.compile captures guarded regions, stages forward and backward, and lowers stable graphs. Graph breaks and recompiles are expected diagnostics, not automatic model failures.
• c10d collectives and DeviceMesh define communication ownership. DTensor placements make replicate, shard, and partial states explicit.
• FSDP2 limits parameter and optimizer residency by all-gathering around use and resharding afterward. Peak memory still includes active parameters and buffers.
• 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; torchrun restarts workers. Recovery works only when training code reloads complete state without assuming stable ranks.
• PyTorch owns tensor and training semantics. Serving schedulers, rollout orchestration, and experiment tracking may belong to vLLM, SGLang, Ray, slime, or another layer.