fix: 修复一些运行时问题

khaosz/trainer/schedule.py

@@ -151,8 +151,8 @@ class SchedulerFactory:
     """
 
     @staticmethod
-    def load(optimizer, scedule_config: ScheduleConfig) -> BaseScheduler:
-        kwargs = scedule_config.get_kwargs()
+    def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
+        kwargs = schedule_config.get_kwargs()
         schedule_type = kwargs.pop("schedule_type")
         
         if schedule_type == "cosine":

khaosz/trainer/train_callback.py

@@ -108,10 +108,10 @@ class CheckpointCallback(TrainCallback):
     
     def _save_checkpoint(self, context: 'TrainContext'):
         save_path = os.path.join(self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}")
-        state_dict = self.state_dict_fn(context.model) if self.state_dict_fn else context.optimizer.state_dict()
+        state_dict = self.state_dict_fn(context.model) if self.state_dict_fn else context.model.state_dict()
+        
         context.checkpoint = Checkpoint(
-            optimizer_state_dict=state_dict,
-            scheduler_state_dict=context.scheduler.state_dict() if context.scheduler else None,
+            state_dict=state_dict,
             epoch=context.epoch,
             iteration=context.iteration
         )

khaosz/trainer/train_context.py

@@ -53,15 +53,13 @@ class TrainContextBuilder:
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
         if checkpoint is None:
             checkpoint = Checkpoint(
-                optimizer_state_dict=self._context.optimizer.state_dict(),
-                scheduler_state_dict=self._context.scheduler.state_dict(),
+                state_dict=self._context.model.state_dict(),
             )
         else:
             # resume from the assigned checkpoint or assigned iteration
             self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
             self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
-            self._context.optimizer.load_state_dict(checkpoint.optimizer_state_dict)
-            self._context.scheduler.load_state_dict(checkpoint.scheduler_state_dict)
+            self._context.model.load_state_dict(checkpoint.state_dict)
         
         self._context.checkpoint = checkpoint
         return self