hyperparameter_serch() does not consider LoRA parameters like r  to be finetuned. 
Author: rajumishra8107
Pull request: no
Labels: name: Feature request

### Feature request

Hyperparameter search is a must known activity to get hyperparameters  which are for optimised machine learning or deep learning model output. I was trying hyperparameter_search() method to get optimal hyperparameters values. I wanted to tune LoRA confugarition parameters like rank r and alpha too. But i found that it is not able to finetune LoRA configuration paramaters. 

### Motivation

Following code part will give you the idea 
```python
def model_init():
    device_map = {"": torch.cuda.current_device(
                 )} if torch.cuda.is_available() else None
​
    model_kwargs_dict = dict(
    # set this to True if your GPU supports it 
    #(Flash Attention drastically speeds up model computations)
    #attn_implementation="flash_attention_2", 
    torch_dtype="auto",
    # set to False as we're going to use gradient checkpointing
    use_cache=False, 
    device_map=device_map,
    )
    device_map = {"": torch.cuda.current_device(
             )} if torch.cuda.is_available() else None
​
    bnb_config_args = dict(load_in_4bit = True,
                   bnb_4bit_quant_type = "nf4",
                   bnb_4bit_compute_dtype = torch.bfloat16,
                   bnb_4bit_use_double_quant = False)
    bnb_config = BitsAndBytesConfig(
                                **bnb_config_args
                     )
    model_kwargs_dict["quantization_config"] = bnb_config
    model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m", 
                                                return_dict=True, 
                                                #**model_kwargs_dict
                                               )
    print(peft_config)
    model = get_peft_model(model, peft_config = peft_config)
    return model

dataset = load_dataset("imdb", split="train")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")


dataset1 = dataset.select([0, 10, 20, 30, 40, 50])
dataset2 = dataset.select([0, 10, 20, 30, 40, 50])


trainer = SFTTrainer(
    model=None,
    args=training_args,
    model_init=model_init,
    tokenizer=tokenizer,
    train_dataset=dataset1,
    eval_dataset= dataset2,
    dataset_text_field="text",
    max_seq_length=512,)


def optuna_hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [16, 32, 64, 128]),
        "r": trial.suggest_float("r", 2, 4, log=True),
    }


trainer.hyperparameter_search(direction=["minimize"],
                            backend="optuna",
                            hp_space=optuna_hp_space,
                            n_trials=2)
``` 

The code above has resulted in output like 

[I 2024-02-29 09:42:36,869] A new study created in memory with name: no-name-331fbdff-6465-42f8-9c97-ad5c6c8c4703
**Trying to set r in the hyperparameter search but there is no corresponding field in `TrainingArguments`.**

### Your contribution

NA

Comments:
FYI @younesbelkada 

I also found out that `hyperparameter_serch()` can only tune the parameters in `TrainingArguments`, not the model parameters or the LoRA parameters. This is because `Trainer` only knows the `TrainingArguments` parameters. To tune the model parameters, I use each call of model_init to generate a model with new parameters (using `optuna` to provide the parameters for the trial), part of the code implementation is as follows:
```python
def model_init(trial: Optional[Union[optuna.trial.Trial,optuna.trial.FrozenTrial]]):
    # seed for model
    if training_args.tuner:
        transformers.enable_full_determinism(training_args.seed)  # each trial should have same seed
    else:
        pass # for run ordinary model

    config_cp = copy.deepcopy(config) # model default params
    if trial is None:
        model = BiEncoderForClassification(config=config_cp)
        return model

    # update params
    update_params = {key.replace('tuner_',''):value
                     for key,value in tokenizer_model_dict.items() if key.startswith('tuner_')}
    update_params = modify_dict_with_trial(update_params, trial) # trial.suggest_xxx
    logger.info(f"Trial Model: {update_params}")

    # update config_cp
    config_cp.update(update_params)

    # save config for trial
    # model/data/training
    all_config = combine_dict(config_cp.to_dict(), asdict(data_args), training_args.to_dict())
    trial.set_user_attr('config', all_config)

    return BiEncoderForClassification(config=config_cp)
```
This works fine on a single process, but can go wrong in a distribution with multiple processes. There are two possible problems. First, the non-main process creates the study again (`optuna` backend and with the `storage=optuna.storages.RDBStorage` parameter option). Second, the non-main process does not correctly identify the model parameters that need to be tuned. I've found the cause of the second issue: in the source code below, it's clear that the non-main process needs the master process to pass `optuna.trial.FixedTrial(trial.params, trial.number)`, but `model_init` is called in the `trainer.train` method. This time, `FixedTrial` does not determine the model parameters suggested by `optuna`.
```python
def run_hp_search_optuna(trainer, n_trials: int, direction: str, **kwargs) -> BestRun:
    import optuna
    from accelerate.utils.memory import release_memory

    if trainer.args.process_index == 0:

        def _objective(trial: optuna.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_space(trial)
                fixed_trial = optuna.trial.FixedTrial(trial.params, trial.number)
                trial_main_rank_list = [fixed_trial]
                torch.distributed.broadcast_object_list(trial_main_rank_list, src=0)
                trainer.train(resume_from_checkpoint=checkpoint, trial=trial)
            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)

            # Free GPU memory
            trainer.model_wrapped, trainer.model = release_memory(trainer.model_wrapped, trainer.model)
            trainer.accelerator.clear()

            return trainer.objective

        timeout = kwargs.pop("timeout", None)
        n_jobs = kwargs.pop("n_jobs", 1)
        gc_after_trial = kwargs.pop("gc_after_trial", False)
        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)
        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
            trial_main_rank_list = [None]
            if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED:
                raise RuntimeError("only support DDP optuna HPO for ParallelMode.DISTRIBUTED currently.")
            torch.distributed.broadcast_object_list(trial_main_rank_list, src=0)
            trainer.train(resume_from_checkpoint=None, trial=trial_main_rank_list[0])
            # 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
```
Is it possible for the developer to put three of these lines of code after the `model_init` call in the `trainer.train` method?
```python
fixed_trial = optuna.trial.FixedTrial(trial.params, trial.number)
trial_main_rank_list = [fixed_trial]
torch.distributed.broadcast_object_list(trial_main_rank_list, src=0)
```

