refactor: 检查点加载重构，路径替代对象传递

- model: nn.Module -> model_fn 工厂函数，spawn 边界只传字符串
- Trainer.train(resume_dir=path) — Checkpoint 不再通过 pickle 传递
- TrainContextBuilder.with_resume_dir(path) — 自动检测 meta.json 分流 resume/from-scratch
- CheckpointCallback: 拆分 state_dict 收集（全 rank）与磁盘写入（rank-0），修复 FSDP 死锁
- serialization: load_torch 支持 broadcast，消除 _load_extra/_load_torch_broadcast
- optimizer/scheduler 恢复逻辑内联到 build()，在 executor.prepare() 之后执行
- pyproject.toml: ruff exclude build/ 避免 CI 扫描构建产物

astrai/trainer/train_callback.py

@@ -15,7 +15,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
+from astrai.parallel.setup import get_current_device, get_rank
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_max,
@@ -139,42 +139,36 @@ class CheckpointCallback(TrainCallback):
         weight_only: bool = False,
         state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
         save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
-        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
     ):
         self.save_dir = save_dir
         self.interval = interval
         self.weight_only = weight_only
         self.state_dict_fn = state_dict_fn
         self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
-        self.load_extra_fn = load_extra_fn or CheckpointCallback.load_extra
         self.last_ckpt_iter = 0
 
-    @only_on_rank(0)
     def _save_checkpoint(self, context: TrainContext):
-        save_path = os.path.join(
-            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
-        )
+        # All ranks gather state_dict — collective for FSDP, local for DDP
         state_dict = (
             self.state_dict_fn(context.model)
             if self.state_dict_fn
             else context.model.state_dict()
         )
-
-        extra = self.save_extra_fn(context)
-        context.checkpoint = Checkpoint(
-            state_dict=state_dict,
-            epoch=context.epoch,
-            iteration=context.iteration,
-            extra=extra,
-            meta=context.config.to_dict(),
-        )
-
-        context.checkpoint.save(save_path)
         self.last_ckpt_iter = context.iteration
 
-    def on_train_begin(self, context: TrainContext):
-        if context.checkpoint and context.checkpoint.extra:
-            self.load_extra_fn(context.checkpoint.extra, context)
+        if get_rank() == 0:
+            save_path = os.path.join(
+                self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
+            )
+            extra = self.save_extra_fn(context)
+            context.checkpoint = Checkpoint(
+                state_dict=state_dict,
+                epoch=context.epoch,
+                iteration=context.iteration,
+                extra=extra,
+                meta=context.config.to_dict(),
+            )
+            context.checkpoint.save(save_path)
 
     def on_batch_end(self, context: TrainContext):
         if context.iteration - self.last_ckpt_iter >= self.interval:
@@ -196,12 +190,6 @@ class CheckpointCallback(TrainCallback):
                 extra[name] = obj.state_dict()
         return extra
 
-    @staticmethod
-    def load_extra(extra: dict, context: TrainContext):
-        for name in CheckpointCallback.extra_keys:
-            if name in extra:
-                getattr(context, name).load_state_dict(extra[name])
-
 
 @CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):