Specify device in DataCollator 
Author: mfuntowicz
Pull request: yes

By setting `device` parameter in `DataCollator` we're able to allocate tensors directly on the right device at tensor creation and avoid moving data around afterwards.

This effectively avoid some snippet like this in the Trainer: 

```python

for k, v in some_dict.items():
     v.to(device)
```

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=h1) Report
> Merging [#4519](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/a08652772791fdaeed6f263b1a99926ca64be5dc&el=desc) will **decrease** coverage by `0.01%`.
> The diff coverage is `81.81%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4519/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #4519      +/-   ##
==========================================
- Coverage   77.83%   77.82%   -0.02%     
==========================================
  Files         123      123              
  Lines       20514    20513       -1     
==========================================
- Hits        15968    15964       -4     
- Misses       4546     4549       +3     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/data/data\_collator.py](https://codecov.io/gh/huggingface/transformers/pull/4519/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFfY29sbGF0b3IucHk=) | `89.70% <80.00%> (+0.47%)` | :arrow_up: |
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/4519/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `39.37% <100.00%> (-0.11%)` | :arrow_down: |
| [src/transformers/hf\_api.py](https://codecov.io/gh/huggingface/transformers/pull/4519/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9oZl9hcGkucHk=) | `93.06% <0.00%> (-4.96%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=footer). Last update [a086527...7f579cb](https://codecov.io/gh/huggingface/transformers/pull/4519?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


@mfuntowicz Perhaps it's useful to pull this through to the encode methods of the tokenizers so that you can pass a device and if return_tensors is used, the tensors are automatically pushed to the correct device? 

I tend to agree we need to have a device argument on the tokenizer `encode_like` methods, that would remove the need to iterate over the items to relocate on GPU/TPU.

Is there a common scheme we can use to do this on both Pytorch & TensorFlow (_I'm not very familiar with TensorFlow_)? 

I suspect strings might be the easiest way to handle this: 

- PyTorch: `device="cpu"`  || `device="cuda:0"`
- TensorFlow: `device="/device:cpu"`  ||  `device="/device:gpu:0"` 

I need to test things out futher here to see what the API may looks like! 👍 
Let's see what the other think of the propal for the Trainer and we can follow up on a dedicated PR asap.

> I tend to agree we need to have a device argument on the tokenizer `encode_like` methods, that would remove the need to iterate over the items to relocate on GPU/TPU.
> 
> Is there a common scheme we can use to do this on both Pytorch & TensorFlow (_I'm not very familiar with TensorFlow_)?
> 
> I suspect strings might be the easiest way to handle this:
> 
> * PyTorch: `device="cpu"`  || `device="cuda:0"`
> * TensorFlow: `device="/device:cpu"`  ||  `device="/device:gpu:0"`
> 
> I need to test things out futher here to see what the API may looks like! 👍
> Let's see what the other think of the propal for the Trainer and we can follow up on a dedicated PR asap.

Sorry, I also don't have much experience with TF so I can't really chip in on that. I guess the encode methods can accept a `device=` property that can be of type `Union[str, int, tf.device, torch.device]`. The device constructor (pt or tf) can depend on the already available `return_tensors` part which can be None, pt, or tf.

- `str`: use it to initialize the `torch.device(str)` or `tf.device(str))` - allows the users with a lot of freedom in case they  want do something like `tf.device('/job:bar/task:0/device:gpu:2')` (example from [docs](https://www.tensorflow.org/api_docs/python/tf/device))
- `int`: assume this is a GPU device id: `torch.device(f"cuda:{int}")` or `torch.device(f"/device:gpu:{int}")`
- a `device`: use the device

@mfuntowicz is the performance improvement of this change significant? How can we measure it?

Let me run some epochs to report quantitative numbers

Perhaps useful discussions here: https://stackoverflow.com/questions/28597014/python-why-is-accessing-instance-attribute-is-slower-than-local

In 100 calls to the variable (self or local) the answer reports 3 seconds difference. Not sure if that is worth it.

I'm all for speed improvements but I'm also really fond of readability. 

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


cc @LysandreJik 

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Should this still be implemented, or is this PR superseded by another one? @mfuntowicz @julien-c 

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.

Bump @julien-c  @mfuntowicz 

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.

Should I keep bumping this? @sgugger 

I don't think it's useful: the Trainer handles it already and for people that want to run their custom training loop, Accelerate can handle that.

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.