the pytorch path · 0/12
start the path

the PyTorch path · Autograd · lesson 03 of 4

What the tape kept

Between the forward pass and the backward pass, some tensors stay alive because a node is holding them. Which ones, how many bytes, and what happens when you write into one are all readable from Python before backward ever runs.

the goal List what a given node saved and whether it saved an input or an output, measure the bytes a forward pass holds alive and explain why freezing a layer changes that number, and say exactly what the version counter compares and the one situation in which it compares nothing.

mastery work · this chapter0/3
  1. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
§ 01

Every node advertises what it saved

The derivative of sin is cos, which needs the input. The derivative of exp is exp, which is the output already computed. PyTorch encodes that choice per op and then exposes it: any _saved_ attribute on a node names a tensor that node is keeping alive, and dir() lists them.

So SinBackward0 carries _saved_self and ExpBackward0 carries _saved_result, and ReluBackward0 carries _saved_result too, because a relu's gradient only needs to know where the output was positive. MulBackward0 carries both operands. Nothing here is documentation you have to trust; it is the node telling you.

One detail in the last two lines of the run is the reason lesson one bothered with output slots. y.grad_fn._saved_result is not y. It is a different Python tensor over the same storage, rebuilt on each access, because a node holding a strong reference to its own output would be a reference cycle that never collects. The rebuild is why the saved value can be checked against the live one at all.

run it (verified, torch 2.2.2 CPU): four nodes naming what they kept
import torch

x = torch.ones(3, requires_grad=True)
saved = lambda t: [a for a in dir(t.grad_fn) if a.startswith("_saved")]
print(x.sin().grad_fn.name(), saved(x.sin()))       # SinBackward0 ['_saved_self']
print(x.exp().grad_fn.name(), saved(x.exp()))       # ExpBackward0 ['_saved_result']
print(torch.relu(x).grad_fn.name(), saved(torch.relu(x)))
# ReluBackward0 ['_saved_result']
print((x * x).grad_fn.name(), saved(x * x))
# MulBackward0 ['_saved_other', '_saved_self']

y = x.exp()
s = y.grad_fn._saved_result
print(s is y, s.data_ptr() == y.data_ptr())   # False True
§ 02

What gets saved depends on what you asked for

saved_tensors_hooks installs a pack function that fires once per save, so a forward pass under it prints its own memory bill. Run a three-layer MLP with a 256 by 512 input and the bill comes to five distinct storages and 2,623,488 bytes, itemized below.

The interesting entry is the one that is absent. 0.weight is 512 by 512 and never gets saved, while 2.weight of exactly the same shape does. The first layer's matmul needs to produce a gradient for its weight, which needs the input; it does not need to produce a gradient for its input, because the input is data. The second layer needs both, so it keeps its weight too.

Turn that around and it is a memory knob with a measured size. Set requires_grad on the input and 0.weight joins the list, taking the total to 3,672,064 bytes. Freeze the first layer instead and the total drops to 1,050,624, because two of the five storages stop being needed at once. Same model, same batch, a factor of 3.5 between the two ends, decided entirely by which gradients the program asked for.

run it (verified, torch 2.2.2 CPU): the forward pass printing what it holds alive
import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(),
                      nn.Linear(512, 512), nn.ReLU(),
                      nn.Linear(512, 1))
x = torch.randn(256, 512)
names = {p.data_ptr(): n for n, p in model.named_parameters()}

held = {}
def pack(t):
    held.setdefault(t.data_ptr(), (tuple(t.shape), t.numel() * t.element_size()))
    return t

with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
    model(x).sum()

for ptr, (shape, nbytes) in held.items():
    print(shape, nbytes, names.get(ptr, "x" if ptr == x.data_ptr() else "an activation"))
print("total", sum(b for _, b in held.values()))
what the program asks fordistinct storages heldbytes
as written: data in, all parameters training52,623,488
input carries requires_grad too63,672,064
first Linear frozen, data in31,050,624
saved storages for one forward pass of the same 3-layer MLP on a 256x512 input, three ways (verified, torch 2.2.2 CPU)
§ 03

The counter that notices a rewrite

Every tensor carries a version counter, readable as _version, that starts at zero and increments on each in-place write. Views share the counter with their base, so writing through a slice bumps the number the base reports, which is how a mutation reaches a node that never saw the view.

