Allow additional keyword args to be passed to optuna hyperparameter search
Author: JanetVictorious
Pull request: no
Labels: name: Feature request

### Feature request

The issue with CUDA out of memory due to hyperparameter optimization has been addressed before ([old issue](https://github.com/huggingface/transformers/issues/13019)) but no implementation has been made to remedy this.

The fix would be quite simple to allow for the additional argument `gc_after_trial` to be passed to `hyperparameter_search()`:

```python
def run_hp_search_optuna(trainer, n_trials: int, direction: str, **kwargs) -> BestRun:
    import optuna

    if trainer.args.process_index == 0:

        def _objective(trial, checkpoint_dir=None):
            checkpoint = None
            if checkpoint_dir:
                for subdir in os.listdir(checkpoint_dir):
                    if subdir.startswith(PREFIX_CHECKPOINT_DIR):
                        checkpoint = os.path.join(checkpoint_dir, subdir)
            trainer.objective = None
            if trainer.args.world_size > 1:
                if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED:
                    raise RuntimeError("only support DDP optuna HPO for ParallelMode.DISTRIBUTED currently.")
                trainer._hp_search_setup(trial)
                torch.distributed.broadcast_object_list(pickle.dumps(trainer.args), src=0)
                trainer.train(resume_from_checkpoint=checkpoint)
            else:
                trainer.train(resume_from_checkpoint=checkpoint, trial=trial)
            # If there hasn't been any evaluation during the training loop.
            if getattr(trainer, "objective", None) is None:
                metrics = trainer.evaluate()
                trainer.objective = trainer.compute_objective(metrics)
            return trainer.objective

        timeout = kwargs.pop("timeout", None)
        n_jobs = kwargs.pop("n_jobs", 1)
        gc_after_trial = kwargs.pop("gc_after_trial", False)  # <--- Added arg
        directions = direction if isinstance(direction, list) else None
        direction = None if directions is not None else direction
        study = optuna.create_study(direction=direction, directions=directions, **kwargs)
        study.optimize(_objective, n_trials=n_trials, timeout=timeout, n_jobs=n_jobs, gc_after_trial=gc_after_trial)  # <--- Added arg
        if not study._is_multi_objective():
            best_trial = study.best_trial
            return BestRun(str(best_trial.number), best_trial.value, best_trial.params)
        else:
            best_trials = study.best_trials
            return [BestRun(str(best.number), best.values, best.params) for best in best_trials]
    else:
        for i in range(n_trials):
            trainer.objective = None
            args_main_rank = list(pickle.dumps(trainer.args))
            if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED:
                raise RuntimeError("only support DDP optuna HPO for ParallelMode.DISTRIBUTED currently.")
            torch.distributed.broadcast_object_list(args_main_rank, src=0)
            args = pickle.loads(bytes(args_main_rank))
            for key, value in asdict(args).items():
                if key != "local_rank":
                    setattr(trainer.args, key, value)
            trainer.train(resume_from_checkpoint=None)
            # If there hasn't been any evaluation during the training loop.
            if getattr(trainer, "objective", None) is None:
                metrics = trainer.evaluate()
                trainer.objective = trainer.compute_objective(metrics)
        return None
```

Then in the training script, calling `trainer.hyperparameter_search()` would allow for the additional arg `gc_after_trial`:

```python
best_trial = trainer.hyperparameter_search(
    direction='minimize',
    backend='optuna',
    hp_space=_hp_space,
    n_trials=model_args.hpo_trials,
    gc_after_trial=False,
)
```

### Motivation

I'm experiencing CUDA out of memory issues when running `trainer.hyperparameter_search()` with the optuna backend. After investigating a bit, the recommendation from optuna is to pass `gc_after_trial=True` as parameter to `study.optimize()` ([optuna reference](https://optuna.readthedocs.io/en/stable/faq.html#how-do-i-avoid-running-out-of-memory-oom-when-optimizing-studies)).

My proposal is to allow `gc_after_trial` to be passed as a `kwarg` and picked up in the `study.optimize()` call in `run_hp_search_optuna()` method ([source code](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/integration_utils.py#L258)).

I would like to be able to pass this argument as part of `trainer.hyperparameter_search()` with optuna backend.

### Your contribution

I can submit a PR for this change if there is a value in this feature.

Comments: