That’s not necessarily the case since the forward activations (which might be needed for the gradient calculation) could increase the memory usage by a significantly larger factor (depending on the model architecture). E.g. take a look at this post describing the memory usage for a ResNet.
This shouldn’t be the case as you should be able to simply wrap your forward pass into the context manager. Here is a small example showing the hooks (which are also used for CPU offloading) for a ResNet:
import torch
import torchvision.models as models
def pack_hook(x):
print("Packing", x.sum())
return x
def unpack_hook(x):
print("Unpacking", x.sum())
return x
device = "cuda"
model = models.resnet50()
model.to(device)
x = torch.randn(1, 3, 224, 224, device=device)
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
out = model(x)
out.sum().backward()
Besides that, you could also try to use activation checkpointing as described here.