Autograd

Active contributors: albanD, soulitzer

Purpose

autograd is PyTorch's reverse-mode automatic differentiation engine. It records every differentiable operation onto a dynamically-built DAG (the "tape") and walks it backward when you call .backward() or torch.autograd.grad. Unlike static-graph systems (TF1, Theano), the tape is recreated every forward pass, so control flow can change between calls.

The autograd implementation is split across:

• Python frontend — torch/autograd/ — user-facing functions (grad, backward, Function, hooks, anomaly detection, profiler).
• C++ engine — torch/csrc/autograd/ — the actual graph data structures and the multi-threaded backward executor.
• Generated code — torch/csrc/autograd/generated/ — per-op forward/backward kernels and Python bindings produced by torchgen from tools/autograd/derivatives.yaml.

Directory layout

Path	Contents
torch/autograd/	Python frontend
torch/autograd/__init__.py	backward, grad, top-level utilities
torch/autograd/function.py	torch.autograd.Function user API
torch/autograd/profiler.py	Legacy autograd profiler
torch/autograd/gradcheck.py	gradcheck and gradgradcheck
torch/autograd/forward_ad.py	Forward-mode autodiff (dual numbers)
torch/csrc/autograd/	C++ engine
torch/csrc/autograd/engine.cpp	The multithreaded backward executor
torch/csrc/autograd/function.h	Node (a.k.a. Function) base class for backward graph nodes
torch/csrc/autograd/variable.cpp	AutogradMeta per-tensor metadata
torch/csrc/autograd/saved_variable.cpp	Saved-tensor mechanics
torch/csrc/autograd/python_function.cpp	Python-side Function glue
torch/csrc/autograd/profiler*.cpp	Profiler integration
tools/autograd/derivatives.yaml	The list of differentiation rules
tools/autograd/templates/	Templates that torchgen renders into per-op forward+backward kernels

Key abstractions

Type	File	Purpose
torch::autograd::Node	torch/csrc/autograd/function.h	Backward graph node; produces grad inputs from grad outputs
torch::autograd::AutogradMeta	torch/csrc/autograd/variable.h	Per-tensor metadata: grad_, grad_fn_, version_counter_
torch::autograd::SavedVariable	torch/csrc/autograd/saved_variable.h	Tensor saved for backward; tracks version counter
torch::autograd::Engine	torch/csrc/autograd/engine.h	The multithreaded backward executor
torch.autograd.Function	torch/autograd/function.py	User-facing extension hook
torch::autograd::Variable	torch/csrc/autograd/variable.h	Historical alias — now identical to at::Tensor
*Backward0 classes	torch/csrc/autograd/generated/Functions.h	Per-op generated backward Nodes

How it works

Forward: building the tape

When a differentiable op (e.g., at::matmul) is called with inputs that have requires_grad=True, the dispatcher routes it to the Autograd* kernel for that op. The autograd kernel:

1. Saves whatever inputs/intermediates the backward will need (SavedVariable).
2. Constructs a *Backward0 Node (e.g., MatmulBackward0) holding those saved values.
3. Sets the result tensor's grad_fn_ to that node.
4. Calls at::redispatch::matmul(...) with the Autograd* keys removed, so a backend kernel runs.

The "edges" of the DAG come from the saved next_edges_: each Node holds a list of (Node, output_index) pairs pointing to its parents. This is built up automatically from the requires_grad flags of the inputs.

Backward: walking the tape

tensor.backward() and torch.autograd.grad(outputs, inputs) both call Engine::execute_with_graph_task. The engine:

graph TD
    Start[backward call] --> Init[Build GraphTask<br/>from output tensors]
    Init --> Queue[Push roots into ready queue]
    Queue --> Pop[Pop a Node]
    Pop --> Run[Call node.apply(grad_in)<br/>returns grad_out]
    Run --> Accum[Accumulate grad_out into parents]
    Accum --> Check{Parent ready?}
    Check -- yes --> Queue
    Check -- no --> Pop
    Run --> Done{All nodes done?}
    Done -- yes --> Finish[Write grads<br/>into leaf .grad]

