Repository: pytorch
Issue: RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward
Labels: module: autograd,triaged,module: nn.utils.parametrize

### 🐛 Describe the bug

I am reopening the issue https://github.com/pytorch/pytorch/issues/71062 with the following additional informations:
When training a network which contains parametrization in DP (Distributed Parallel) and the caching mode (torch.nn.utils.parametrize.cached), a race condition occurs at a random point in the training causing the error:


RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

Anomaly detection points to the following culprit:
```
File "/mnt/task_runtime/src/dibnn_ext/nn/modules.py", line 832, in forward
    weight=self.weight.to(input.device, copy=True) * self.gain, bias=self.bias,
  File "/miniconda/lib/python3.7/site-packages/torch/nn/utils/parametrize.py", line 339, in get_parametrized
    return get_cached_parametrization(parametrization)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/utils/parametrize.py", line 324, in get_cached_parametrization
    tensor = parametrization()
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/utils/parametrize.py", line 260, in forward
    x = self[0](self.original)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/task_runtime/src/dibnn_ext/nn/parametrization.py", line 36, in forward
    weight = DF.weight_standardization(weight, eps=self.eps)
  File "/mnt/task_runtime/src/dibnn_ext/nn/functional.py", line 660, in weight_standardization
    weight = (weight - mean) / (torch.sqrt(var * fan_in + eps))
 (function _print_stack)
```

The training step is defined as follow:
```
def training_step(...):
   ...
   with torch.nn.utils.parametrize.cached():
            pred = self.forward(data.detach(), data2.detach())
            loss = torch.abs(pred - data).mean()

    return { 'loss': loss }
```

I suspect the following:
The training step is handle per thread and per GPU, but the caching mechanism is global. What can happen is GPU 0 starts to enter to the forward pass and cache the weights, GPU 1 enters later in the game but GPU 0 has finished its backward pass and the gradients of the weights cached are not longer available, this causes issue on GPU 1 since its backward pass now is invalid.

This also explain why the following line is required to transfer weights cached to the current valid device/GPU when caching is used:
```
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return TF.conv2d(input,
                         weight=self.weight.to(input, copy=True) * self.gain, bias=self.bias,  # Due to the cache, weights can be on a different device than input
                         stride=self._stride,
                         padding=self._padding,
                         dilation=self._dilation,
                         groups=self._groups)
```

Attached a dummy example to reproduce the issue using DP and autograd-detect-anomaly.


