From b092316385c253ead8f0f308467db0258142e361 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 13 Jul 2026 12:27:09 +0800
Subject: [PATCH] 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        | 25 ++++++++++++++++++++-
 astrai/serialization/checkpoint.py |  3 ---
 astrai/trainer/train_callback.py   | 36 +++++++++++++++---------------
 3 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py
index 571365d..2d43a73 100644
--- a/astrai/parallel/executor.py
+++ b/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()
 
diff --git a/astrai/serialization/checkpoint.py b/astrai/serialization/checkpoint.py
index 0079f35..37ef04c 100644
--- a/astrai/serialization/checkpoint.py
+++ b/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,
diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py
index ec6156f..828194f 100644
--- a/astrai/trainer/train_callback.py
+++ b/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: