Fix for the issue of device-id getting hardcoded for token_type_ids during Tracing [WIP]
Author: HamidShojanazeri
Pull request: yes

# What does this PR do?

Using Torchscript Trace API to convert the HF models is creating an issue , where during the tracing some of the name/id of the device getting hardcoded for some of the tensors. As a result position_embedding or token_type_embedding fail when model is loaded for inference on another device, as they were tied to the device that they were traced on (e.g cpu, gpu id). This issue will raise when one needs to switch between devices and specially for multi-gpu inference where model is Torchscripted/traced.

For Bert model, device name get hardcoded for token_type_ids, this has been previously addressed for position_ids in [the merged PR](https://github.com/huggingface/transformers/pull/5773). This PR, fixes this issue by registering a buffer for token_type_ids. Similar changes for other models are required as well, that will submit PRs accordingly. 

The following code snippet can be used to test the issue and suggested fix.

```
import transformers
from pathlib import Path
import os
import json
import torch
from transformers import (AutoModelForSequenceClassification, AutoTokenizer, AutoModelForQuestionAnswering,
 AutoModelForTokenClassification, AutoConfig)
device1 = torch.device('cuda') # this can be changed to cuda:0 in multi-gpu use-case
device2 = torch.device('cpu')# this can be changed to cuda:1 in multi-gpu use-case
model_name = 'bert-base-uncased'
config = AutoConfig.from_pretrained(model_name,num_labels=2,torchscript=True)
model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)
tokenizer = AutoTokenizer.from_pretrained(model_name,do_lower_case=True)

dummy_input = "This is a dummy input for torch jit trace"
max_length = 20

inputs = tokenizer.encode_plus(dummy_input,max_length = int(max_length),pad_to_max_length = True, add_special_tokens = True, return_tensors = 'pt')

input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
print('device1 {}, device2 {}'.format(device1,device2))
outputs = model(**inputs)

model.to(device1).eval()

last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

traced_model = torch.jit.trace(model,(input_ids.to(device1),attention_mask.to(device1)))

torch.jit.save(traced_model, "bert.pt")
print("*************************** traced model graph on device 1 ***************************")
print(traced_model.graph)
print("\n")
loaded = torch.jit.load("bert.pt", map_location=device2)
print("\n")
print("*************************** model graph on loaded on device 2 ***************************")
print(loaded.graph)
outputs = loaded(input_ids.to(device2),attention_mask.to(device2))
print(outputs)
```
Error log :
[bert_gpu_to_cpu.logs.txt](https://github.com/huggingface/transformers/files/6312550/bert_gpu_to_cpu.logs.txt)

Fix log:
[bert_gpu_to_cpu_fixed.logs.txt](https://github.com/huggingface/transformers/files/6312560/bert_gpu_to_cpu_fixed.logs.txt)


Fixes # (issue)
Registering a buffer for token_type_ids in the constructor and then resizing it in the forward method based on input-shape. 

## Before submitting

- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
issues #5664 and  #976


## Who can review?

@LysandreJik



Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

@LysandreJik as discussed off-line, I would appreciate if you could help with re-opening the PR, thanks.

> Looks very cool! We can deploy the same fix on the other models.

Thanks @sgugger, once this can be merged, I will start the other models and submit separate PRs for them.

Great, thanks a lot @HamidShojanazeri. Before we merge; could you apply the same changes to all models affected by the `fix-copies` script in your PR? If we merge this as is, we'll have the implementation be partially supported for these models, which is unwanted.

Thank you!

@LysandreJik sure, I will update the affected models as well. 

> Great, thanks a lot @HamidShojanazeri. Before we merge; could you apply the same changes to all models affected by the `fix-copies` script in your PR? If we merge this as is, we'll have the implementation be partially supported for these models, which is unwanted.
> 
> Thank you!

@LysandreJik  Updated!

Ran the GPU tests, works like a charm. As this gets implemented in other models, let's think of a test we can add similar to your snippet @HamidShojanazeri to ensure there is no regression.

Merging!

Thanks @LysandreJik sure, will sync off-line.

This introduced an issue in our slow tests that I'm patching in https://github.com/huggingface/transformers/pull/12336

> This introduced an issue in our slow tests that I'm patching in #12336

Thank a lot! @LysandreJik.