Full example to reproduce the issue
```python
from typing import Optional, List, Dict

import os
import argparse
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as TF
import torch.nn.utils.parametrize as TP

import torchvision
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

import pytorch_lightning as pl


def weight_standardization(weight: torch.Tensor, eps: float = 1e-6):
    """
    Weight standardization as defined by:
    Characterizing signal propagation to close the performance gap in unnormalized ResNets, https://arxiv.org/abs/2101.08692

    :param weight: weight to standardize
    :param eps: epsilon to ensure safe division
    :return: weight to standardized
    """

    fan_in = nn.init._calculate_fan_in_and_fan_out(weight)[0]

    var, mean = torch.var_mean(weight, dim=list(range(1, weight.ndim)), keepdim=True)
    weight = (weight - mean) / (torch.sqrt(var * fan_in + eps))
    return weight


class WeightStandardization(nn.Module):
    def __init__(self, eps: float=1e-6):
        super().__init__()

        self.eps = eps

    def forward(self, weight: torch.Tensor) -> torch.Tensor:
        weight = weight_standardization(weight, eps=self.eps)
        return weight


class SWSConv2d(nn.Module):
    r"""
    2D Conv layer with Scaled Weight Standardization.

    Characterizing signal propagation to close the performance gap in unnormalized ResNets
    https://arxiv.org/abs/2101.08692
    """
    def __init__(self,
                 in_channels: int, out_channels: int, kernel_size: int,
                 stride: int=1, padding: int=0, padding_mode: str='zeros', dilation=1, groups: int=1,
                 bias: bool=True):

        super().__init__()

        self._stride = stride
        self._padding = padding
        self._padding_mode = padding_mode
        self._dilation = dilation
        self._groups = groups

        self.weight = nn.Parameter(torch.empty(out_channels, in_channels, kernel_size, kernel_size))
        self.bias = nn.Parameter(torch.empty(out_channels)) if bias else None
        self.gain = nn.Parameter(torch.ones(out_channels, 1, 1, 1))

        # Init weights
        k = groups / (in_channels * kernel_size * kernel_size)
        torch.nn.init.uniform_(self.weight, -k, k)
        if bias:
            torch.nn.init.uniform_(self.bias, -k, k)

        self.register_parametrization()

    def register_parametrization(self):
        if not TP.is_parametrized(self, 'weight'):
            TP.register_parametrization(self, 'weight', WeightStandardization(eps=1e-6))

    def remove_parametrization(self, leave_parametrized: bool=True):
        if TP.is_parametrized(self, 'weight'):
            TP.remove_parametrizations(self, 'weight', leave_parametrized=leave_parametrized)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return TF.conv2d(input,
                         weight=self.weight.to(input, copy=True) * self.gain, bias=self.bias,  # Due to the cache, weights can be on a different device than input
                         stride=self._stride,
                         padding=self._padding,
                         dilation=self._dilation,
                         groups=self._groups)


class FeatHeadLayer(nn.Module):
    def __init__(self,
                 in_channels: int,
                 out_channels: int = 16,
                 hidden_features: int = 24):

        super().__init__()

        activation = nn.LeakyReLU(negative_slope=0.1)

        self._in_channels = in_channels

        # TODO: Improve hidden channels as multiple of 16
        self._feat_head = nn.Sequential(
            SWSConv2d(in_channels, hidden_features, 3, stride=1, bias=True, padding=1), activation,
            SWSConv2d(hidden_features, out_channels, 1, bias=True)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._feat_head(x)


class RegressionLayer(nn.Module):
    def __init__(self):
        super().__init__()

        self._regres = SWSConv2d(32, 1, 1)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        pred = self._regres(input)
        return pred


class DummyModel(nn.Module):
    def __init__(self):
        super().__init__()

        self._features_head = FeatHeadLayer(in_channels=1, out_channels=24)
        self._cv = SWSConv2d(24 * 2, 32, 1)
        self._regression_layer = RegressionLayer()

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        feat0_mc2 = self._features_head(x1)
        feat1_mc2 = self._features_head(x2)
        corr2 = self._cv(torch.cat((feat0_mc2, feat1_mc2), dim=1))
        pred = self._regression_layer(corr2)
        return pred


class MNISTDataModule(pl.core.LightningDataModule):
    def __init__(self, hparams):
        super().__init__()

        self._hparams = hparams

        self._transform = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(
                (0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        # Download the datasets
        MNIST(os.getcwd(), train=True, download=True, transform=None)
        MNIST(os.getcwd(), train=False, download=True, transform=None)

    def train_dataloader(self):
        import torch.utils.data

        num_samples = self.hparams.num_steps * self.hparams.batch_size

        dataset = MNIST(os.getcwd(), train=True, download=False, transform=self._transform)  # Only access to the dataset previously downloaded
        sampler = torch.utils.data.RandomSampler(dataset, replacement=True, num_samples=num_samples)
        loader = DataLoader(dataset, sampler=sampler, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
        return loader

    def val_dataloader(self):
        dataset = MNIST(os.getcwd(), train=False, download=False, transform=self._transform)
        loader = DataLoader(dataset, batch_size=1, num_workers=0)
        return loader

    @staticmethod
    def add_specific_args(parser: argparse.ArgumentParser):
        g_training = parser.add_argument_group('training')
        g_training.add_argument('--num-steps', dest='num_steps', type=int, default=1000)
        g_training.add_argument('--batch-size', dest='batch_size', type=int, default=64)
        g_training.add_argument('--num-workers', dest='num_workers', type=int, default=16)


class MNISTModel(pl.core.LightningModule):
    def __init__(self, hparams):
        super().__init__()

        self.save_hyperparameters(hparams)

        self.model = DummyModel()

    def forward(self, x1, x2):
        pred = self.model(x1, x2)
        return pred

    def training_step(self, batch, batch_idx):
        data = batch[0]
        target = batch[1]

        data2 = torch.randn(data.size())

        with torch.nn.utils.parametrize.cached():
            pred = self.forward(data.detach(), data2.detach())
            loss = torch.abs(pred - data).mean()

        return { 'loss': loss }

    def validation_step(self, batch, batch_idx, *args):
        data = batch[0]
        target = batch[1]

        data2 = torch.randn(data.size())

        pred = self.forward(data, data2)
        loss = torch.abs(pred - data).mean()
        return {
            'loss': loss.item(),
            'batch': batch,
            'pred': pred
        }

    def test_step(self, batch, batch_idx, *args):
        return self.validation_step(batch, batch_idx, *args)

    def validation_visualization_groundtruth(self, batch, batch_idx, dataloader_idx: int):
        viz_dict = dict()

        data = batch[0]
        data_img = data.detach().cpu().squeeze(0).numpy().transpose((1, 2, 0))
        viz_dict['data'] = np.clip(data_img * 255.0, 0.0, 255.0).astype(np.uint8)
        return viz_dict

    def configure_optimizers(self):
        from torch.optim.sgd import SGD

        optim = SGD(self.parameters(), lr=self.hparams.learning_rate, momentum=self.hparams.momentum)
        return {
            'optimizer': optim
        }

    @staticmethod
    def add_specific_args(parser: argparse.ArgumentParser):
        default_log_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '71062', 'log')
        parser.add_argument('--log-path', dest='log_path', type=str, default=default_log_path)
        parser.add_argument('--num-epochs', dest='num_epochs', type=int, default=100)

        g_training = parser.add_argument_group('training')
        g_training.add_argument('--learning-rate', dest='learning_rate', type=float, default=0.01)
        g_training.add_argument('--momentum', dest='momentum', type=float, default=0.5)


if __name__ == '__main__':
    args_parser = argparse.ArgumentParser(description='MNIST example')
    MNISTModel.add_specific_args(args_parser)
    MNISTDataModule.add_specific_args(args_parser)

    hparams = args_parser.parse_args()
    os.makedirs(hparams.log_path, exist_ok=True)

    import pytorch_lightning as pl
    import pytorch_lightning.loggers
    import pytorch_lightning.plugins

    model = MNISTModel(hparams)
    data_module = MNISTDataModule(hparams)

    pl.seed_everything(42)

    num_epochs = hparams.num_epochs
    trainer = pl.trainer.Trainer(default_root_dir=hparams.log_path,
                                 enable_checkpointing=True,
                                 enable_model_summary=False,
                                 strategy=pl.plugins.DataParallelPlugin(parallel_devices=None),
                                 accelerator='gpu',
                                 devices='auto',
                                 min_epochs=num_epochs, max_epochs=num_epochs,
                                 num_sanity_val_steps=0,
                                 check_val_every_n_epoch=1,
                                 val_check_interval=1.0,
                                 limit_val_batches=100,
                                 limit_test_batches=50,
                                 detect_anomaly=True)

    # Run training
    trainer.fit(model, datamodule=data_module)

    # Test loop
    trainer.test(model, dataloaders=data_module.val_dataloader())
```

