Repository: pytorch
Issue: expected scalar type Half but found Float with torch.cuda.amp and torch.nn.DataParallel
Labels: triaged,module: data parallel,module: amp (automated mixed precision)

## 🐛 Bug

I try to using amp in Pytorch core with torch.nn.DataParallel for multi-gpu training. I wrap forward pass in model in autocast, but get error  

## To Reproduce

Steps to reproduce the behavior:

1. install sinkhorn_transformer==0.4.4, transformers
2. run following script

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import sinkhorn_transformer
from sinkhorn_transformer import SinkhornTransformerLM
from transformers import BertTokenizer, AdamW
from torch.cuda.amp import autocast, GradScaler 

def mask_tokens(inputs: torch.Tensor, tokenizer, mlm_probability=0.15, pad=True):
    """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
    labels = inputs.clone()
    # mlm_probability defaults to 0.15 in Bert
    probability_matrix = torch.full(labels.shape, mlm_probability)
    special_tokens_mask = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
    ]
    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)
    if tokenizer._pad_token is not None:
        padding_mask = labels.eq(tokenizer.pad_token_id)
        probability_matrix.masked_fill_(padding_mask, value=0.0)
    masked_indices = torch.bernoulli(probability_matrix).bool()
    labels[~masked_indices] = -100  # We only compute loss on masked tokens

    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)

    # 10% of the time, we replace masked input tokens with random word
    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
    random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long)
    inputs[indices_random] = random_words[indices_random]

    if pad:
        input_pads = tokenizer.max_len - inputs.shape[-1]
        label_pads = tokenizer.max_len - labels.shape[-1]
        
        inputs = F.pad(inputs, pad=(0,input_pads), value=tokenizer.pad_token_id)
        labels = F.pad(labels, pad=(0,label_pads), value=tokenizer.pad_token_id)
    
    # The rest of the time (10% of the time) we keep the masked input tokens unchanged
    return inputs, labels

def read_batch_from_dataset_restored(tokenizer):
    batch = []
    iterat = 0
    for it in range(1000):
        yield torch.randint(3, tokenizer.vocab_size, (4, 40960))

class LossWraper(nn.Module):
  def __init__(self, model, loss=None):
    super(LossWraper, self).__init__()
    self.model = model
    self.loss = loss
    
  @autocast()
  def forward(self, inputs, labels=None):
    loss_mx = labels != -100
    output = self.model(inputs)
    output = output[loss_mx].view(-1, tokenizer.vocab_size)
    labels = labels[loss_mx].view(-1)
    loss = self.loss(output, labels)
    return loss, output

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = SinkhornTransformerLM(
    num_tokens= tokenizer.vocab_size,
    dim = 768,
    depth = 12,
    max_seq_len = 40960,
    heads = 16,
    buckets = 64,
    causal = False,           # auto-regressive or not
    sinkhorn_iter = 7,        # number of sinkhorn iterations - default is set at reported best in paper
    n_sortcut = 2,            # use sortcut to reduce complexity to linear time
    temperature = 0.75,       # gumbel temperature - default is set at reported best in paper
    non_permutative = False,  # allow buckets of keys to be sorted to queries more than once
    attn_sort_net = True,     # attention to reorder the buckets, unlocks flexible sequence lengths
    ff_chunks = 10,           # feedforward chunking, from Reformer paper
    reversible = True,        # make network reversible, from Reformer paper
    ff_dropout = 0.1,         # feedforward dropout
    attn_dropout = 0.1,       # post attention dropout
    attn_layer_dropout = 0.1, # post attention layer dropout
    layer_dropout = 0.1,      # add layer dropout, from 'Reducing Transformer Depth on Demand' paper
    weight_tie = True,        # tie layer parameters, from Albert paper
    emb_dim = 128,            # embedding factorization, from Albert paper
    ff_glu = True,            # use GLU in feedforward, from paper 'GLU Variants Improve Transformer'
    n_local_attn_heads = 4,   # replace N heads with local attention, suggested to work well from Routing Transformer paper
)
device = torch.device("cuda")
loss_fn = nn.CrossEntropyLoss()
model = LossWraper(model, loss_fn)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
model.to(device)
scaler = GradScaler()
model = torch.nn.DataParallel(model)
model.train()
for ep in range(50):
    for it in read_batch_from_dataset_restored(tokenizer):
        optimizer.zero_grad()
        torch.cuda.empty_cache()
        inp, mask = mask_tokens(it, tokenizer)
        inputs, labels = inp.to("cuda"), mask.to("cuda")
        with autocast():
            output = model(inputs, labels)
            loss = output[0]
            loss = loss.mean()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.25)
        scaler.step(optimizer)
        scaler.update()
```


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```python
RuntimeError  Traceback (most recent call last)
<ipython-input-7-7a29bfd8a4a8> in <module>
     63             loss = loss.mean()
---> 64         scaler.scale(loss).backward()
     65         scaler.unscale_(optimizer)

/opt/anaconda/envs/pytorch_apex_core_cuda_10.1/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    182                 products. Defaults to ``False``.
    183         """
--> 184         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    185 
    186     def register_hook(self, hook):

/opt/anaconda/envs/pytorch_apex_core_cuda_10.1/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    119         retain_graph = create_graph
    120 
--> 121     Variable._execution_engine.run_backward(
    122         tensors, grad_tensors, retain_graph, create_graph,
    123         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: expected scalar type Half but found Float
```

## Environment

PyTorch version: 1.6.0.dev20200514
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: Could not collect

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Quadro RTX 8000
GPU 1: Quadro RTX 8000

Nvidia driver version: 440.33.01
cuDNN version: Could not collect

Versions of relevant libraries:
[pip] numpy==1.18.1
[pip] torch==1.6.0.dev20200514
[pip] torchvision==0.7.0.dev20200514
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.1.243             h6bb024c_0  
[conda] mkl                       2020.1                      217  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.0.15           py38ha843d7b_0  
[conda] mkl_random                1.1.0            py38h962f231_0  
[conda] numpy                     1.18.1           py38h4f9e942_0  
[conda] numpy-base                1.18.1           py38hde5b4d6_1  
[conda] pytorch                   1.6.0.dev20200514 py3.8_cuda10.1.243_cudnn7.6.3_0    pytorch-nightly
[conda] torchvision               0.7.0.dev20200514      py38_cu101    pytorch-nightly



cc @mcarilli