feat: 新增训练后端工厂框架
- BaseTrainingBackend 定义 prepare/accumulate/unwrap_model 抽象 - DDPTrainingBackend 支持全部 DDP 参数并通过 BackendFactory 注册 - unwrap_model 改为实例方法,由子类各自实现
This commit is contained in:
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user