AlbertModel output is not HalfTensor when using apex fp16
Author: rasoolims
Pull request: no

# 🐛 Bug
Despite the fact that I turned the model to fp16 with apex, the hidden representation output is not half tensor (see code snippet for details) while class heads are half tensors.

## Information

Model I am using AlbertModel.

Language I am using the model on arbitrary:

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

```py
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AlbertModel, AlbertConfig
from transformers.modeling_albert import AlbertMLMHead
import apex

pad_token_id = 0
bos_token_id = 2
eos_token_id = 3
vocab_size = 20

config = {
             "attention_probs_dropout_prob": 0.1,
             "hidden_act": "gelu_new",
             "hidden_dropout_prob": 0.1,
             "embedding_size": 64,
             "hidden_size": 256,
             "initializer_range": 0.02,
             "intermediate_size": 1024,
             "max_position_embeddings": 512,
             "num_attention_heads": 2,  # smaller than usual
             "num_hidden_layers": 2,  # smaller than usual
             "num_hidden_groups": 1,
             "net_structure_type": 0,
             "gap_size": 0,
             "num_memory_blocks": 0,
             "inner_group_num": 1,
             "down_scale_factor": 1,
             "type_vocab_size": 2,
             "vocab_size": vocab_size,
             "pad_token_id": pad_token_id,
             "bos_token_id": bos_token_id,
             "eos_token_id": eos_token_id,
         }

albert_config = AlbertConfig(**config)
encoder: AlbertModel = AlbertModel(albert_config)
masked_lm = AlbertMLMHead(albert_config)
optimizer = torch.optim.Adam(encoder.parameters(), lr=0.0001)
encoder = encoder.cuda()
model, optimizer = apex.amp.initialize(encoder, optimizer, opt_level="O1", )

# When giving LongTensor as input, the class heads are half tensors,
# but hidden representations are not half!
long_input = torch.randint(1, 10, (10,5)).cuda()
f1, m1= encoder(long_input)

f1.type()
"""
'torch.cuda.FloatTensor'
"""
m1.type()
"""
'torch.cuda.HalfTensor'
"""

# When giving HalfTensor as input, it crashes
half_input = long_input.half()
f2, m2= encoder(half_input)

"""
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/transformers/modeling_albert.py", line 570, in forward
    input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/transformers/modeling_bert.py", line 173, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/torch/nn/modules/sparse.py", line 114, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/mnt/castor/seas_home/r/rasooli/torch_env/lib64/python3.6/site-packages/torch/nn/functional.py", line 1484, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError:
"""
```

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 2.7.0
- Platform:
- Python version: Python 3.6.10
- PyTorch version (GPU?): 1.4.0
- Tensorflow version (GPU?):
- Using GPU in script?: Yes, NVIDIA-SMI 418.67; GeForce RTX 208
- Using distributed or parallel set-up in script?: Optional parallel


Comments:
Hi @rasoolims, 

This seems correct. Basically, using `opt_level="O1"` means apex will add some downcasting to `float16` from a set of whitelisted operations, such as GEMM or Conv. This way, these operations will benefit from using Tensor Cores on latest devices, achieving higher throughput.. 

In the other hand, for some operations you want the whole spectrum of representable values to keep a very high accuracy in the output, this is true for many activation functions such as `Softmax` or `GeLU`.

What you observe here with your ouputs (`f1, m1`) is directly related to the dynamic downcasting of apex: 

- `f1` : Comes from a GeLU operation from AlbertaLayer, which is not downcasted to `float16`
- `m1`: Comes from a Linear layer operation, which is implemented through `gemm` and would greatly benefits from using Tensor Cores.

In addition, you should not be doing any input type conversion / downcasting when using `opt_level="O1"`.

Hi @mfuntowicz 
Thanks for the response.
Do you mean with the current setting, it is better to just use fp32? Or do you recommend changing the opt_level or activation function?


It really depends what you want to do and which level of Mixed-Precision training you want to achieve. 

With O1, only the operations which can expected improvements by running specialised CUDA/CuDNN fp16 kernels (on Tensor Cores) will be patched to have fp16 weights and input/output conversions.

With 02, all the weights of the model are going to be converted to fp16 with the exception of some layers like Batch Norm, so you have a quasi complete fp16 training. 

With O3, everything is run through fp16.

#thanks @mfuntowicz