Hyperparameter search w/ Optuna CUDA out of memory
Author: marcoabrate
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.1.1
- Platform: Google Colab and Ubuntu
- Python version: 3.6.9
- PyTorch version (GPU?): 1.7
- Using GPU in script?: Yes

### Who can help

- trainer: @sgugger
- Optuna: ??? 

## Information

Model I am using (Bert, XLNet ...): sshleifer/distilbart-cnn-12-6

## To reproduce

I am running an hyperparameter search with Optuna. I get a CUDA OOM error even if `per_device_train_batch_size` is set to 1 and the only parameters that I change are `learning_rate` and `gradient_accumulation_steps`. I have the same problem both with Google Colab and Ubuntu. Both of this environments have a 15 GB GPU.

The code I am running:

```
def hp_objective(metrics):
    loss = metrics.pop('eval_loss', None)
    _ = metrics.pop('epoch', None)
    _ = metrics.pop('eval_gen_len', None)

    return np.sum(list(metrics.values()))

def hp_space(trial):

    return {
    'learning_rate': trial.suggest_float('learning_rate', 1e-5, 1e-2, log=True),
    'gradient_accumulation_steps':\
        trial.suggest_categorical('gradient_accumulation_steps', [4, 8]),
    }

def model_init():
            model = AutoModelForSeq2SeqLM.from_pretrained(
                model_args.model_name_or_path,
                config=config,
                cache_dir=model_args.cache_dir,
            )
            # use task specific params
            use_task_specific_params(model, data_args.task)

            # set num_beams for evaluation
            if data_args.eval_beams is None:
                data_args.eval_beams = model.config.num_beams

            # set decoder_start_token_id for MBart
            if model.config.decoder_start_token_id is None and isinstance(tokenizer, MBartTokenizer):
                assert (
                    data_args.tgt_lang is not None and data_args.src_lang is not None
                ), "mBart requires --tgt_lang and --src_lang"
                model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.tgt_lang]

            if model_args.freeze_embeds:
                freeze_embeds(model)
            if model_args.freeze_encoder:
                freeze_params(model.get_encoder())
                assert_all_frozen(model.get_encoder())

            return model

trainer = Seq2SeqTrainer(
            model_init=model_init,
            config=config,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=Seq2SeqDataCollator(tokenizer, data_args, training_args.tpu_num_cores),
            compute_metrics=compute_metrics_fn,
            data_args=data_args,
        )

logger.info("*** Hyperparameters Search ***")
start_time = time.time()

trainer.hyperparameter_search(
            direction = "maximize",
            compute_objective = hp_objective,
            hp_space = hp_space,
            backend = "optuna")
```
The error:

```
[INFO|modeling_utils.py:1149] 2021-02-01 15:24:41,640 >> All the weights of BartForConditionalGeneration were initialized from the model checkpoint at sshleifer/distilbart-cnn-12-6.
If your task is similar to the task the model of the checkpoint was trained on, you can already use BartForConditionalGeneration for predictions without further training.
02/01/2021 15:24:41 - INFO - utils -   using task specific params for summarization: {'early_stopping': True, 'length_penalty': 2.0, 'max_length': 142, 'min_length': 56, 'no_repeat_ngram_size': 3, 'num_beams': 4}

[W 2021-02-01 15:24:42,103] Trial 8 failed because of the following error: RuntimeError('CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 14.76 GiB total capacity; 13.66 GiB already allocated; 13.75 MiB free; 13.83 GiB reserved in total by PyTorch)',)
Traceback (most recent call last):
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/_optimize.py", line 211, in _run_trial
    value_or_values = func(trial)
  File "/home/ubuntu/transformers/src/transformers/integrations.py", line 168, in _objective
    trainer.train(model_path=model_path, trial=trial)
  File "/home/ubuntu/transformers/src/transformers/trainer.py", line 622, in train
    self.model = model.to(self.args.device)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 612, in to
    return self._apply(convert)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  [Previous line repeated 2 more times]
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 381, in _apply
    param_applied = fn(param)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 610, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 14.76 GiB total capacity; 13.66 GiB already allocated; 13.75 MiB free; 13.83 GiB reserved in total by PyTorch)
Traceback (most recent call last):
  File "/home/ubuntu/transformers/examples/seq2seq/finetune_trainer.py", line 435, in <module>
    main()
  File "/home/ubuntu/transformers/examples/seq2seq/finetune_trainer.py", line 350, in main
    backend = "optuna")
  File "/home/ubuntu/transformers/src/transformers/trainer.py", line 1077, in hyperparameter_search
    best_run = run_hp_search(self, n_trials, direction, **kwargs)
  File "/home/ubuntu/transformers/src/transformers/integrations.py", line 178, in run_hp_search_optuna
    study.optimize(_objective, n_trials=n_trials, timeout=timeout, n_jobs=n_jobs)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/study.py", line 385, in optimize
    show_progress_bar=show_progress_bar,
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/_optimize.py", line 73, in _optimize
    progress_bar=progress_bar,
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/_optimize.py", line 164, in _optimize_sequential
    trial = _run_trial(study, func, catch)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/_optimize.py", line 262, in _run_trial
    raise func_err
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/optuna/_optimize.py", line 211, in _run_trial
    value_or_values = func(trial)
  File "/home/ubuntu/transformers/src/transformers/integrations.py", line 168, in _objective
    trainer.train(model_path=model_path, trial=trial)
  File "/home/ubuntu/transformers/src/transformers/trainer.py", line 622, in train
    self.model = model.to(self.args.device)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 612, in to
    return self._apply(convert)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 359, in _apply
    module._apply(fn)
  [Previous line repeated 2 more times]
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 381, in _apply
    param_applied = fn(param)
  File "/home/ubuntu/miniconda3/envs/magma/lib/python3.6/site-packages/torch/nn/modules/module.py", line 610, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 14.76 GiB total capacity; 13.66 GiB already allocated; 13.75 MiB free; 13.83 GiB reserved in total by PyTorch)
```


## Expected behavior

The GPU should never go OOM, since the batch size is 1 in all trials.


Comments:
I don't think optuna properly optimizes GPU memory. We don't have support from them so you may be better using ray-tune, where the maintainers happily reply to question on our GitHub in case of problems. 

Thank you. I was using Optuna because with RayTune I get an error even before the first trial starts. I will open an issue about the RayTune error.