It seems that we can use the `hp_space` parameter of the `hyperparameter_search` function without changing the source code. For the first problem mentioned before (the duplicate creation of the study), we can just create sqlite.db in the main process beforehand. For the second problem, we can use the `hp_space` parameter but return an empty dictionary (if we only tune the model parameters). In the `hp_space_hook` function, we can use `trial.suggest_xxx` to suggest the model parameters and not have to modify the `model_init` function (since the model_init function suggests the same value for the same parameter name in one experiment). To train the LoRA parameter, just add it to the model's parameters and modify it in both the hp_space and model_init functions.
```python
if training_args.tuner:
    # only tuner
    training_args.do_train = False
    training_args.do_eval = False

    # set hp_space_hook to initial trial callback each trial
    def hp_space_hook(trial: Union[optuna.trial.Trial,optuna.trial.FrozenTrial]):
        # set pruner callback for optuna
        callback_str_list: List[str] = trainer.callback_handler.callback_list.split('\n')
        current_module = sys.modules[__name__]
        trial_callback_class = getattr(current_module, "PruningCallback")
        if str_if_contain_in_str_list('PruningCallback', callback_str_list, mode='contained'):
            trainer.callback_handler.remove_callback(trial_callback_class)
        trainer.add_callback(trial_callback_class(trial, monitor=training_args.metric_for_best_model))

        # tuner model hyperparameters hook for trial
        # and in model_init func, trial will suggest the same hyperparameters
        # but we do not need to return model hyperparameters (because it will cause errors.)
        # update params
        update_params = {key.replace('tuner_', ''): value
                         for key, value in tokenizer_model_dict.items() if key.startswith('tuner_')}
        update_params = modify_dict_with_trial(update_params, trial)  # trial.suggest_xxx
        return dict()

    # set hp search
    # set url_path
    if not os.path.exists(training_args.tuner_results_dir):
        os.makedirs(training_args.tuner_results_dir)
    url_obj = os.path.join(training_args.tuner_results_dir, "tuner.db")
    url_path = "sqlite:///" + url_obj
    # create study in rank 0
    if not os.path.exists(url_obj):
        with training_args.main_process_first(local=False, desc='create study'):
            study = optuna.create_study(
                storage=optuna.storages.RDBStorage(
                    url_path
                )
            )
    best_trial = trainer.hyperparameter_search(
        #hp_space=lambda trial: dict(),
        hp_space=hp_space_hook,
        compute_objective=lambda metrics: metrics[training_args.metric_for_best_model],
        n_trials=training_args.n_trials,
        direction="minimize" if training_args.metric_for_best_model.endswith('loss') else "maximize",
        backend='optuna',
        study_name='CSTS',
        storage=optuna.storages.RDBStorage(
            url_path,
            heartbeat_interval=60,
            grace_period=120,
            failed_trial_callback=optuna.storages.RetryFailedTrialCallback(3),
        ),
        load_if_exists=True,
        sampler=optuna.samplers.TPESampler(),
        gc_after_trial=True,
        pruner=CombinePruner([
            optuna.pruners.MedianPruner(n_warmup_steps=30),
            optuna.pruners.PercentilePruner(0.1, n_warmup_steps=30)
        ])
        # pruner=optuna.pruners.NopPruner(),
    )
```
I tried a simple example that can be found here: https://github.com/032004129xuzhiyong/TransformerTest