Fixing case in which `Trainer` hung while saving model in distributed training
Author: TevenLeScao
Pull request: yes

As found thanks to the great @mfuntowicz , the call to `store_flos` in `Trainer` can hang indefinitely, as it was only executed in the main thread and in some cases the other threads were already past this point. This PR moves this call in order to avoid this behaviour.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/7365?src=pr&el=h1) Report
> Merging [#7365](https://codecov.io/gh/huggingface/transformers/pull/7365?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/28cf873036d078b47fb9dd38ac3421a7c874da44?el=desc) will **increase** coverage by `2.52%`.
> The diff coverage is `50.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #7365      +/-   ##
==========================================
+ Coverage   76.58%   79.11%   +2.52%     
==========================================
  Files         181      181              
  Lines       34828    34827       -1     
==========================================
+ Hits        26674    27552     +878     
+ Misses       8154     7275     -879     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/7365?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `55.70% <50.00%> (+0.08%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_lxmert.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9seG1lcnQucHk=) | `22.14% <0.00%> (-72.41%)` | :arrow_down: |
| [src/transformers/modeling\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yYWcucHk=) | `25.39% <0.00%> (-51.59%)` | :arrow_down: |
| [src/transformers/modeling\_marian.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19tYXJpYW4ucHk=) | `60.00% <0.00%> (-30.00%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9iZXJ0LnB5) | `69.10% <0.00%> (-29.80%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_xlm.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl94bG0ucHk=) | `88.37% <0.00%> (-4.87%)` | :arrow_down: |
| [src/transformers/activations.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9hY3RpdmF0aW9ucy5weQ==) | `81.81% <0.00%> (-4.55%)` | :arrow_down: |
| [src/transformers/configuration\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3JhZy5weQ==) | `97.77% <0.00%> (-2.23%)` | :arrow_down: |
| [src/transformers/modeling\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19hdXRvLnB5) | `82.62% <0.00%> (-1.41%)` | :arrow_down: |
| [src/transformers/modeling\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iYXJ0LnB5) | `93.93% <0.00%> (-0.34%)` | :arrow_down: |
| ... and [15 more](https://codecov.io/gh/huggingface/transformers/pull/7365/diff?src=pr&el=tree-more) | |

------

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


Reminder that we need to add some CI infra (and tests) for multi-gpu and/or multi-node setups