refactor: 重构训练后端为 Executor 模式

- backend.py → executor.py,BaseTrainingBackend → BaseExecutor
- 新增 NoneExecutor(单卡)和 DDPExecutor(DDP,world_size=1 自动降级)
- 新增 GradientState 分离梯度同步状态,AccumOptimizer/AccumScheduler 包裹拦截
- 新增 astrai/protocols.py:OptimizerProtocol/SchedulerProtocol 结构子类型
- TrainContext.backend → executor,TrainConfig 移除 parallel_wrapper/state_dict_fn,新增 parallel_mode/executor_kwargs
- 训练循环用 accumulate() 包裹,on_optimizer_step 命名约定=gate
- scripts/tools/train.py 移除 ddp_wrap/prepare_checkpoint,新增 --parallel_mode
This commit is contained in:
2026-05-24 20:35:44 +08:00
parent 8cbf3f36e2
commit 3ab4f237e5
8 changed files with 153 additions and 104 deletions
+9 -5
View File
@@ -1,8 +1,10 @@
from astrai.parallel.backend import (
from astrai.parallel.executor import (
AccumOptimizer,
AccumScheduler,
BackendFactory,
BaseTrainingBackend,
BaseExecutor,
ExecutorFactory,
GradientState,
NoneExecutor,
)
from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
from astrai.parallel.setup import (
@@ -23,8 +25,10 @@ __all__ = [
"spawn_parallel_fn",
"RowParallelLinear",
"ColumnParallelLinear",
"BackendFactory",
"BaseTrainingBackend",
"ExecutorFactory",
"BaseExecutor",
"GradientState",
"AccumOptimizer",
"AccumScheduler",
"NoneExecutor",
]