Repository: huggingface/transformers (Python)
Task: Bug Fix

Issue:
FunnelTransformerForSequenceClassification crashes when fine tuning with mixed precision flag
## 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: 3.2.0
- Platform: Linux-4.15.0-45-generic-x86_64-with-debian-buster-sid
- Python version: Python 3.7.7
- PyTorch version (GPU?): 1.5.1 (True)
- Tensorflow version (GPU?): 2.2.0 (True)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: No

### Who can help
@sgugger As I saw you were the one who worked on the PR implementing Funnel Transformer

## Information

Model I am using: Funnel Transformer

The problem arises when using:
* [ o ] the official example scripts: (give details below)
* [ x ] my own modified scripts:
Only when enabling the mixed precision flag. I am now training the model without it, but I had to lower the batch size, thus increasing the training time.
I have to mention that I just fined tuned a `roberta-base` model using `fp16=True` and `fp16_opt_level='O1'`, thus nvidia APEX is properly installed/configured.

The tasks I am working on is:
* [ o ] an official GLUE/SQUaD task: (give the name)
* [ x ] my own task or dataset:
Basically I am trying to fine tune `FunnelForSequenceClassification` using my own custom data-set:
```python
# some code to load data from CSV
# ...
# wrapper around PyTorch for holding datasets
class IMDbDataset(torch.utils.data.Dataset):
    # same code as in the Huggingface docs
    # ...

# load tokenizer
tokenizer = FunnelTokenizer.from_pretrained('funnel-transformer/large-base')

# tokenize texts
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
val_encodings = tokenizer(val_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)

train_dataset = IMDbDataset(train_encodings, train_labels)
val_dataset = IMDbDataset(val_encodings, val_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)

# training args used
training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total number of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=64,  # batch size for evaluation
    #learning_rate=35e-6,
    weight_decay=0.01,               # strength of weight decay
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    logging_dir='./logs',            # directory for storing logs
    logging_steps=10,
    fp16=True,
    fp16_opt_level='O1'       # here I tried both O1 and O2 with the same result
)

model = FunnelForSequenceClassification.from_pretrained('funnel-transformer/large-base',
                                                        return_dict=True,
                                                        num_labels=max(train_labels)+1)

trainer = Trainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset             # evaluation dataset
)

trainer.train()

trainer.save_model('funnel')
```


## To reproduce

Steps to reproduce the behavior:

1. Run script
2. Wait for script to reach the training part

Stacktrace:
```
  File "funnel.py", line 89, in <module>
    trainer.train()
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 741, in train
    tr_loss += self.training_step(model, inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 1046, in training_step
    loss = self.compute_loss(model, inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 1070, in compute_loss
    outputs = model(**inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 1263, in forward
    return_dict=return_dict,
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 950, in forward
    return_dict=return_dict,
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 655, in forward
    layer_output = layer(query, key, value, attention_inputs, output_attentions=output_attentions)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 602, in forward
    attn = self.attention(query, key, value, attention_inputs, output_attentions=output_attentions)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 548, in forward
    content_score = torch.einsum("bind,bjnd->bnij", q_head + r_w_bias, k_head)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/functional.py", line 292, in einsum
    return _VF.einsum(equation, operands)
RuntimeError: Expected object of scalar type Float but got scalar type Half for argument #2 'mat2' in call to _th_bmm
```

