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

scripts/tools/train.py

@@ -8,7 +8,6 @@ import torch.optim as optim
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory
 from astrai.model import AutoRegressiveLM
-from astrai.serialization import Checkpoint
 from astrai.trainer import SchedulerFactory, Trainer
 
 
@@ -166,6 +165,10 @@ def parse_args() -> argparse.Namespace:
     return args
 
 
+def create_model(config):
+    return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
+
+
 def create_optimizer(model, **kwargs) -> optim.Optimizer:
     return optim.AdamW(model.parameters(), fused=True, **kwargs)
 
@@ -238,15 +241,6 @@ def train(
     if window_size is None:
         window_size = config.max_len
 
-    # Create model and load full checkpoint (state_dict + optimizer + scheduler + meta)
-    checkpoint = Checkpoint.load(param_path)
-    model = AutoRegressiveLM(config).to(dtype=torch.bfloat16)
-    model.load_state_dict(checkpoint.state_dict, strict=False)
-
-    # Strip state_dict to avoid pickling ~7GB through mp.spawn pipe
-    # (model weights already loaded into model above)
-    checkpoint.state_dict = {}
-
     strategy_kwargs = {
         "beta": dpo_beta,
         "label_smoothing": label_smoothing,
@@ -261,6 +255,7 @@ def train(
         "broadcast_buffers": False,
     }
 
+    model_fn = partial(create_model, config)
     dataset = DatasetFactory.load(
         train_type=train_type,
         load_path=data_root_path,
@@ -292,7 +287,7 @@ def train(
     )
 
     train_config = TrainConfig(
-        model=model,
+        model_fn=model_fn,
         strategy=train_type,
         dataset=dataset,
         optimizer_fn=optimizer_fn,
@@ -317,7 +312,7 @@ def train(
     )
 
     trainer = Trainer(train_config)
-    trainer.train(checkpoint=checkpoint)
+    trainer.train(resume_dir=param_path)
 
 
 if __name__ == "__main__":