In-place writes are not banned, and the run below is the case that works. y = x * 2 then y.add_(1) bumps y._version to 1 and also rebases y.grad_fn from MulBackward0 to AddBackward0, and the backward pass afterwards is correct, because MulBackward0 saved x rather than y and nothing it kept was touched.

Swap exp for the multiply and the same two lines fail, because exp saved its own output and that output is what got rewritten. The museum wing holds that failure with its verbatim error. The mechanism worth carrying is the one this run makes visible: the counter is always running, and whether a write matters depends only on whether some node saved that exact tensor first.

run it (verified, torch 2.2.2 CPU): an in-place write that rebases the node and keeps the gradient right
import torch

x = torch.ones(3, requires_grad=True)
y = x * 2
print(y.grad_fn.name(), y._version)   # MulBackward0 0
y.add_(1)
print(y.grad_fn.name(), y._version)   # AddBackward0 1
y.sum().backward()
print(x.grad)                         # tensor([2., 2., 2.])

b = torch.ones(4)
v = b[:2]
v.add_(1)
print(b._version, v._version)         # 1 1
§ 04

What the check actually compares

Two lines of C++ from opposite ends of SavedVariable are the whole comparison. On the way in, the constructor copies the tensor's current version into saved_version_. On the way out, unpack reads the version again and raises if the two differ. Nothing is checksummed and nothing is copied; a saved tensor is a pointer plus an integer taken at save time.

That is why the failure is reported at the backward pass rather than at the write. Nobody is watching the tensor. The counter increments silently, and the mismatch is discovered when the node finally asks for what it saved, which can be many lines later in a place with no obvious connection to the mutation.

It also means you can run the check yourself, early. _version on a tensor you know a node saved tells you before you ever call backward whether the number has moved.

verbatim, torch/csrc/autograd/saved_variable.cpp at pytorch v2.2.2: line 50 from the constructor, then lines 160-167 from unpack, joined here under added headings
// on the way in, at SavedVariable::SavedVariable
    saved_version_ = version_counter.current_version();

// on the way out, at SavedVariable::unpack
  // Only check version counter in the case without hooks
  // If user provides hooks, we can't track versions through the hooks
  if (!hooks_) {
    auto current_version = saved_original_
        ? impl::version_counter(data_).current_version()
        : version_counter_.current_version();

    if (saved_version_ != current_version) {
§ 05

And the hook that turns the check off

The if (!hooks_) in that excerpt is a guard with a consequence most people meet by accident. Install saved_tensors_hooks and the version check is skipped, because a pack function may have returned something that is not a tensor at all and there is nothing left to read a version from.

So the exact program that raises without hooks computes a wrong answer with them. Under the hooks, e.grad comes back 3.7183, which is the mutated saved value; the gradient of sum(exp(e)) at e = 1 is 2.7183. No warning, no error, a number that is off by exactly the size of the in-place add.

The same hook pair is the mechanism behind offloading saved tensors to host memory, which is the reason it exists and a good thing to do. What it costs is the guardrail, and any pack and unpack pair you write should be treated as code that has to be right on its own, because nothing downstream is going to check it.

run it (verified, torch 2.2.2 CPU): hooks installed, check skipped, answer silently wrong
import torch

e = torch.ones(3, requires_grad=True)
with torch.autograd.graph.saved_tensors_hooks(lambda t: t, lambda t: t):
    z = e.exp()
    z.add_(1)           # the write the museum's version-counter exhibit catches
    z.sum().backward()
print(e.grad)           # tensor([3.7183, 3.7183, 3.7183]); the answer is 2.7183
before you move on

Check yourself

01 Two Linear layers in the same model have identically shaped weights. Why is only one of them saved for backward?

Because a matmul saves its weight only when a gradient for its input is needed, and the first layer feeds on data that carries no gradient. Measured on a 3-layer MLP, 0.weight was absent and 2.weight present, a difference of 1,048,576 bytes.

02 What exactly does the version counter compare, and when is the comparison made?

The integer recorded when the tensor was saved against the integer the tensor reports when the node unpacks it, which happens during the backward pass. Nothing is checksummed, so a mutation is discovered at unpack rather than at the write that caused it.

03 You installed saved_tensors_hooks and a program that used to raise now returns a number. Should you trust it?

No. The version check runs only when no hooks are installed, so the same in-place write that raised before now passes silently. The measured case returned 3.7183 where the correct gradient is 2.7183.

assigned

Readings