Fix hyperparameter search when optuna+deepseed
Author: corentin-ryr
Pull request: yes

## Problem Overview

In a distributed DeepSpeed training setting, using Optuna for hyperparameter optimization (HPO) causes hyperparameter inconsistencies across processes. The root of the problem is that `trainer._hp_search_setup`, which applies the hyperparameters for each trial, only executes on the main process. This causes discrepancies in settings such as `gradient_accumulation_steps` across processes, leading to timeout errors due to process misalignment during training.

The current code broadcasts the trainer arguments and applies them to the auxiliary processes but doesn't execute the code to update the deepspeed/accelerate configuration. I also think that the current code would not account for hyperparameter used in the `model_init` on the auxiliary processes.

## Solution

To ensure consistent hyperparameter settings across all processes, the following changes were implemented:

1. **Broadcast a FixedTrial object**: After sampling a trial on the main process, we create a `FixedTrial` object on the main process and broadcast it to the other processes.
   
2. **Update `_report_to_hp_search`**: The `_report_to_hp_search` method was modified to not update the study in case we are on a process that got a FixedTrial (the FixedTrial doesn't have a reference to the study object which is necessary to report to Optuna).

## Implementation Details

**In `integration_utils.py`**:

After initializing the trial on the main process, we create a `FixedTrial` object with the same hyperparameters, we serialize it and broadcast it to all other processes. The train method is called with the trial argument.

On the main process we sample a trial (using the `hp_space` function) and create the `FixedTrial` object (lines [240](https://github.com/huggingface/transformers/blob/33868a057c02f0368ba63bd1edb746be38fe3d90/src/transformers/integrations/integration_utils.py#L240)):
```python
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)
```

On the auxiliary processes we retrieve the `FixedTrial` and call `train` by passing the trial as an argument (lines [269](https://github.com/huggingface/transformers/blob/33868a057c02f0368ba63bd1edb746be38fe3d90/src/transformers/integrations/integration_utils.py#L269)):
```python
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)
```

**Other changes**:

The `_hp_search_setup` method now checks if it receives a dictionary of hyperparameters directly, allowing each process to configure itself independently (line [1751](https://github.com/huggingface/transformers/blob/33868a057c02f0368ba63bd1edb746be38fe3d90/src/transformers/trainer.py#L1750)):

```diff
- if not trial.study._is_multi_objective():
+ if hasattr(trial, "study") and not trial.study._is_multi_objective():
```

In the `hp_params` the check for the trial object class is broaden to `BaseTrial` (line [211](https://github.com/huggingface/transformers/blob/33868a057c02f0368ba63bd1edb746be38fe3d90/src/transformers/integrations/integration_utils.py#L211)):
```diff
- if isinstance(trial, optuna.Trial):
+ if isinstance(trial, optuna.trial.BaseTrial):
```

## Additional Notes

This change should ensures consistent hyperparameter application and prevents deadlocks during distributed training.

Let me know whether there might be a more efficient or cleaner way to handle the broadcast and reinitialization across processes. I think Marc Sun worked on the Optuna integration before.




Comments:
Hey,
Thanks for the review.

The tests in `tests/trainer/test_trainer.py` run without issues but I'm having trouble running the deepspeed tests (due to NCCL it seems). 



The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34642). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.