move dir creation to initialize checkpoint

Author: Leahlijuan

src/ml_flashpoint/core/checkpoint_saver.py

@@ -338,10 +338,9 @@ def __setstate__(self, state):
     @override
     @log_execution_time(logger=_LOGGER, name="initialize_checkpoint")
     def initialize_checkpoint(self, checkpoint_id: CheckpointContainerId) -> None:
+        os.makedirs(checkpoint_id.data, exist_ok=True)
         self._create_dirty_checkpoint_marker(checkpoint_id)
-        if self._local_rank_getter() == 0:
-            os.makedirs(checkpoint_id.data, exist_ok=True)
-            _LOGGER.info("Created checkpoint directory: '%s'", checkpoint_id.data)
+        _LOGGER.info("Initialized checkpoint: '%s'", checkpoint_id.data)
 
     @override
     @log_execution_time(logger=_LOGGER, name="stage_data", level=logging.INFO)