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 扫描构建产物

tests/trainer/conftest.py

@@ -27,7 +27,7 @@ class TrainerDataset(Dataset):
 
 
 def create_train_config(
-    model: torch.nn.Module,
+    model_fn,
     dataset: Dataset,
     test_dir: str,
     device: str,
@@ -43,7 +43,7 @@ def create_train_config(
     """Factory function to create common TrainConfig for tests.
 
     Args:
-        model: The model to train
+        model_fn: Model factory (callable returning nn.Module)
         dataset: Training dataset
         test_dir: Checkpoint directory
         device: Device type ("cuda" or "cpu")
@@ -70,7 +70,7 @@ def create_train_config(
 
     return TrainConfig(
         strategy=strategy,
-        model=model,
+        model_fn=model_fn,
         dataset=dataset,
         optimizer_fn=optimizer_fn,
         scheduler_fn=scheduler_fn,

tests/trainer/test_callbacks.py

@@ -106,7 +106,7 @@ def test_gradient_checkpointing_trainer_integration(base_test_env, random_datase
         )
 
     train_config = TrainConfig(
-        model=base_test_env["model"],
+        model_fn=lambda: base_test_env["model"],
         strategy="seq",
         dataset=random_dataset,
         optimizer_fn=optimizer_fn,
@@ -140,7 +140,7 @@ def test_callback_integration(base_test_env, random_dataset):
         )
 
     train_config = TrainConfig(
-        model=base_test_env["model"],
+        model_fn=lambda: base_test_env["model"],
         strategy="seq",
         dataset=random_dataset,
         optimizer_fn=optimizer_fn,

tests/trainer/test_early_stopping.py

@@ -4,7 +4,6 @@ import numpy as np
 import torch
 
 from astrai.config.train_config import TrainConfig
-from astrai.serialization import Checkpoint
 from astrai.trainer.schedule import SchedulerFactory
 from astrai.trainer.trainer import Trainer
 
@@ -24,7 +23,7 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
         strategy="seq",
         optimizer_fn=optimizer_fn,
         scheduler_fn=scheduler_fn,
-        model=base_test_env["model"],
+        model_fn=lambda: base_test_env["model"],
         dataset=early_stopping_dataset,
         ckpt_dir=base_test_env["test_dir"],
         log_dir=os.path.join(base_test_env["test_dir"], "logs"),
@@ -39,17 +38,20 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
     trainer = Trainer(train_config)
 
     # Should handle early stopping gracefully
-    checkpoint = None
     try:
-        checkpoint = trainer.train()
+        trainer.train()
     except Exception:
-        # Handle any exceptions
         pass
 
+    # Resume from latest checkpoint
     load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2")
-    checkpoint = Checkpoint.load(load_dir)
-    trainer.train(checkpoint)
+    trainer = Trainer(train_config)
+    trainer.train(resume_dir=load_dir)
 
+    # Verify checkpoint was saved at expected iteration
     load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10")
-    checkpoint = Checkpoint.load(load_dir)
-    assert checkpoint.iteration == 10
+    import json
+
+    with open(os.path.join(load_dir, "meta.json")) as f:
+        meta = json.load(f)
+    assert meta["iteration"] == 10

tests/trainer/test_trainer.py

@@ -9,7 +9,7 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
 
     for batch_per_device in batch_sizes:
         train_config = train_config_factory(
-            model=base_test_env["model"],
+            model_fn=lambda: base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
@@ -25,7 +25,7 @@ def test_gradient_accumulation(base_test_env, random_dataset, train_config_facto
 
     for grad_accum_steps in grad_accum_steps_list:
         train_config = train_config_factory(
-            model=base_test_env["model"],
+            model_fn=lambda: base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
@@ -50,7 +50,7 @@ def test_memory_efficient_training(base_test_env, random_dataset, train_config_f
 
     for config in small_batch_configs:
         train_config = train_config_factory(
-            model=base_test_env["model"],
+            model_fn=lambda: base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],