feat: 新增训练后端工厂框架

- BaseTrainingBackend 定义 prepare/accumulate/unwrap_model 抽象
- DDPTrainingBackend 支持全部 DDP 参数并通过 BackendFactory 注册
- unwrap_model 改为实例方法,由子类各自实现
This commit is contained in:
2026-05-24 15:15:14 +08:00
parent 0594ce1017
commit 8cbf3f36e2
2 changed files with 193 additions and 0 deletions
+10
View File
@@ -1,3 +1,9 @@
from astrai.parallel.backend import (
AccumOptimizer,
AccumScheduler,
BackendFactory,
BaseTrainingBackend,
)
from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
from astrai.parallel.setup import (
get_current_device,
@@ -17,4 +23,8 @@ __all__ = [
"spawn_parallel_fn",
"RowParallelLinear",
"ColumnParallelLinear",
"BackendFactory",
"BaseTrainingBackend",
"AccumOptimizer",
"AccumScheduler",
]