fix deadlock; make style

finetrainers/trainer.py

@@ -423,10 +423,15 @@ class Trainer:
             if not self.state.accelerator.distributed_type == DistributedType.DEEPSPEED:
                 while len(models) > 0:
                     model = models.pop()
-                    if isinstance(unwrap_model(self.state.accelerator, model), type(unwrap_model(self.state.accelerator, self.transformer))):
+                    if isinstance(
+                        unwrap_model(self.state.accelerator, model),
+                        type(unwrap_model(self.state.accelerator, self.transformer)),
+                    ):
                         transformer_ = unwrap_model(self.state.accelerator, model)
                     else:
-                        raise ValueError(f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}")
+                        raise ValueError(
+                            f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}"
+                        )
             else:
                 transformer_ = unwrap_model(self.state.accelerator, self.transformer).__class__.from_pretrained(
                     self.args.pretrained_model_name_or_path, subfolder="transformer"
@@ -561,10 +566,15 @@ class Trainer:
         initial_global_step = 0
 
         # Potentially load in the weights and states from a previous save
-        resume_from_checkpoint_path, initial_global_step, global_step, first_epoch = get_latest_ckpt_path_to_resume_from(
+        (
+            resume_from_checkpoint_path,
+            initial_global_step,
+            global_step,
+            first_epoch,
+        ) = get_latest_ckpt_path_to_resume_from(
             resume_from_checkpoint=self.args.resume_from_checkpoint,
             num_update_steps_per_epoch=self.state.num_update_steps_per_epoch,
-            output_dir=self.args.output_dir
+            output_dir=self.args.output_dir,
         )
         if resume_from_checkpoint_path:
             self.state.accelerator.load_state(resume_from_checkpoint_path)
@@ -701,7 +711,9 @@ class Trainer:
                     if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
                         if global_step % self.args.checkpointing_steps == 0:
                             save_path = get_intermediate_ckpt_path(
-                                checkpointing_limit=self.args.checkpointing_limit, step=global_step, output_dir=self.args.output_dir
+                                checkpointing_limit=self.args.checkpointing_limit,
+                                step=global_step,
+                                output_dir=self.args.output_dir,
                             )
                             accelerator.save_state(save_path)
 
@@ -750,7 +762,9 @@ class Trainer:
                 transformer_lora_layers=transformer_lora_layers,
             )
 
-            self.validate(step=global_step, final_validation=True)
+        self.validate(step=global_step, final_validation=True)
+
+        if accelerator.is_main_process:
             if self.args.push_to_hub:
                 upload_folder(
                     repo_id=self.state.repo_id, folder_path=self.args.output_dir, ignore_patterns=["checkpoint-*"]

finetrainers/utils/checkpointing.py

@@ -7,6 +7,7 @@ from ..utils.file_utils import find_files, delete_files
 logger = get_logger("finetrainers")
 logger.setLevel(FINETRAINERS_LOG_LEVEL)
 
+
 def get_latest_ckpt_path_to_resume_from(
     resume_from_checkpoint: str, num_update_steps_per_epoch: int, output_dir: str
 ) -> Tuple[str, int, int, int]:
@@ -26,9 +27,7 @@ def get_latest_ckpt_path_to_resume_from(
             path = dirs[-1] if len(dirs) > 0 else None
 
         if path is None:
-            logger.info(
+            logger.info(f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run.")
-                f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run."
-            )
             resume_from_checkpoint = None
             initial_global_step = 0
             global_step = 0