[WIP] TF Trainer with TPUs
Author: jplu
Pull request: yes
Labels: name: wontfix

This PR is to make the TF trainer fully compliant with the TPUs.

Should fix #5042, #4996 and #4994.

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

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

```diff
@@            Coverage Diff             @@
##           master    #5065      +/-   ##
==========================================
- Coverage   79.02%   78.96%   -0.06%     
==========================================
  Files         138      138              
  Lines       24064    24089      +25     
==========================================
+ Hits        19017    19023       +6     
- Misses       5047     5066      +19     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/5065?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/5065/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyX3RmLnB5) | `17.92% <8.00%> (-0.77%)` | :arrow_down: |
| [src/transformers/training\_args\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/5065/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmluZ19hcmdzX3RmLnB5) | `53.19% <42.85%> (+2.02%)` | :arrow_up: |
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/5065/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `39.38% <0.00%> (-0.24%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5065/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `87.92% <0.00%> (+0.29%)` | :arrow_up: |

------

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


@Colanim Can you summarize your findings here? It would be better suited than going across multiple issues.

@jplu Sure !

* [x] Optimizer was passed as arguments to `_step` function, but we can pass only Tensors (issue #4994). This is fixed by ef02be8.

* [x] Optimization steps were not computed correctly : it was ignoring gradient accumulation (issue #5072). This is fixed by a897e09.

* [x] Computation of number of steps should take into account `dataloader_drop_last` (mentioned in #5072), and use `math.floor` in that case (instead of `math.ceil`). Fixed in #5051

* [ ] I'm still having error as described in #4996, but I didn't figure the reason yet. Maybe it's my model that has a problem, not `TFTrainer`.

Thanks for having sum up the issues here.

> * [ ]  Computation of number of steps should take into account `dataloader_drop_last` (mentioned in #5072), and use `math.floor` in that case (instead of `math.ceil`). Not fixed yet.

This is fixed in PR #5051.

> * [ ]  I'm still having error as described in #4996, but I didn't figure the reason yet. Maybe it's my model that has a problem, not `TFTrainer`.

This is what I'm currently looking for, but still not figuring out why :/

Which TPUs are you using, over ctpu, Cloud AI or Colab?

> Which TPUs are you using, over ctpu, Cloud AI or Colab?

I'm currently using ctpu, but I could see similar issue when using Colab.

For me, this is the problem : https://github.com/huggingface/transformers/blob/5f721ad6e48c9d846de25c3fefa0e50a306cbf10/src/transformers/trainer_tf.py#L388-L389

Somehow the Exception is not catched on TPU, which crash the training. Using `max_steps` argument instead of `num_train_epochs` fix the problem because we repeat the dataset, and therefore we never have an out of range error.

In the eval code, since we iterate the dataset without repeat, it's causing the same problem.

Can I see all your logs output during the training?

What is the command line you are using to create your TPU and run the process? In order to be aligned with the same errors :)

@jplu finally I was wrong : even when using `max_steps` I'm having the problem.

It always happen at the beginning of the end of first epoch (for both validation and training). I think it's due to how the dataset is iterated. 

According to [this link](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu), we should have a single iterator for the whole training procedure (iterator should not be reset after each epoch), and always `repeat` the dataset indefinitely. Even for validation dataset.

Unfortunately I'm having trouble with GCP right now, I can't try things on my end..

Ok, can you share your command lines here please. Because even the size computation cannot be run on TPU with TF 2.2.

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.