How PyTorch Autograd Dynamic DAG Engine Works

PyTorch relies on its autograd engine to calculate gradients automatically using a dynamic Directed Acyclic Graph (DAG). Unlike static computational graph frameworks that define the graph structure before execution, PyTorch generates a new graph on-the-fly during every forward pass. This article explains how the autograd engine tracks tensor operations, builds graph nodes and edges via dynamic execution, traverses these structures during backpropagation using reverse-mode automatic differentiation, and cleans up memory dynamically between training steps.

Tensors and the grad_fn Reference

The foundation of the autograd graph is the torch.Tensor. When a tensor is initialized with requires_grad=True, PyTorch begins monitoring all mathematical operations applied to it.

Tensors created directly by the user are designated as "leaf nodes." Leaf nodes do not have an operation that created them, so their grad_fn attribute remains None. However, whenever an operation transforms a tensor that tracks gradients, the resulting output tensor receives a grad_fn attribute. This attribute points to a C++ Node object (such as AddBackward0 or MulBackward0) that encapsulates the backward formula for that specific operation.

Dynamic Graph Construction in the Forward Pass

PyTorch employs a "define-by-run" execution model. The computational graph is not compiled ahead of time; instead, it is constructed eagerly as Python executes operations sequentially:

1. Operation Execution: When an operation like z = x * y executes, PyTorch's C++ backend performs the forward numerical computation.
2. Node Creation: Concurrently, an operation-specific Node instance is instantiated in memory.
3. Tracking Dependencies: The output tensor z is assigned this new node as its grad_fn. The new node records pointers to the grad_fn instances of its inputs (x and y) through an internal list called next_edges.
4. Context Storage: Operations that require forward-pass values to compute derivatives (such as activations or activations' inputs) cache these values using an internal context object (ctx.save_for_backward).

Through this mechanism, the graph represents operations as nodes and data dependencies as directed edges flowing from outputs back toward the original inputs.

Representation of the Graph

In the internal C++ core (torch::autograd), the graph is structured with clear roles:

• Nodes (torch::autograd::Node): Represent elementary operations. Each node defines an apply() method that accepts incoming gradients (vector-Jacobian products) and computes gradients with respect to its inputs.
• Edges (torch::autograd::Edge): Directed connections consisting of a target Node pointer and an input index, specifying which input of the parent operation the gradient corresponds to.
• AccumulateGrad Nodes: Special terminal nodes attached to leaf tensors. When backward execution reaches a leaf, the AccumulateGrad node adds the computed gradient directly into the tensor's .grad attribute.

The Backward Pass and Graph Traversal

Backpropagation begins when .backward() is called on a scalar tensor (typically the loss). The autograd engine evaluates the graph through reverse-mode automatic differentiation:

1. Root Initialization: The engine initializes the gradient of the root tensor, defaulting to torch.tensor(1.0) for scalar outputs.
2. Topological Sorting and Queuing: The C++ Engine::execute method manages graph traversal using a task queue. It tracks node dependencies to ensure a node is only evaluated once all its dependent downstream gradients have been calculated.
3. Applying Derivatives: The engine pops nodes from the queue, passes incoming gradients to their apply() method, and receives the resulting gradients for input variables.
4. Accumulation: As gradients reach AccumulateGrad nodes, they are written to the .grad fields of the respective leaf tensors using an in-place addition (+=) operation.

Dynamic Behavior and Memory Management

A critical design feature of PyTorch's DAG is its ephemerality. By default, once .backward() finishes traversing the graph, the intermediate buffers and non-leaf node references are immediately freed from memory (retain_graph=False).

Because the graph is dismantled after every backward pass, control flow statements such as Python if conditions, for loops, and dynamic recursion naturally alter the graph topology across iterations. In the subsequent training step, the forward pass dynamically constructs an entirely new DAG reflecting whatever execution path the Python interpreter takes.