[This](https://github.com/NVIDIA/apex/issues/302#issuecomment-552198322) seems like a very similar issue.

## Expected behavior

We should be able to train the model with mixed precision to use VRAM more efficiently.
FunnelTransformerForSequenceClassification crashes when fine tuning with mixed precision flag
## 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: 3.2.0
- Platform: Linux-4.15.0-45-generic-x86_64-with-debian-buster-sid
- Python version: Python 3.7.7
- PyTorch version (GPU?): 1.5.1 (True)
- Tensorflow version (GPU?): 2.2.0 (True)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: No

### Who can help
@sgugger As I saw you were the one who worked on the PR implementing Funnel Transformer

## Information

Model I am using: Funnel Transformer

The problem arises when using:
* [ o ] the official example scripts: (give details below)
* [ x ] my own modified scripts:
Only when enabling the mixed precision flag. I am now training the model without it, but I had to lower the batch size, thus increasing the training time.
I have to mention that I just fined tuned a `roberta-base` model using `fp16=True` and `fp16_opt_level='O1'`, thus nvidia APEX is properly installed/configured.

The tasks I am working on is:
* [ o ] an official GLUE/SQUaD task: (give the name)
* [ x ] my own task or dataset:
Basically I am trying to fine tune `FunnelForSequenceClassification` using my own custom data-set:
```python
# some code to load data from CSV
# ...
# wrapper around PyTorch for holding datasets
class IMDbDataset(torch.utils.data.Dataset):
    # same code as in the Huggingface docs
    # ...

# load tokenizer
tokenizer = FunnelTokenizer.from_pretrained('funnel-transformer/large-base')

# tokenize texts
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
val_encodings = tokenizer(val_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)

train_dataset = IMDbDataset(train_encodings, train_labels)
val_dataset = IMDbDataset(val_encodings, val_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)

# training args used
training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total number of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=64,  # batch size for evaluation
    #learning_rate=35e-6,
    weight_decay=0.01,               # strength of weight decay
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    logging_dir='./logs',            # directory for storing logs
    logging_steps=10,
    fp16=True,
    fp16_opt_level='O1'       # here I tried both O1 and O2 with the same result
)

model = FunnelForSequenceClassification.from_pretrained('funnel-transformer/large-base',
                                                        return_dict=True,
                                                        num_labels=max(train_labels)+1)

trainer = Trainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset             # evaluation dataset
)

trainer.train()

trainer.save_model('funnel')
```


## To reproduce

Steps to reproduce the behavior:

1. Run script
2. Wait for script to reach the training part

Stacktrace:
```
  File "funnel.py", line 89, in <module>
    trainer.train()
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 741, in train
    tr_loss += self.training_step(model, inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 1046, in training_step
    loss = self.compute_loss(model, inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/trainer.py", line 1070, in compute_loss
    outputs = model(**inputs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 1263, in forward
    return_dict=return_dict,
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 950, in forward
    return_dict=return_dict,
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 655, in forward
    layer_output = layer(query, key, value, attention_inputs, output_attentions=output_attentions)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 602, in forward
    attn = self.attention(query, key, value, attention_inputs, output_attentions=output_attentions)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/transformers/modeling_funnel.py", line 548, in forward
    content_score = torch.einsum("bind,bjnd->bnij", q_head + r_w_bias, k_head)
  File "/root/anaconda/envs/ai/lib/python3.7/site-packages/torch/functional.py", line 292, in einsum
    return _VF.einsum(equation, operands)
RuntimeError: Expected object of scalar type Float but got scalar type Half for argument #2 'mat2' in call to _th_bmm
```

[This](https://github.com/NVIDIA/apex/issues/302#issuecomment-552198322) seems like a very similar issue.

## Expected behavior

We should be able to train the model with mixed precision to use VRAM more efficiently.


Discussion:

Patch:
diff --git a/src/transformers/modeling_funnel.py b/src/transformers/modeling_funnel.py
--- a/src/transformers/modeling_funnel.py
+++ b/src/transformers/modeling_funnel.py
@@ -367,7 +367,6 @@ def pool_tensor(self, tensor, mode="mean", stride=2):
         # Stride is applied on the second-to-last dimension.
         stride = (stride, 1)
 
-        tensor = tensor.float()
         if mode == "mean":
             tensor = F.avg_pool2d(tensor, stride, stride=stride, ceil_mode=True)
         elif mode == "max":
@@ -554,7 +553,7 @@ def forward(self, query, key, value, attention_inputs, output_attentions=False):
         attn_score = attn_score.float()
         # perform masking
         if attention_mask is not None:
-            attn_score = attn_score - INF * attention_mask[:, None, None].float()
+            attn_score = attn_score - INF * (1 - attention_mask[:, None, None].float())
         # attention probability
         attn_prob = torch.softmax(attn_score, dim=-1, dtype=dtype)
         attn_prob = self.attention_dropout(attn_prob)
@@ -856,7 +855,9 @@ class FunnelForPreTrainingOutput(ModelOutput):
         attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
             Mask to avoid performing attention on padding token indices.
             Mask values selected in ``[0, 1]``:
-            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **maked**.
 
             `What are attention masks? <../glossary.html#attention-mask>`__
         token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
diff --git a/src/transformers/modeling_tf_funnel.py b/src/transformers/modeling_tf_funnel.py
--- a/src/transformers/modeling_tf_funnel.py
+++ b/src/transformers/modeling_tf_funnel.py
@@ -555,7 +555,7 @@ def call(self, query, key, value, attention_inputs, output_attentions=False, tra
             attn_score = tf.cast(attn_score, tf.float32)
         # perform masking
         if attention_mask is not None:
-            attn_score = attn_score - INF * tf.cast(attention_mask[:, None, None], tf.float32)
+            attn_score = attn_score - INF * (1 - tf.cast(attention_mask[:, None, None], tf.float32))
         # attention probability
         attn_prob = tf.nn.softmax(attn_score, axis=-1)
         if dtype != tf.float32:


Test patch:
diff --git a/tests/test_modeling_funnel.py b/tests/test_modeling_funnel.py
--- a/tests/test_modeling_funnel.py
+++ b/tests/test_modeling_funnel.py
@@ -428,16 +428,16 @@ def test_inference_tiny_model(self):
         model = FunnelModel.from_pretrained("sgugger/funnel-random-tiny")
         output = model(input_ids, token_type_ids=token_type_ids)[0].abs()
 
-        expected_output_sum = torch.tensor(2344.9023)
-        expected_output_mean = torch.tensor(0.8053)
+        expected_output_sum = torch.tensor(2344.8352)
+        expected_output_mean = torch.tensor(0.8052)
         self.assertTrue(torch.allclose(output.sum(), expected_output_sum, atol=1e-4))
         self.assertTrue(torch.allclose(output.mean(), expected_output_mean, atol=1e-4))
 
         attention_mask = torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]])
         output = model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)[0].abs()
 
-        expected_output_sum = torch.tensor(2363.2178)
-        expected_output_mean = torch.tensor(0.8115)
+        expected_output_sum = torch.tensor(2343.8425)
+        expected_output_mean = torch.tensor(0.8049)
         self.assertTrue(torch.allclose(output.sum(), expected_output_sum, atol=1e-4))
         self.assertTrue(torch.allclose(output.mean(), expected_output_mean, atol=1e-4))
 
@@ -448,7 +448,7 @@ def test_inference_model(self):
         inputs = tokenizer("Hello! I am the Funnel Transformer model.", return_tensors="pt")
         output = model(**inputs)[0]
 
-        expected_output_sum = torch.tensor(235.7827)
+        expected_output_sum = torch.tensor(235.7246)
         expected_output_mean = torch.tensor(0.0256)
         self.assertTrue(torch.allclose(output.sum(), expected_output_sum, atol=1e-4))
         self.assertTrue(torch.allclose(output.mean(), expected_output_mean, atol=1e-4))