### Versions

Versions

Collecting environment information...
PyTorch version: 1.10.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: Could not collect
Libc version: glibc-2.10

Python version: 3.7.7 (default, Mar 23 2020, 22:36:06) [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.19.56-1.el7.x86_64-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 418.40.04
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] pytorch-lightning==1.5.5
[pip3] torch==1.10.0+cu113
[pip3] torchaudio==0.10.0+cu113
[pip3] torchmetrics==0.6.2
[pip3] torchvision==0.11.1+cu113
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.3.1 h2bc3f7f_2
[conda] libblas 3.9.0 12_linux64_mkl conda-forge
[conda] libcblas 3.9.0 12_linux64_mkl conda-forge
[conda] liblapack 3.9.0 12_linux64_mkl conda-forge
[conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py37h7f8727e_0
[conda] mkl_fft 1.3.1 py37hd3c417c_0
[conda] mkl_random 1.2.2 py37h51133e4_0
[conda] numpy 1.19.5 pypi_0 pypi
[conda] numpy-base 1.21.2 py37h79a1101_0
[conda] pytorch-cpu 1.1.0 py37he1b5a44_0 conda-forge
[conda] pytorch-lightning 1.5.5 pyhd8ed1ab_0 conda-forge
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torch 1.10.0+cu113 pypi_0 pypi
[conda] torchaudio 0.10.0+cu113 pypi_0 pypi
[conda] torchmetrics 0.6.2 pyhd8ed1ab_0 conda-forge
[conda] torchvision 0.11.1+cu113 pypi_0 pypi

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7