Enable wandb logging for Ray Tune HPO runs
Author: krfricke
Pull request: yes

<!-- This line specifies which issue to close after the pull request is merged. -->
Currently wandb logging does not work with Ray Tune, as each process tries to call `wandb.log` without initializing wandb first.

This PR fixes this by wrapping the objective in Ray Tune's `wandb_mixin`, which makes sure each trial initializes wandb and logs to it.


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

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

```diff
@@            Coverage Diff             @@
##           master    #6791      +/-   ##
==========================================
+ Coverage   79.36%   79.78%   +0.41%     
==========================================
  Files         157      157              
  Lines       28569    28578       +9     
==========================================
+ Hits        22675    22800     +125     
+ Misses       5894     5778     -116     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6791?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `52.90% <0.00%> (+39.68%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9tb2JpbGViZXJ0LnB5) | `24.55% <0.00%> (-72.36%)` | :arrow_down: |
| [src/transformers/tokenization\_mbart.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fbWJhcnQucHk=) | `35.93% <0.00%> (-59.38%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_distilbert.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9kaXN0aWxiZXJ0LnB5) | `64.47% <0.00%> (-34.36%)` | :arrow_down: |
| [src/transformers/tokenization\_xlnet.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25feGxuZXQucHk=) | `66.66% <0.00%> (-23.43%)` | :arrow_down: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `72.25% <0.00%> (-8.71%)` | :arrow_down: |
| [src/transformers/configuration\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2JhcnQucHk=) | `90.00% <0.00%> (-4.00%)` | :arrow_down: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `85.21% <0.00%> (-1.26%)` | :arrow_down: |
| [src/transformers/generation\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3V0aWxzLnB5) | `96.66% <0.00%> (-0.28%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.66% <0.00%> (+0.25%)` | :arrow_up: |
| ... and [14 more](https://codecov.io/gh/huggingface/transformers/pull/6791/diff?src=pr&el=tree-more) | |

------

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


making sure that @borisdayma sees this PR.

Thanks, I've not played too much with ray tune yet but there seems to be 2 ways to integrate through Ray Tune libraries as per [the docs](https://docs.wandb.com/library/integrations/ray-tune).

However ideally, the setup, logging, etc would be handled directly by `Trainer` existing functions for clarity and concision (and also to support all existing loggers). Handling multiple logging runs should be done within `hyperparameter_search` if possible.

Could the setup methods be wrapped in a new function and called during the search, in order to avoid duplicating the same logic.
For wandb, forcing a new run just requires `wandb.init(reinit=True)` so it works both in notebooks and scripts.

Note: Use this argument **only** while using `hyperparameter_search` as users can currently call manually `wandb.init` before (for example when using pytorch-lightning, sweeps, or keras + huggingface), making the call within the `Trainer` a "noop" (because it does not have `reinit=True`).

Thanks for your comments.

@borisdayma, simply moving the logger setup to `train()` would do the trick in any case, as it is called from the hyperparameter search methods. This should also work for Optuna, not only for Ray Tune.

I created a PR for that here: #6850. Is this what you meant?

Closed in favor of #6850.