fix HPO DDP GPU problem (#19168)

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

src/transformers/integrations.py

@@ -23,6 +23,7 @@ import pickle
 import shutil
 import sys
 import tempfile
+from dataclasses import asdict
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, Optional
 
@@ -195,9 +196,10 @@ def run_hp_search_optuna(trainer, n_trials: int, direction: str, **kwargs) -> Be
             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)
-            local_rank = trainer.args.local_rank  # backup the local_rank info
-            trainer.args = pickle.loads(bytes(args_main_rank))
-            trainer.args.local_rank = local_rank
+            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:
@@ -429,9 +431,10 @@ def run_hp_search_sigopt(trainer, n_trials: int, direction: str, **kwargs) -> Be
             if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED:
                 raise RuntimeError("only support DDP Sigopt HPO for ParallelMode.DISTRIBUTED currently.")
             torch.distributed.broadcast_object_list(args_main_rank, src=0)
-            local_rank = trainer.args.local_rank  # backup the local_rank info
-            trainer.args = pickle.loads(bytes(args_main_rank))
-            trainer.args.local_rank = local_rank
+            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:
@@ -470,7 +473,6 @@ def run_hp_search_wandb(trainer, n_trials: int, direction: str, **kwargs) -> Bes
         sweep_config["name"] = name
 
     def _objective():
-
         run = wandb.run if wandb.run else wandb.init()
         trainer.state.trial_name = run.name
         run.config.update({"assignments": {}, "metric": metric})