The engine runs multiple worker threads, one per device, with per-thread ready_queues_. CUDA backwards run on a CUDA worker that submits kernels to the device's autograd stream. CPU work runs on a CPU worker pool. Synchronization between devices is handled by InputBuffer and the GraphTask's dependencies_ map.

Saved tensors and version counters

When a tensor is saved for backward, autograd records its version_counter_. If that storage is mutated in place after the save, the version counter increments and backward will raise RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation. This is the layer that prevents silently-wrong gradients from in-place ops.

torch.autograd.graph.saved_tensors_hooks lets a user run code on save/load, used by activation checkpointing and CPU offload.

torch.autograd.Function extension

User code can extend autograd by subclassing torch.autograd.Function:

class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.sin()

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return grad_out * x.cos()

torch/csrc/autograd/python_function.cpp glues this onto the C++ Node interface.

derivatives.yaml

The bulk of in-tree backward rules live in tools/autograd/derivatives.yaml. Each entry pairs an op with closed-form derivatives:

- name: add(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  self: handle_r_to_c(self.scalar_type(), grad)
  other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj()))

torchgen renders these into Functions.cpp and VariableType_*.cpp under torch/csrc/autograd/generated/. The latter contains the per-op autograd kernels that record onto the tape.

Forward-mode autodiff

torch/autograd/forward_ad.py and aten/src/ATen/native/ForwardADGrad.cpp implement dual-number based forward-mode autodiff. It is implemented by carrying a tangent on each tensor through AutogradMeta::fw_grad_. Forward-mode is younger than reverse-mode and supports a smaller op set; it is the primitive used by torch.func.jvp and torch.func.jacfwd.

Anomaly detection and gradcheck

• with torch.autograd.detect_anomaly(): records Python tracebacks where each Node was created and re-raises them with that traceback when backward errors. Implemented in torch/csrc/autograd/anomaly_mode.cpp.
• torch.autograd.gradcheck/gradgradcheck numerically verify gradients against finite differences and are used pervasively in tests.

Integration points

• Dispatcher. Autograd kernels are dispatched on the Autograd* keys; see Dispatcher.
• AOT Autograd. A different system that pre-traces autograd into a single forward+backward FX graph. See AOT Autograd.
• JIT/TorchScript. TorchScript integrates with the autograd engine for differentiable scripted functions.
• Profiler. Autograd Nodes and saved tensors show up in profiler traces.
• Distributed. DDP and FSDP hook into autograd via comm_hooks_ and register_post_accumulate_grad_hook.

Entry points for modification

• To add a derivative for an existing op, edit tools/autograd/derivatives.yaml. Re-run the build; torchgen regenerates torch/csrc/autograd/generated/.
• To customize the saved-tensor mechanics for an op, write a custom Backward class in torch/csrc/autograd/FunctionsManual.cpp and reference it from derivatives.yaml.
• To add user-extensible behaviour, work in torch/autograd/function.py (Python) and torch/csrc/autograd/python_function.cpp (binding).
• For engine-level work (parallelism, device synchronization), torch/csrc/autograd/engine.cpp is the file.

Key source files

File	Purpose
torch/csrc/autograd/engine.cpp	Multithreaded backward executor
torch/csrc/autograd/function.h	Node base class
torch/csrc/autograd/variable.cpp	AutogradMeta per-tensor state
torch/csrc/autograd/saved_variable.cpp	Saved tensors
torch/csrc/autograd/python_function.cpp	Function Python binding
tools/autograd/derivatives.yaml	Differentiation rules
tools/autograd/templates/Functions.cpp	Template for generated backward Nodes
tools/autograd/templates/VariableType.cpp	Template for generated autograd kernels
torch/autograd/__init__.py	Python user API
torch/autograd/function.py	Function extension API
torch/autograd/gradcheck.py	Numerical gradient check