Inconsistent Hyperparameter Application in Distributed DeepSpeed + Optuna Setup for Hyperparameter Search
Author: corentin-ryr
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.47.0.dev0
- Platform: Linux-4.18.0-513.11.1.el8_9.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.14
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- PyTorch version (GPU?): 2.2.1+cu121 (False)

### Who can help?

@SunMarc @ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

This issue can be reproduced by running `trainer.hyperparameter_search` with any DeepSpeed configuration:

```json
{
    "fp16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true
    }
}
```

```python
# region Imports
import logging
from datasets import load_dataset
from transformers import AutoModelForCausalLM
import torch.distributed
from transformers.utils.logging import disable_progress_bar
import torch
from transformers import AutoTokenizer, TrainingArguments
from accelerate.logging import get_logger
from transformers import Trainer
from optuna import Trial

logging.basicConfig(level=logging.INFO)

IS_BF16_AVAILABLE = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
# endregion

# region Trainer parameters

# Compute the batch size per device depending on the number of GPUs
batch_size_for_device_count = {"2": 1, "4": 4, "8": 16}
per_device_batch_size = batch_size_for_device_count[str(torch.cuda.device_count())]
batch_size_for_device_count_eval = {"2": 2, "4": 8, "8": 32}
per_device_batch_size_eval = batch_size_for_device_count_eval[
    str(torch.cuda.device_count())
]

train_args = TrainingArguments(
    per_device_train_batch_size=per_device_batch_size,
    per_device_eval_batch_size=per_device_batch_size_eval,
    dataloader_drop_last=False,
    output_dir="output-hyper-param-search",
    logging_dir="tensorboard_logs/hyper-param-search",
    bf16=IS_BF16_AVAILABLE,
    bf16_full_eval=IS_BF16_AVAILABLE,
    gradient_checkpointing=True,
    num_train_epochs=1,
    save_total_limit=1,
    run_name="hyper-param-search",
    deepspeed="deepspeed_config.json",
    log_level="info",
    logging_steps=5,
    save_steps=100,
    eval_steps=200,
    report_to="tensorboard",
    learning_rate=5e-6,
    warmup_ratio=0.1,
)
# endregion

if torch.distributed.get_rank() != 0:
    verbose = False
    disable_progress_bar()
else:
    verbose = True
logger = get_logger("Trainer")
logger.info(f"Bfloat enabled: {IS_BF16_AVAILABLE}")

# region Tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
tokenizer.truncation_side = "left"
tokenizer.model_max_length = 4096


def collate_fn(batch: list):
    input_ids = []
    labels = []
    for item in batch:
        input_ids.append(item["input_ids"])
        labels.append(item["labels"])
    return {"input_ids": torch.tensor(input_ids), "labels": torch.tensor(labels)}


# endregion


# region Dataset setup
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding=True)


def group_texts(examples):
    examples["labels"] = examples["input_ids"].copy()
    return examples


dataset = (
    load_dataset("rotten_tomatoes", split="validation")
    .map(tokenize_function, batched=True)
    .map(group_texts)
)
dataset_test = (
    load_dataset("rotten_tomatoes", split="validation")
    .map(tokenize_function, batched=True)
    .map(group_texts)
)

logger.info(f"Dataset size:\t {len(dataset)}")
# endregion


def model_init(trial: Trial):
    return AutoModelForCausalLM.from_pretrained("gpt2", torch_dtype=torch.bfloat16)


# region Define the trainer & train
trainer_dict = {
    "tokenizer": tokenizer,
    "args": train_args,
    "train_dataset": dataset,
    "data_collator": collate_fn,
    "eval_dataset": dataset_test,
    "model_init": model_init,
}
trainer = Trainer(**trainer_dict)


def optuna_hp_space(trial: Trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
        "gradient_accumulation_steps": trial.suggest_categorical(
            "gradient_accumulation_steps", [1, 2, 4]
        )
    }


best_trial = trainer.hyperparameter_search(
    backend="optuna", hp_space=optuna_hp_space, n_trials=10
)

logger.info(best_trial)

```

### Expected behavior

## Problem Description

When using DeepSpeed and Optuna for hyperparameter optimization (HPO) in a distributed setting, an inconsistency occurs due to `trainer._hp_search_setup` only executing on the main process. This function applies hyperparameters from each trial and reinitializes the DeepSpeed configurations. However, since the setup is not performed on all processes, only the main process is correctly configured, resulting in parameter mismatches (e.g., `gradient_accumulation_steps`) across processes.

This inconsistency leads to synchronization issues where some processes start backpropagation while others are still in the forward pass or accumulating gradients, causing timeouts or deadlocks.

## Expected Behavior

The trial parameters and DeepSpeed/Accelerate configurations should be synchronized across all processes to ensure consistency in hyperparameter application and prevent process misalignment during training.


Comments:
I proposed a fix in #34642.