Repository: pytorch
Issue: Performance: Autodiff stores inputs of convolution layers with non-trainable weight
Labels: module: autograd,module: memory usage,triaged,enhancement,needs design

### 🐛 Describe the bug

**Description:** Assume we have a concolution layer that processes an input `x` into an output `y` using a kernel `w`, i.e. `y = conv(x, w)`. Further, assume that the input is differentiable, but the weight is not, i.e. `w.requires_grad=False` and `x.requires_grad=True`. 
- **Desired behavior:** To backpropagate through this layer, we only need to store `w`, but not `x`.
- **Observed behavior:** It seems that PyTorch stores `x`, too, even when `w.requires_grad=False`.

**Summary:** Since `x` is unnecessarily kept in memory, this increases memory usage.

The same reasoning applies to other layers whose output depends linearly on both the input and the parameters, such as fully-connected layers, transpose convolutions, and batch normalization in evaluation mode.

We found the following layers to be affected by the above behavior:

- `torch.nn.Conv{1,2,3}d`
- `torch.nn.ConvTranspose{1,2,3}d`
- `torch.nn.BatchNormNd` (evaluation mode)

**Solution:** Store only the necessary tensors depending on the `requires_grad` values of layer inputs and weights. See e.g. [here](https://github.com/plutonium-239/memsave_torch/blob/fef6003dc84111de5fa8993f39e6a94194e6382d/memsave_torch/nn/functional/Conv.py#L40-L43) for an example fix that simply defines a new `torch.autograd.Function` for convolution with the correct logic.

**More details:** We (@f-dangel and me) describe this finding in more detail in our [WANT@ICML'24 paper](https://openreview.net/pdf?id=KsUUzxUK7N) along with visualizations demonstrating this behaviour.

**MWE:** Let's consider a sequence of 5 size-preserving convolutions (each intermediate takes 512 MiB storage). If we set all parameters except for the first layer's weight to non-trainable, PyTorch should discard the non-trainable layers' inputs. However, the memory consumption seems to be the same as if all weights were trainable. See Figure 1 in the [paper](https://openreview.net/pdf?id=KsUUzxUK7N).

Example requirements: `pip install memory_profiler==0.61.0 torch==2.0.1`

```python
"""PyTorch's convolutions store their input when `weight.requires_grad=False`."""

from collections import OrderedDict

from memory_profiler import memory_usage
from torch import rand
from torch.nn import Conv2d, Sequential

SHAPE_X = (256, 8, 256, 256)  # shape of the input
MEM_X = 512  # requires 512 MiB storage
NUM_LAYERS = 5


def setup():
    """Create a deep linear CNN with size-preserving convolutions and an input."""
    layers = OrderedDict()
    for i in range(NUM_LAYERS):
        layers[f"conv{i}"] = Conv2d(8, 8, 3, padding=1, bias=False)
    return rand(*SHAPE_X), Sequential(layers)


# Consider three different scenarios: 1) no parameters are trainable, 2) all
# layers are trainable, 3) only the first layer is trainable
def non_trainable():
    """Forward pass through the CNN with all layers non-trainable."""
    X, net = setup()
    for i in range(NUM_LAYERS):
        getattr(net, f"conv{i}").weight.requires_grad = False

    for name, param in net.named_parameters():
        print(f"{name}, requires_grad={param.requires_grad}")

    return net(X)


def all_trainable():
    """Forward pass through the CNN with all layers trainable."""
    X, net = setup()
    for i in range(NUM_LAYERS):
        getattr(net, f"conv{i}").weight.requires_grad = True

    for name, param in net.named_parameters():
        print(f"{name}, requires_grad={param.requires_grad}")

    return net(X)


def first_trainable():
    """Forward pass through the CNN with first layer trainable."""
    X, net = setup()
    for i in range(NUM_LAYERS):
        getattr(net, f"conv{i}").weight.requires_grad = i == 1

    for name, param in net.named_parameters():
        print(f"{name}, requires_grad={param.requires_grad}")

    return net(X)


if __name__ == "__main__":
    kwargs = {"interval": 1e-4, "max_usage": True}  # memory profiler settingss

    # measure memory and print
    mem_setup = memory_usage(setup, **kwargs)
    print(f"Weights+input: {mem_setup:.1f} MiB.")

    mem_non = memory_usage(non_trainable, **kwargs) - mem_setup
    print(
        f"Non-trainable: {mem_non:.1f} MiB (≈{mem_non / MEM_X:.1f} hidden activations)."
    )

    mem_all = memory_usage(all_trainable, **kwargs) - mem_setup
    print(
        f"All-trainable: {mem_all:.1f} MiB (≈{mem_all / MEM_X:.1f} hidden activations)."
    )

    mem_first = memory_usage(first_trainable, **kwargs) - mem_setup
    print(
        f"First-trainable: {mem_first:.1f} MiB (≈{mem_first / MEM_X:.1f} hidden activations)."
    )
```

Output:
```bash
Weights+input: 702.5 MiB.
conv0.weight, requires_grad=False
conv1.weight, requires_grad=False
conv2.weight, requires_grad=False
conv3.weight, requires_grad=False
conv4.weight, requires_grad=False
Non-trainable: 1535.5 MiB # (≈3.0×hidden activations).
conv0.weight, requires_grad=True
conv1.weight, requires_grad=True
conv2.weight, requires_grad=True
conv3.weight, requires_grad=True
conv4.weight, requires_grad=True
All-trainable: 3071.6 MiB # (≈6.0×hidden activations).
conv0.weight, requires_grad=False
conv1.weight, requires_grad=True
conv2.weight, requires_grad=False
conv3.weight, requires_grad=False
conv4.weight, requires_grad=False
First-trainable: 3071.6 MiB # (≈6.0×hidden activations).
```




### Versions

```
Collecting environment information...
PyTorch version: 2.0.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home Single Language
GCC version: (tdm64-1) 10.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.6 (tags/v3.8.6:db45529, Sep 23 2020, 15:52:53) [MSC v.1927 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19041-SP0
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1660 Ti
Nvidia driver version: 546.17
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2592
DeviceID=CPU0
Family=198
L2CacheSize=1536
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2592
Name=Intel(R) Core(TM) i7-10750H CPU @ 2.60GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] flake8==7.0.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.4
[pip3] pytorch-lightning==1.2.8
[pip3] pytorch-memlab==0.3.0
[pip3] pytorch-metric-learning==1.2.1
[pip3] torch==2.0.1+cu118
[pip3] torch-cluster==1.6.1+pt20cu118
[pip3] torch-geometric==2.3.1
[pip3] torch-scatter==2.1.1+pt20cu118
[pip3] torch-sparse==0.6.17+pt20cu118
[pip3] torch-spline-conv==1.2.2+pt20cu118
[pip3] torchaudio==2.0.2+cu118
[pip3] torcheval==0.0.7
[pip3] torchinfo==1.8.0
[pip3] torchmetrics==0.8.2
[pip3] torchsample==0.1.3
[pip3] torchview==0.2.6
[pip3] torchvision==0.15.2+cu118
[pip3] torchviz==0.0.2
[conda] Could not collect
```

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan