Repository: huggingface/transformers
Pull request: Fix scheduler initialization order

# What does this PR do?

This PR put the `lr_scheduler` initialization after the optimizer is prepared to deal with any modification that might happen after. This is mainly to deal with deepspeed as the optimizer is modified and starting with torch 2.10, checks are more strict. 

Fixes https://github.com/huggingface/transformers/pull/43692

Changed files:
filename: src/transformers/integrations/deepspeed.py
patch: @@ -549,8 +549,6 @@ def _lr_scheduler_callable(optimizer):
                 return lr_scheduler
 
             lr_scheduler = DummyScheduler(optimizer, lr_scheduler_callable=_lr_scheduler_callable)
-        else:
-            lr_scheduler = trainer.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)
 
     return optimizer, lr_scheduler
 

filename: src/transformers/trainer.py
patch: @@ -1176,12 +1176,7 @@ def create_optimizer_and_scheduler(self, num_training_steps: int):
         `create_scheduler`) in a subclass.
         """
         self.create_optimizer()
-        if is_sagemaker_mp_enabled() and smp.state.cfg.fp16:
-            # If fp16 is enabled, we unwrap the optimizer
-            optimizer = self.optimizer.optimizer
-        else:
-            optimizer = self.optimizer
-        self.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)
+        self.create_scheduler(num_training_steps=num_training_steps)
 
     def get_decay_parameter_names(self, model) -> list[str]:
         """
@@ -1756,9 +1751,15 @@ def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optim
             num_training_steps (int): The number of training steps to do.
         """
         if self.lr_scheduler is None:
+            if optimizer is None:
+                if is_sagemaker_mp_enabled() and smp.state.cfg.fp16:
+                    # If fp16 is enabled, we unwrap the optimizer
+                    optimizer = self.optimizer.optimizer
+                else:
+                    optimizer = self.optimizer
             self.lr_scheduler = get_scheduler(
                 self.args.lr_scheduler_type,
-                optimizer=self.optimizer if optimizer is None else optimizer,
+                optimizer=optimizer,
                 num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
                 num_training_steps=num_training_steps,
                 scheduler_specific_kwargs=self.args.lr_scheduler_kwargs,
@@ -2287,7 +2288,7 @@ def _inner_training_loop(
             self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)
 
         if not delay_optimizer_creation:
-            self.create_optimizer_and_scheduler(num_training_steps=max_steps)
+            self.create_optimizer()
 
         self.state = TrainerState(
             stateful_callbacks=[
@@ -2322,21 +2323,26 @@ def _inner_training_loop(
                 self._fsdp_qlora_plugin_updates()
                 if self.accelerator.mixed_precision != "fp8":
                     self.model = self.accelerator.prepare(self.model)
-            self.create_optimizer_and_scheduler(num_training_steps=max_steps)
+            self.create_optimizer()
 
         # prepare using `accelerator` prepare
         if use_accelerator_prepare:
             self.model.train()
-            if hasattr(self.lr_scheduler, "step"):
-                model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
+            if self.is_deepspeed_enabled:
+                from accelerate.utils import DummyScheduler
+
+                if isinstance(self.lr_scheduler, DummyScheduler):
+                    model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
+                        self.model, self.optimizer, self.lr_scheduler
+                    )
             else:
-                # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config.
-                model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
-                    self.model, self.optimizer, self.lr_scheduler
-                )
+                model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
         else:
             self.optimizer = self.accelerator.prepare(self.optimizer)
 
+        # Create scheduler now that the optimizer won't change anymore
+        self.create_scheduler(num_training_steps=max_steps)
+
         # since DataLoader was Accelerate prepared w/o a model arg in the same call, we now have to complete the DL wrapping for ALST/UlyssesSP, after model has been prepared
         pc = getattr(self.accelerator, "parallelism_config", None)
         if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled: