Add barriers and fix pyright

llmfoundry/callbacks/hf_checkpointer.py

@@ -308,6 +308,8 @@ def _save_checkpoint(self, state: State, logger: Logger):
                     if isinstance(v, torch.Tensor):
                         state_dict[k] = v.to(dtype=self.dtype)
 
+            new_model_instance = None  # Need this for pyright because variable could be unbound
+
             if dist.get_global_rank() == 0:
                 log.debug('Saving Hugging Face checkpoint in global rank 0')
 
@@ -371,6 +373,9 @@ def _save_checkpoint(self, state: State, logger: Logger):
                             overwrite=self.overwrite,
                         )
 
+            dist.barrier()
+
+            if dist.get_global_rank() == 0:
                 if self.mlflow_registered_model_name and self._is_last_batch(
                         state):
                     components = {'model': new_model_instance}
@@ -435,3 +440,4 @@ def _save_checkpoint(self, state: State, logger: Logger):
                             })
                         process.start()
                         self.child_processes.append(process)
+            dist.barrier()