feat : add distributed checkpoint via executor checkpoint_context

- Add checkpoint_context context manager to BaseExecutor with entry/exit barrier
- Add _gather_state_dict hook overridden per executor (template method)
- DDPExecutor skips unwrap on non-rank-0 to avoid redundant state_dict gather
- FSDPExecutor uses rank0_only=True to reduce memory on non-writers
- Remove redundant rank-0 guard from Checkpoint.save and manual barrier from Callback

astrai/trainer/train_callback.py

@@ -14,7 +14,7 @@ from tqdm import tqdm
 
 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
-from astrai.parallel.setup import get_current_device, get_rank
+from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_norm,
@@ -142,25 +142,25 @@ class CheckpointCallback(TrainCallback):
         self.last_ckpt_step = 0
 
     def _save_checkpoint(self, context: TrainContext):
-        state_dict = context.executor.unwrap_model(context.model)
         self.last_ckpt_step = context.optimizer_step
 
-        if get_rank() == 0:
-            save_path = os.path.join(
-                self.save_dir,
-                f"epoch_{context.epoch}_step_{context.optimizer_step}",
-            )
-            extra = self.save_extra_fn(context)
-            meta = context.config.to_dict()
-            context.checkpoint = Checkpoint(
-                state_dict=state_dict,
-                epoch=context.epoch,
-                consumed_samples=context.consumed_samples,
-                config=context.model_config,
-                extra=extra,
-                meta=meta,
-            )
-            context.checkpoint.save(save_path)
+        with context.executor.checkpoint_context(context.model) as state_dict:
+            if state_dict is not None:
+                save_path = os.path.join(
+                    self.save_dir,
+                    f"epoch_{context.epoch}_step_{context.optimizer_step}",
+                )
+                extra = self.save_extra_fn(context)
+                meta = context.config.to_dict()
+                context.checkpoint = Checkpoint(
+                    state_dict=state_dict,
+                    epoch=context.epoch,
+                    consumed_samples=context.consumed_samples,
+                    config=context.model_config,
+                    extra=extra,
+                    meta=meta,
+                )
+                context.checkpoint.save(save_path)
 
     def on_batch_end(self, context: TrainContext):
         if context.optimizer_step - self.last_ckpt_step >= self.interval: