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):

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):

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):

- 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):

- 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.