Repository: pytorch
Issue: Multi-GPU training stuck when using grad scaler
Labels: oncall: distributed,module: cuda,module: amp (automated mixed precision),Stale

## 🐛 Bug

When trying to run multi-GPU training using amp, I notice that training sometimes freezes. The GPUs have one of the two following stack traces:

Current thread 0x00007f4aed99f740 (most recent call first):
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py", line 284 in <genexpr>
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py", line 284 in _maybe_opt_step
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py", line 338 in step

Current thread 0x00007f0a91baa740 (most recent call first):
  File "/data/hturki/my_code/models/my_model.py", line 28 in forward
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102 in _call_impl
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 886 in forward
  File "/home/hturki/miniconda3/envs/my-env/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102 in _call_impl

## To Reproduce

I'm having a hard time pinning a non-data-dependent repro, but my training code looks something like:

```
            for dataset_index, item in enumerate(data_loader):
                with torch.cuda.amp.autocast(enabled=self.hparams.amp):
                    metrics, bg_rays_present = self._training_step(
                        item['rgbs'].to(self.device, non_blocking=True),
                        item['rays'].to(self.device, non_blocking=True),
                        item['image_indices'].to(self.device, non_blocking=True))

                for optimizer in optimizers.values():
                    optimizer.zero_grad(set_to_none=True)

                scaler.scale(metrics['loss']).backward()

                for key, optimizer in optimizers.items():
                    scaler.step(optimizer)

                scaler.update()

                for scheduler in schedulers.values():
                    scheduler.step()
```

And my model code is:

```
class MyModel(nn.Module):
    def __init__(self, sub_modules: List[nn.Module], centroids: torch.Tensor, boundary_margin: int, xyz_real: bool):
        super(MegaNeRF, self).__init__()
        self.sub_modules = nn.ModuleList(sub_modules)
        self.register_buffer('centroids', centroids)
        self.boundary_margin = boundary_margin
        self.xyz_real = xyz_real

    def forward(self, x: torch.Tensor, sigma_only: bool = False,
                sigma_noise: Optional[torch.Tensor] = None) -> torch.Tensor:
        cluster_distances = torch.cdist(x[:, 1:3], self.centroids[:, 1:3])
        inverse_cluster_distances = 1 / (cluster_distances + 1e-5)

        min_cluster_distances = cluster_distances.min(dim=1)[0].unsqueeze(-1).repeat(1, cluster_distances.shape[1])
        inverse_cluster_distances[cluster_distances > self.boundary_margin * min_cluster_distances] = 0
        weights = inverse_cluster_distances / (inverse_cluster_distances.sum(dim=-1).unsqueeze(-1) + 1e-8)

        results = torch.empty(0)

        for i, child in enumerate(self.sub_modules):
            cluster_mask = weights[:, i] > 0
            sub_input = x[cluster_mask, 3:] if self.xyz_real else x[cluster_mask] # Corresponds to "/data/hturki/my-code/models/my_model.py", line 28 in forward in the second stuck stacktrace

            if sub_input.shape[0] > 0:
                sub_result = child(sub_input, sigma_only,
                                   sigma_noise[cluster_mask] if sigma_noise is not None else None)

                if results.shape[0] == 0:
                    results = torch.zeros(x.shape[0], sub_result.shape[1], device=sub_result.device,
                                          dtype=sub_result.dtype)

                results[cluster_mask] += sub_result * weights[cluster_mask, i].unsqueeze(-1)

        return results
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

I feel like this is likely due to user error, but it would be great to have a more informative description of what's incorrect vs things silently hanging indefinitely.

## Environment

PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 8.4.0
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.17

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.1.3.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti
GPU 8: GeForce RTX 2080 Ti
GPU 9: GeForce RTX 2080 Ti

Nvidia driver version: 455.38
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] torch==1.10.0
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py39h7f8727e_0
[conda] mkl_fft                   1.3.1            py39hd3c417c_0
[conda] mkl_random                1.2.2            py39h51133e4_0
[conda] numpy                     1.21.2           py39h20f2e39_0
[conda] numpy-base                1.21.2           py39h79a1101_0
[conda] pytorch                   1.10.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchvision               0.11.1               py39_cu113    pytorch



cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @ngimel @mcarilli @ptrblck