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 扫描构建产物
This commit is contained in:
2026-05-27 20:15:29 +08:00
parent 34c6c45bd6
commit 4145d35e3c
10 changed files with 170 additions and 116 deletions
+3 -3
View File
@@ -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,