[bugfix] fix optimizer deepspeed (#9173)

Author: Jintao-Huang

swift/optimizers/base.py

@@ -1,4 +1,5 @@
 from torch.optim import Optimizer
+from transformers.trainer import Trainer as HfTrainer
 from typing import TYPE_CHECKING
 
 try:
@@ -48,7 +49,7 @@ def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
         trainer.scheduler = self.create_scheduler(num_training_steps, trainer.optimizer)
 
     def create_optimizer(self) -> Optimizer:
-        return self.trainer.create_optimizer()
+        return HfTrainer.create_optimizer(self.trainer)
 
     def create_scheduler(self, num_training_steps: int, optimizer: Optimizer) -> LRScheduler:
-        return self.trainer.create_scheduler(num_training_steps, optimizer)
+        return HfTrainer.create_scheduler(self.trainer, num_training_steps, optimizer)

swift/trainers/mixin.py

@@ -77,6 +77,7 @@ def __init__(self,
         self.padding_free = self.template.padding_free
         self.task_type = self.template.task_type
         self.problem_type = getattr(model.config, 'problem_type', None)
+        self.optimizer_callback = optimizers_map[args.optimizer or 'default'](args, self)
         if args.check_model and hasattr(model, 'model_dir'):
             with ms_logger_context(logging.CRITICAL), patch_modelscope_hub_timeout():
                 config_info = self._collect_config_info()
@@ -983,8 +984,17 @@ def create_loss_and_eval_metric(self, args):
         return res
 
     def create_optimizer_and_scheduler(self, num_training_steps: int):
-        optimizer_callback: OptimizerCallback = optimizers_map[self.args.optimizer or 'default'](self.args, self)
-        optimizer_callback.create_optimizer_and_scheduler(num_training_steps)
+        self.optimizer_callback.create_optimizer_and_scheduler(num_training_steps)
+
+    def create_optimizer(self):
+        self.optimizer = self.optimizer_callback.create_optimizer()
+        if self.optimizer is not None:
+            self.optimizer.param_groups = [pg for pg in self.optimizer.param_groups if len(pg['params']) > 0]
+        return self.optimizer
+
+    def create_scheduler(self, num_training_steps: int, optimizer=None):
+        self.lr_scheduler = self.optimizer_callback.create_scheduler(num_training_steps, optimizer)
+        return self.lr_scheduler
 
     @staticmethod
     def _get_listwise_reranker_preds(logits, labels):