Repository: pytorch
Issue: Optimizer step stuck during distributed training
Labels: oncall: distributed,module: data parallel,module: amp (automated mixed precision)

## 🐛 Bug

Model training gets stuck when using mixed precision training on single node multiple GPUs 

## To Reproduce

Steps to reproduce the behavior:

```python
#!/usr/bin/env python
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import argparse
import numpy as np
import wandb
import pickle
import clip
import torch.backends.cudnn as cudnn

from sentence_transformers import SentenceTransformer
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm as tqdm
from sentence_transformers import SentenceTransformer, models
from transformers import AutoTokenizer, AutoModel
from dataloader import  get_iterable_dataloaders

class Model(nn.Module):
    def __init__(self):
        super(type(self), self).__init__()
        self.lang_model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
        
    def forward(self, other_lang):
        lang_emb = self.lang_model(other_lang)['sentence_embedding']
        return  lang_emb


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


def init_distributed_mode(args):
    # launched with torch.distributed.launch
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    dist.init_process_group(
        backend="nccl",
        init_method=args.dist_url,
        world_size=args.world_size,
        rank=args.rank,
    )

    torch.cuda.set_device(args.gpu)
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    dist.barrier()
    setup_for_distributed(args.rank == 0)


def run(args):
    init_distributed_mode(args)
    cudnn.benchmark = True
    rank = args.gpu
    num_epochs = args.num_epochs
    torch.cuda.set_device(rank)
    
    clip_model, _ = clip.load("ViT-B/32", device='cuda')
    for param in clip_model.parameters():
        param.requires_grad = False
    clip_model.eval()
    model = Model()
    model = model.cuda()
    model = DDP(model, device_ids=[rank], find_unused_parameters=True)
    torch.autograd.set_detect_anomaly(True)
    
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, betas=(0.9,0.98), eps=1e-6, weight_decay=0.2)
    train_loader, val_loader = get_iterable_dataloaders(args)
    best_val_loss = np.inf
    fp16_scaler = torch.cuda.amp.GradScaler()
    for i in range(0, args.num_epochs):
        total_loss = 0
        step = 0
        logs = {}
        model.train()
        for english_title, other_lang_title in tqdm(train_loader):
            other_lang_title = {k: v.cuda() for k, v in other_lang_title.items()}
            english_title= english_title.cuda()
            with torch.no_grad():
                english_emb = clip_model.encode_text(english_title)
                english_emb = english_emb / english_emb.norm(dim=-1, keepdim=True)
            english_embedding = english_emb.detach().clone()
            with torch.cuda.amp.autocast():
                other_lang_emb = model(other_lang_title) 
                output = loss_fn(other_lang_emb, english_embedding)
            model.zero_grad()
            print("Completed zero grad")
            fp16_scaler.scale(output).backward()
            print("Completed backward pass")
            fp16_scaler.step(optimizer)
            print("Complted step")
            fp16_scaler.update()
            print("Complted update")

 

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int)
    parser.add_argument("--batch_size", type=int, default=150)
    parser.add_argument("--num_epochs", type=int, default=200)
    parser.add_argument("--dist_url", type=str)
    args = parser.parse_args()
    run(args)
    
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

### Console log output
```
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
*****************************************
| distributed init (rank 1): tcp://x.x.x.x:29500
| distributed init (rank 0): tcp://x.x.x.x:29500

0it [00:03, ?it/s]
Completed zero grad
Completed backward pass
````

## Command to launch in distributed mode
```
python -m torch.distributed.launch  --nproc_per_node 2 --node_rank 0 --nnodes 1 train.py --dist_url tcp://x.x.x.x:29500 --batch_size 512
```

## Expected behavior
Model gets stuck after the backward pass. However, the training works fine in single GPU training launched in distributed mode

<!-- A clear and concise description of what you expected to happen. -->

## Environment

 - PyTorch Version (e.g., 1.0): 1.7.1
 - OS (e.g., Linux): Debian
 - How you installed PyTorch (`conda`, `pip`, source): `pip`
 - Build command you used (if compiling from source):
 - Python version: 3.8.5
 - CUDA/cuDNN version: 11.0
 - GPU models and configuration: A100, 16GPUs
 - Any other relevant information:

## Additional context

Dataloader used during training is a `IterableDataset` with custom `worker_init_fn` which sees a subset of files per worker. 


cc @mcarilli @ptrblck @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23