fix : 并行训练 state_dict 收集与训练/推理并发缺陷

- FSDPExecutor: unwrap_model 返回全量 state_dict (state_dict_type FULL);use_orig_params=True
- DDPExecutor/BaseExecutor: unwrap_model 统一返回 model.module.state_dict() / model.state_dict()
- CheckpointCallback: 走 executor.unwrap_model 拿完整 state_dict
- strategy.py: 移除 FSDP/DDp 依赖;create_ref_model(model_fn, state_dict) 纯函数
- TrainContextBuilder: 传递 model_fn + executor 到 strategy
- GRPOStrategy.sync_ref_model: 通过 executor.unwrap_model 获取完整权重
- TaskManager.wait_for_tasks: 锁内检查队列,消除 clear/set 竞态
- ProtocolHandler: stop token 不再计入 completion_tokens(流式/非流式)
This commit is contained in:
2026-05-29 21:12:52 +08:00
parent a3275423a4
commit d4451f6afb
7 changed files with 39 additions and 41 deletions
-2
View File
@@ -1,4 +1,3 @@
import json
import os
import numpy as np
@@ -8,7 +7,6 @@ import torch
from astrai.dataset.dataset import DatasetFactory, SEQDataset
from astrai.dataset.storage import (
H5Store,
MmapStore,
StoreFactory,
detect_format,
load_bin,