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.
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 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.
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 for | distinct storages held | bytes |
|---|---|---|
| as written: data in, all parameters training | 5 | 2,623,488 |
| input carries requires_grad too | 6 | 3,672,064 |
| first Linear frozen, data in | 3 | 1,050,624 |
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.
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 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.
// 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) { 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.
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 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.
Readings
- Hooks for autograd saved tensors ↗ the official tutorial for pack and unpack, including the offload-to-CPU pattern
- saved_variable.cpp at v2.2.2 ↗ the save, the unpack, the version comparison, and the message it formats, in 250 lines
- Autograd mechanics: in-place operations ↗ the maintainers' own case against in-place ops, and the two reasons they still support them