Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
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][2]
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:
- Local reasoning. A line of Python should execute now, produce a normal traceback, and allow a debugger or print statement.
- 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][2]
The project is more than torch.Tensor:
| Layer | Responsibility | LLM consequence |
|---|---|---|
| Python API | Modules, optimizers, tensor expressions, user ergonomics | Model code stays inspectable and composable |
| ATen operators | Schemas, tensor kernels, shape and dtype contracts | mm, view, softmax, and attention share one operator surface |
| Dispatcher and c10 | Select backend, autograd, functionalization, and other keys | One call can route to CUDA, a sparse backend, or a tracing layer |
| Autograd | Build a dynamic graph and execute reverse-mode derivatives | Training can branch on data and still backpropagate |
| Compiler stack | Capture, transform, and lower Python programs | Stable regions can become fused kernels |
| c10d and distributed | Process groups, collectives, meshes, sharding | Parameters and activations can span a GPU cluster |
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.
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. "It is PyTorch" is not a diagnosis.
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;
TensorImplcarries 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.
1import torch
2
3x = torch.arange(6, dtype=torch.float32).reshape(2, 3)
4y = x.transpose(0, 1)
5print(x.shape, x.stride(), x.storage_offset())
6print(y.shape, y.stride(), y.storage_offset())1torch.Size([2, 3]) (3, 1) 0
2torch.Size([3, 2]) (1, 3) 0The 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.
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
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.
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, 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] 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:
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][5] 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:
| Signal | What it can reveal |
|---|---|
| Q/K/V strides | A transpose or reshape inserted a copy before the kernel |
| Dtype and head dimension | Hardware kernel constraints or fallback to math path |
| Causal or arbitrary mask | Supported fused path versus generic masked matmul |
| Sequence length | Tiling, memory pressure, and launch overhead |
| Autograd state | Training backward path versus inference-only fast path |
| Kernel logs or profiler | Actual implementation, rather than preferred implementation |
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.
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. 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][7]
The sequence is easier to reason about as four contracts:
- Capture. Dynamo specializes on Python values, tensor metadata, and control-flow assumptions. It emits guards that must hold for a cached graph.
- 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.
- AOTAutograd. Forward and backward graphs are staged so compiler passes can see saved values and gradient dependencies.
- 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] 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.
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 is not proof that model math is wrong.
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]
The primitive operations are familiar:
| Collective | Result | LLM use |
|---|---|---|
all_reduce | Every rank receives reduced value | Data-parallel gradient sum or TP partial output |
reduce_scatter | Reduce then give each rank one slice | FSDP2 gradient sharding |
all_gather | Every rank receives all slices | Rebuild a parameter or activation |
all_to_all | Each rank exchanges distinct slices | MoE token dispatch or sequence partition |
broadcast | One source sends to all | Initial state or control metadata |
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]
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][2]
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][2]
These axes solve different limits and compose through different communication:
| Axis | What each rank owns | Main communication | Typical reason to use it |
|---|---|---|---|
| FSDP2 / data | Parameter, gradient, and optimizer shards | Parameter all-gather, gradient reduce-scatter | Full model state doesn't fit per rank |
| Tensor | Slice of one layer's matrix or activation | All-reduce, reduce-scatter, or all-gather inside layer | One layer's compute or weights are too large |
| Pipeline | Consecutive model stages | Activations forward, gradients backward | Model depth spans devices or slower links |
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
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]
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.
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, 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][2]
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.
1from torch.distributed.checkpoint import load, save
2from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
3
4model_state, optim_state = get_state_dict(model, optimizer)
5state = {"model": model_state, "optimizer": optim_state, "step": step}
6save(state, checkpoint_id=checkpoint_dir)
7
8# After constructing the same logical model and optimizer on restarted workers:
9model_state, optim_state = get_state_dict(model, optimizer)
10restored = {"model": model_state, "optimizer": optim_state, "step": 0}
11load(restored, checkpoint_id=checkpoint_dir)
12set_state_dict(
13 model,
14 optimizer,
15 model_state_dict=restored["model"],
16 optim_state_dict=restored["optimizer"],
17)
18step = restored["step"]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]
Test recovery as a normal distributed feature. Save at step , kill one worker, and verify every worker restarts from checkpoint . 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 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
An LLM training step crosses several ownership boundaries:
- The data loader creates token IDs with a known device, dtype, shape, and stride.
- The embedding lookup returns hidden states, possibly sharded across a DTensor mesh.
- Attention calls SDPA, which selects a backend from layout, dtype, mask, and device.
- MLP and normalization operators dispatch through ATen and may be fused by Inductor.
- Autograd saves only what backward needs, then executes reverse nodes.
- FSDP2 all-gathers parameters, runs a module, and reduce-scatters gradients.
- The optimizer updates local state, while c10d synchronizes the replicas required by the chosen parallelism.
- 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.
| Application | PyTorch owns | Another system may own |
|---|---|---|
| Fine-tuning | Model modules, autograd, optimizer, distributed collectives | Dataset streaming, experiment tracking |
| Pretraining | Kernels, graph execution, sharding, distributed checkpoint state | Cluster placement and durable storage service |
| Native inference | Tensor math, SDPA, quantized modules | HTTP, batching, KV cache, autoscaling |
| vLLM or SGLang serving | Custom model components or exported weights | Request scheduling, KV blocks, sampling |
| RL post-training | Policy loss, gradient update, model state | Rollout orchestration, reward service, trainer mesh |
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
| Dimension | Strength | Boundary or weakness |
|---|---|---|
| Research loop | Eager Python, inspectable state, dynamic control flow | Different eager and compiled paths can diverge in performance or diagnostics |
| Operator ecosystem | One schema reaches CPU, CUDA, autograd, tracing, and custom backends | Dispatch and aliasing contracts are deep and easy to violate in extensions |
| Performance | Fused kernels, SDPA backends, compiler lowering, and memory allocators | Small or irregular workloads can lose to compile and launch overhead |
| Autograd | Flexible reverse-mode tape and custom functions | In-place mutation, views, and saved tensors create sharp edges |
| Distributed scale | c10d, DeviceMesh, DTensor, FSDP2, TP, PP, DCP, and elastic workers | Collective ordering, topology, state completeness, and alpha APIs widen failure surface |
| Ecosystem | Broad model, accelerator, and tooling support | Version and backend combinations expand the test matrix |
| Deployment | Export paths, inference mode, quantization, and custom ops | Native PyTorch doesn't automatically provide a serving scheduler |
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:
- Reproduce on one device. Fix seed, dtype, shape, layout, and a small batch. Save input and expected output.
- Inspect tensor contracts. Print device, dtype, shape, strides, storage offset,
requires_grad, and memory format at the failing boundary. - Separate eager from compile. Run eager, compile only the stable module, then inspect graph breaks and guards.
- Name the selected kernel. For SDPA, vary backend controls and record which path succeeds and its peak memory.
- Check autograd versions. Remove in-place writes, compare
cloneand view behavior, and use gradcheck on a tiny double-precision case. - Reduce the distributed world. Run one rank, then two ranks with one collective. Confirm every rank enters calls in the same order.
- Measure communication and memory. Profile all-gathers, reduce-scatters, allocator peaks, compile time, and kernel launches separately.
- 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:
| Symptom | First suspect | Useful proof |
|---|---|---|
| Unexpected copy before attention | Stride or memory format | Compare .stride() and profiler copy event |
| "Modified by an inplace operation" | View alias or saved tensor version | Remove mutation and run anomaly detection |
| Compile produces many variants | Shape, dtype, stride, or Python guard | Recompilation report and guard list |
| Fused SDPA unavailable | Mask, dtype, head dimension, or device | Backend enablement and fallback log |
| FSDP2 OOM during layer entry | Parameter all-gather peak | Peak allocation around fully_shard module |
| Distributed hang | Rank divergence or collective mismatch | Per-rank logs, timeout, and one-collective reproducer |
| Restart begins at step 0 or repeats data | Missing application state or wrong checkpoint path | Restored step, scheduler, RNG, and data-cursor receipt |
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]
| Field | Current project fact |
|---|---|
| Origin | Torch community work became PyTorch through Meta/FAIR incubation and broad external collaboration.[1] |
| Stewardship | The Linux Foundation hosts the PyTorch Foundation. Technical authority belongs to individual maintainers working through project governance, not permanent company seats.[12][13] |
| Contributor path | Public issues, pull requests, maintainer review, code ownership, and technical decision processes provide the contribution path.[13] |
| Project licenses | Main source is BSD 3-Clause. Distributed packages include separately licensed components, so installed license reports contain BSD, Apache, MIT, LLVM-exception, and other notices.[14][15] |
| Commercial boundary | Companies fund contributors and sell products built with PyTorch, but that doesn't turn the foundation project into one vendor's product. |
| Asset boundary | PyTorch's source license doesn't grant rights to model weights, datasets, extensions, or services used with it. Review each asset separately. |
The current repository remains the executable source of truth for operator code, compiler internals, distributed APIs, tests, and release work.[2] Read its governance and license files before describing ownership or redistribution rights.
Important research roots map to present code:
| Research or source | Core idea | Where to connect it |
|---|---|---|
| Paszke et al. (2019) | Imperative tensors plus tape-based autograd | Python API, ATen, autograd, and extension model[1] |
| PyTorch 2 paper (2024) | Preserve eager semantics while adding graph capture and compiler lowering | Dynamo, AOTAutograd, Inductor, and graph-break diagnostics[6] |
| FlashAttention-2 | Tiled attention with lower memory traffic and improved parallel work partitioning | SDPA backend selection and custom attention kernels[5] |
| FSDP2 documentation | Per-parameter sharding around module execution | fully_shard, DTensor placements, and resharding behavior[10] |
| Current distributed docs | Process groups, DeviceMesh, DTensor, and collective contracts | torch.distributed, _tensor, and mesh APIs[9] |
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] Start with a small, traceable operation:
- Read
torch/_tensor.pyfor the Python tensor façade andtorch/overrides.pyfor interception hooks. - Open
aten/src/ATen/core/dispatch/Dispatcher.handc10/core/DispatchKey.hto see schema dispatch and key ordering. - Search generated operator registrations under
aten/src/ATenandtorchgen/to connect schemas to backend implementations. - Follow
torch/autograd/graph.pyand the C++ engine to see nodes, saved tensors, version checks, and backward scheduling. - Read
torch/nn/functional.pyaround scaled dot-product attention, then inspect backend capability checks and profiler traces. - Trace
torch/_dynamo/eval_frame.py,torch/_functorch/aot_autograd.py, andtorch/_inductor/compile_fx.pyfor capture, forward/backward staging, and lowering. - Read
torch/distributed/device_mesh.py,torch/distributed/tensor/parallel/, andtorch/distributed/fsdp/for mesh, TP plans, and sharding contracts. - Follow
torch/distributed/pipelining/from stage construction into GPipe and 1F1B schedules. - Inspect
torch/distributed/checkpoint/andtorch/distributed/run.pyfor sharded state, resharded load, worker restart, and unstable rank assignments. - Run a tiny two-rank collective and a one-layer
fully_shardmodel, 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.compilecaptures 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.