Add drop_last arg for data loader
Author: setu4993
Pull request: yes

Add an extra argument to `TrainingArguments` that would be passed on to `Trainer` for use in DataLoader.

I ran into a problem while using the `Trainer` this week and the GPU expecting the full batch size of vector inputs, and put a workaround in place in the dataset class I was using, but would be useful to have this as an optional argument.

Comments:
Hi! This looks like a cool feature to add, indeed. I'm curious, what was the error you obtained because of an incomplete batch? It shouldn't raise an error if a batch is smaller than what it has previously seen.

I could see it being useful when the framework needs to trace with a given input size though, like with TPUs or with JAX.

Hey @LysandreJik, thanks for taking a look!

This error occurred on the last step of the epoch:
`RuntimeError: Gather got an input of invalid size: got [2, 1, 20, 256, 64], but expected [2, 2, 20, 256, 64] (gather at /AWS-PyTorch/torch/csrc/cuda/comm.cpp:231)`

Because of the nature of the error and it occurring on the last step, my suspicion was it was because of `drop_last`. I implemented a workaround for it and that stopped the error from re-appearing.

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

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

```diff
@@            Coverage Diff             @@
##           master    #4757      +/-   ##
==========================================
- Coverage   77.41%   77.34%   -0.08%     
==========================================
  Files         128      128              
  Lines       21059    21060       +1     
==========================================
- Hits        16302    16288      -14     
- Misses       4757     4772      +15     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4757?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `39.62% <ø> (ø)` | |
| [src/transformers/training\_args.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmluZ19hcmdzLnB5) | `76.66% <100.00%> (+0.26%)` | :arrow_up: |
| [src/transformers/data/processors/squad.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL3Byb2Nlc3NvcnMvc3F1YWQucHk=) | `28.66% <0.00%> (-6.37%)` | :arrow_down: |
| [src/transformers/tokenization\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHMucHk=) | `89.80% <0.00%> (-0.12%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `88.96% <0.00%> (+0.32%)` | :arrow_up: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/4757/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `82.03% <0.00%> (+1.35%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4757?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/4757?src=pr&el=footer). Last update [5bf9afb...9a8fb7a](https://codecov.io/gh/huggingface/transformers/pull/4757?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


I'm assuming you use GPU distribution – what method do you use for distribution? `nn.DataParallel`? 

@julien-c:  Correct, it is using `nn.DataParallel` under the hood.

It should work out of the box with torch.distributed instead of nn.DataParallel.

I have no objection to merging this though :)

Hmmm, this was on AWS SageMaker, so I'll double-check how it is implemented there.

Good recommendation on the change. Also, another question: Should the arg be separate for train and eval data loaders? I assumed not, but just wanted to confirm :).

No I can't think of a scenario where one would want to drop_last in train and not in eval (or inversely)

Thank you, merging