🐛 [TFTrainer] Wrong number of optimization steps
# 🐛 Bug

It seems TFTrainer compute the wrong number of optimization steps :

https://github.com/huggingface/transformers/blob/e4aaa4580515446cd5a2972ab42fec0b95819c84/src/transformers/trainer_tf.py#L79

**It does not take into account the gradient accumulation**.

---

I believe this line should be changed to :

```
if self.args.dataloader_drop_last:
    approx = math.floor
else:
    approx = math.ceil
self.train_steps: int = approx(self.num_train_examples / (self.args.train_batch_size * self.args.gradient_accumulation_steps))
```

---

Also, on TPU `drop_remainder` should be called on gradient accumulation steps as well. 

https://github.com/huggingface/transformers/blob/e4aaa4580515446cd5a2972ab42fec0b95819c84/src/transformers/trainer_tf.py#L81-L86

Should be changed to :

```
ds = (
    self.train_dataset.cache()
    .shuffle(self.num_train_examples)
    .batch(self.args.train_batch_size * self.args.gradient_accumulation_steps, drop_remainder=self.args.dataloader_drop_last)
    .unbatch()
    .batch(self.args.train_batch_size)
    .prefetch(tf.data.experimental.AUTOTUNE)
)
```

---

@jplu Maybe to add in #5065 ?

You are right, the number of optimization steps is already fixed in an another PR (#5051). And indeed the batch size for TPUs has to take into account the gradient accumulation, but I put this for later because we don't have yet a proper way to differenciate a run on TPU/CPU/GPU. Also the `unbatch` then `batch` is not needed. Same thing apply later when computing the loss in the training step.<|||||>Thanks for your answer @jplu !

Can I ask you clarification about this :

>Also the unbatch then batch is not needed. Same thing apply later when computing the loss in the training step.

We still need the data to be batched based on the batch_size_per_device, not the total batch size, right ?
Is the `floor` / `ceil` needed ?

Is there any other changes I should apply locally to make it work ?<|||||>I'm pretty sure that the TPUs as to be set for the full size of batches (including those with the accumulation).

| Is the floor / ceil needed ?
`floor` and `ceil` are also needed in order to be sure you get the proper approximation.

| Is there any other changes I should apply locally to make it work ?

There are certainly other changes to make, but I still have to figure out which one :smile: Don't hesitate to participate on the PR I have opened ^^<|||||>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.
