torch.einsum gets wrong results randomly when training with multi-gpu
Labels: oncall: distributed

### 🐛 Describe the bug

When training with multi-gpu, I find that the loss will grow rapidly and randomly. So I investigated the cause of the problem through debugging and found that one of the element in the result matrix became very large when torch.einsum was executed. And this problem does not appear when DistributedDataParallel is not used.
I implemented the DistributedDataParallel method of multi-gpu training by following the official tutorial, and everything was fine except for this problem, the loss dropped normally in the first few epochs.

here is the code :
``` python
scores: torch.Tensor = torch.einsum("bdhn,bdhm->bhnm", query, key) / dim**0.5
```
query and key both are correct, and this is a part of the first result:
![image](https://user-images.githubusercontent.com/34328687/181249062-125fd8d3-bb24-4c67-8613-9d49513a43f7.png)
I run the code once more, and this is the second result:
![840b95eb642e569f21692900d8c2eef](https://user-images.githubusercontent.com/34328687/181249250-d42c7f5f-2f1f-4879-b9c7-47873810d523.png)

By the way, data type is half. and here is my multi-gpu training script
``` python
def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "12355"
    os.environ["NCCL_SOCKET_IFNAME"] = "docker0"
    os.environ["GLOO_SOCKET_IFNAME"] = "docker0"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size, epoch, lr, graphs_pr: Dict, output: List):
    print(f"Training on rank {rank}.")
    setup(rank, world_size)
    try:
        torch.cuda.set_device(rank)
        with open(os.path.join(cache_path, "traning.plk"), "rb") as f:
            training: Training = dill.load(f)
        training.notify_callback = notify
        training.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(training.model)
        training.to(rank)
        # Train loss
        training.model = MatchingDDP(training.model, device_ids=[rank])
        log_loss = losses.LandmarkLogLoss(training.model)
        log_loss2 = losses.LaneLogLoss()
        train_loss = losses.ProductLoss([log_loss, log_loss2])
        # Dataset
        with open(os.path.join(cache_path, "dataset.plk"), "rb") as f:
            dataset: NewGtDataset = dill.load(f)
            dataset.load(os.path.join(cache_path, "dataset"))
        # Dataloader
        sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
        dataloader = DataLoader(dataset, batch_size, sampler=sampler, pin_memory=True)
        dataloader.collate_fn = Lambda1(training.model, dataset).f
        # Evaluator and update_pr
        with open(os.path.join(cache_path, "evaluator.plk"), "rb") as f:
            evaluator: evaluation.SubmapEvaluator = dill.load(f)

        def update_pr(data):
            graphs_pr.update(evaluator.eval_graphs_pr(data, training.model))

        training.fit(
            train_objective=(dataloader, train_loss),
            epochs=epoch,
            scheduler_name="ConstantLR",
            optimizer_params={"lr": lr},
            use_amp=True,
            callback=update_pr,
            checkpoint_path="saved_models" if rank == 0 else None,
            checkpoint_save_steps=40,
            checkpoint_save_total_limit=1,
            rank=rank,
        )
        if rank == 0:
            training.to("cpu")
            output.append(training.model.state_dict())
            output.append(training.epoch)
            output.append(training.global_step)
    except Exception as e:
        notify(e)
    finally:
        dist.destroy_process_group()
```


### Versions

pytorch = 1.12.0
python = 3.9.12
cuda = 10.2
devices: TITAN X (Pascal) * 8

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501