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/parallel/executor.py

@@ -7,6 +7,7 @@ from contextlib import contextmanager
 from typing import Optional, Tuple
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.fsdp import FullStateDictConfig, StateDictType
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
@@ -120,6 +121,21 @@ class BaseExecutor:
     def unwrap_model(self, model: nn.Module):
         return model.state_dict()
 
+    @contextmanager
+    def checkpoint_context(self, model: nn.Module):
+        if self.use_distributed:
+            dist.barrier()
+        state_dict = self._gather_state_dict(model)
+        yield state_dict
+        if self.use_distributed:
+            dist.barrier()
+
+    def _gather_state_dict(self, model: nn.Module):
+        state_dict = self.unwrap_model(model)
+        if self.use_distributed and get_rank() != 0:
+            return None
+        return state_dict
+
     @property
     def use_distributed(self) -> bool:
         return get_world_size() > 1
@@ -208,6 +224,13 @@ class DDPExecutor(BaseExecutor):
             return model.module.state_dict()
         return model.state_dict()
 
+    def _gather_state_dict(self, model: nn.Module):
+        if not self.use_distributed:
+            return self.unwrap_model(model)
+        if get_rank() != 0:
+            return None
+        return self.unwrap_model(model)
+
 
 @ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
@@ -279,7 +302,7 @@ class FSDPExecutor(BaseExecutor):
             with FSDP.state_dict_type(
                 model,
                 StateDictType.FULL_STATE_DICT,
-                FullStateDictConfig(offload_to_cpu=True, rank0_only=False),
+                FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
             ):
                 return model.state_dict()
 

astrai/serialization/checkpoint.py

@@ -148,9 +148,6 @@ class Checkpoint:
         save_path = Path(save_dir)
         save_path.mkdir(parents=True, exist_ok=True)
 
-        if get_rank() != 0:
-            return
-
         meta = {
             "epoch": self.epoch,
             "consumed_samples": self.consumed_samples,

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: