This wiki page provides a high-level introduction to the PyTorch repository. PyTorch is a comprehensive machine learning framework built around flexible tensor computations and dynamic computation graphs, supporting GPU acceleration and multiple hardware backends. The repository is a complex monorepo containing multiple subsystems spanning Python frontends and extensive C++ core libraries.
Major aspects covered include:
Compilation system: Encompasses torch.compile API, TorchDynamo bytecode analysis, AOTAutograd for joint forward-backward graph tracing and functionalization, and the TorchInductor lowering and code generation backend.
Distributed training: Covers collective communication infrastructure (c10d), the DTensor distributed tensor abstraction with placement and sharding semantics, and sharding-based distributed training via Fully Sharded Data Parallel (FSDP).
Device backends: Includes implementations for CUDA, Apple MPS, Intel XPU, and the ATen native operation dispatch framework for routing operators to backends.
Build/Test infrastructure: Describes the repository’s build system (CMake, Python setups), code generation for ATen ops (torchgen), testing frameworks including OpInfo for operator coverage, and CI workflows.
This page provides an architectural overview and points to child pages for detailed exploration and onboarding.
Sources: torch/_dynamo/output_graph.py1-20 torch/_inductor/ir.py168-207 torch/_inductor/scheduler.py1-40 .ci/pytorch/test.sh1-170
PyTorch translates user-defined dynamic Python code into efficient, device-optimized machine code through a multi-stage compilation and dispatch pipeline.
This pipeline begins with the user’s Python code (typically nn.Module or Python functions). Invoking torch.compile() triggers TorchDynamo to intercept Python bytecode evaluation for symbolic execution and graph construction. The resulting FX graph is processed via AOTAutograd to produce a joint forward-backward trace with functionalization of in-place operations.
Next, the graph is lowered in TorchInductor into an intermediate representation (IR) composed of tensor and buffer abstractions, followed by scheduling that fuses operations and optimizes memory usage. The final step generates optimized kernels with backends such as Triton (GPU) or compiled C++ (CPU), which are then executed.
Sources: torch/__init__.py10-25 torch/_dynamo/output_graph.py5-15 torch/_inductor/compile_fx.py1-45 torch/_inductor/lowering.py120-145 torch/_inductor/scheduler.py10-50
TorchDynamo: The frontend bytecode interpreter uses eval_frame to capture Python execution frames. It employs classes like InstructionTranslator and VariableTracker to perform symbolic execution of operations, building an OutputGraph representing the computation with attached runtime guards to ensure validity torch/_dynamo/output_graph.py5-20
torch.export: Provides a static graph export mechanism producing ExportedProgram artifacts. It enforces strict symbolic shape constraints via ShapeEnv and SymInt abstractions for dynamic dimensions. It leverages FakeTensorMode to propagate metadata during export torch/_dynamo/output_graph.py70-85
AOTAutograd: Ahead-of-time joint tracing of forward and backward passes, including graph functionalization to remove inplace mutations. It partitions the graph for efficient memory use and compiles with caching support torch/_inductor/compile_fx.py10-45
TorchInductor: The default backend compiler that lowers FX graphs into its own tensor IR (TensorBox, Buffer) and applies scheduling passes for fusion and memory planning torch/_inductor/lowering.py120-145 It then generates device-specific code such as Triton GPU kernels and compiled C++ wrappers torch/_inductor/runtime/triton_heuristics.py50-135
For more details, see Compilation System.
ATen native functions: PyTorch’s core tensor operations are registered via the ATen dispatcher using native_functions.yaml schemas. This system dispatches operator calls by device and dtype aten/src/ATen/native/native_functions.yaml
CUDA backend: Includes device management, memory allocation by the CUDACachingAllocator with sanity checks, CUDA graphs for replaying workloads, and optimized BLAS integration for GEMM operations, including newer precisions like FP8 torch/_inductor/scheduler.py117 torch/_inductor/runtime/triton_heuristics.py130-135
MPS backend: Supports Apple Silicon GPU devices implementing key operations and compiling Metal shaders for GPU kernel execution aten/src/ATen/mps/EmptyTensor.cpp
XPU backend: Intel’s GPU backend integrated with corresponding allocator, device/stream APIs, and Inductor support for Triton-XPU integration.
For more, see Device Backends and Native Operations.
c10d: The fundamental collective communication library supports stores, process groups (NCCL, Gloo, etc.), and asynchronous collective operations torch/distributed/c10d/process_group.cpp
DTensor abstraction: Enables single-program-multiple-data (SPMD) style distributed tensors with placement specifications like Shard, Replicate, and Partial. It manages sharding propagation, communication insertion (all-reduce, all-gather), and cost-based strategy selection for distributed execution torch/distributed/_tensor/dtensor.py
Fully Sharded Data Parallel (FSDP): Implements parameter sharding, communication scheduling for all-gather/reduce-scatter, and optimizer state sharding for efficient large-scale training torch/distributed/fsdp/fsdp.py
Pipeline parallelism: Supports pipelined model parallelism via stages, with scheduling modes like 1F1B and GPipe, compatible with FSDP and DTensor abstractions.
For more details, see Distributed Training Systems.
The repository employs a flexible and powerful build and testing infrastructure:
Uses setup.py for Python packaging and CMake for native code compilation. The build scripts configure numerous feature flags (CUDA, ROCm, XPU, etc.) and handle dependencies .ci/pytorch/build.sh1-180
torchgen: The ATen native operator code generator translates operator schema definitions from YAML into dispatch and registration C++ boilerplate tools/codegen/torchgen
The OpInfo framework drives comprehensive operator testing across device types and dtypes, with systematic test generation and automatic skipping/exclusion test/inductor/test_torchinductor.py85-100
CI workflows use Docker images that encapsulate build environments targeting various OSes and hardware (CUDA, ROCm, XPU). Specialized shards for Inductor, Dynamo, and distributed subsystems ensure targeted testing .github/workflows/pull.yml1-150
Automated binary release pipelines produce platform-specific wheels and conda packages with cross-compilation support .ci/docker/build.sh1-150
| System | Primary Entry Point | Code Areas |
|---|---|---|
| Compilation | torch.compile | torch/_dynamo/, torch/_inductor/ |
| IR and Lowering | torch/_inductor/ir.py | torch/_inductor/lowering.py |
| Distributed | torch.distributed | torch/distributed/, torch/_inductor/scheduler.py |
| Device backends | torch.cuda, torch.mps | torch/_inductor/codegen/ |
| Runtime | torch/_inductor/runtime | triton_heuristics.py, triton_helpers.py |
For onboarding and detailed repo layout, see
Refresh this wiki
This wiki was recently refreshed. Please wait 1 day to refresh again.