Add checkpointing to Ray Tune HPO
Author: krfricke
Pull request: yes

Saves checkpoints in Ray Tune HPO, enabling advanced schedulers like PBT.

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

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

```diff
@@            Coverage Diff             @@
##           master    #6747      +/-   ##
==========================================
- Coverage   79.01%   78.18%   -0.83%     
==========================================
  Files         157      157              
  Lines       28739    28782      +43     
==========================================
- Hits        22707    22503     -204     
- Misses       6032     6279     +247     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6747?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyX3V0aWxzLnB5) | `59.57% <0.00%> (-4.87%)` | :arrow_down: |
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `13.57% <7.14%> (+0.45%)` | :arrow_up: |
| [src/transformers/integrations.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9pbnRlZ3JhdGlvbnMucHk=) | `31.11% <9.09%> (-34.61%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9tb2JpbGViZXJ0LnB5) | `24.55% <0.00%> (-72.36%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9hbGJlcnQucHk=) | `21.47% <0.00%> (-69.44%)` | :arrow_down: |
| [src/transformers/tokenization\_mbart.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fbWJhcnQucHk=) | `57.14% <0.00%> (-39.69%)` | :arrow_down: |
| [src/transformers/modeling\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yb2JlcnRhLnB5) | `77.37% <0.00%> (-19.71%)` | :arrow_down: |
| [src/transformers/data/data\_collator.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL2RhdGFfY29sbGF0b3IucHk=) | `91.90% <0.00%> (-0.41%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.41% <0.00%> (-0.26%)` | :arrow_down: |
| [src/transformers/tokenization\_utils\_base.py](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfYmFzZS5weQ==) | `93.90% <0.00%> (-0.14%)` | :arrow_down: |
| ... and [10 more](https://codecov.io/gh/huggingface/transformers/pull/6747/diff?src=pr&el=tree-more) | |

------

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


```diff
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 3470a473..acf1503c 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -544,7 +544,8 @@ class Trainer:
             if trial.should_prune():
                 raise optuna.TrialPruned()
         elif self.hp_search_backend == HPSearchBackend.RAY:
-            self._tune_save_checkpoint()
+            if self.global_step % self.args.save_steps == 0:
+                self._tune_save_checkpoint()
             tune.report(objective=self.objective, **metrics)

     def _tune_save_checkpoint(self):
@@ -911,6 +912,8 @@ class Trainer:
             # search.
             _tb_writer = self.tb_writer
             self.tb_writer = None
+            _model = self.model
+            self.model = None
             # Setup default `resources_per_trial` and `reporter`.
             if "resources_per_trial" not in kwargs and self.args.n_gpu > 0:
                 n_jobs = int(kwargs.pop("n_jobs", 1))
```

This allows us to:

1. Not die when tuning BERT and
2. Not be dominated by saving latency.

Thanks for your suggestions. I moved the bulk of the hp search code to `integrations`, including the objective, since it depends on the search space. Is this what you had in mind?

Yes. I think we can split the function in two: one for ray, one for optuna and avoid a lot of tests this way to have some cleaner code (with a small duplication in the _objective function). I can do it in a separate PR if you want.

That would be great, thanks!

Merging and will follow up then.