fix ds optimizer

hiyouga · Mar 26, 2024 · 3bcd41b · 3bcd41b · hiyouga · Mar 26, 2024
1 parent b29d556
commit 3bcd41b
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/src/llmtuner/train/dpo/trainer.py b/src/llmtuner/train/dpo/trainer.py
@@ -64,10 +64,10 @@ def __init__(
                 self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
 
     def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
-        self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
         if self.optimizer is None:
-            self.create_optimizer()
+            self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
 
+        self.create_optimizer()
         self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
 
     def sft_loss(self, chosen_logits: torch.FloatTensor, chosen_labels: torch.LongTensor) -> torch.Tensor:

diff --git a/src/llmtuner/train/pt/trainer.py b/src/llmtuner/train/pt/trainer.py
@@ -23,8 +23,8 @@ def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
         self.finetuning_args = finetuning_args
 
     def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
-        self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
         if self.optimizer is None:
-            self.create_optimizer()
+            self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
 
+        self.create_optimizer()
         self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
diff --git a/src/llmtuner/train/rm/trainer.py b/src/llmtuner/train/rm/trainer.py
@@ -30,10 +30,10 @@ def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
         self.can_return_loss = True  # override property to return eval_loss
 
     def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
-        self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
         if self.optimizer is None:
-            self.create_optimizer()
+            self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
 
+        self.create_optimizer()
         self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
 
     def compute_loss(

diff --git a/src/llmtuner/train/sft/trainer.py b/src/llmtuner/train/sft/trainer.py
@@ -30,10 +30,10 @@ def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
         self.finetuning_args = finetuning_args
 
     def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
-        self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
         if self.optimizer is None:
-            self.create_optimizer()
+            self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
 
+        self.create_optimizer()
         self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
 
     def prediction_step(