ViperEkura
738cb8f128
- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks - Fix create_ref_model returning None on non-rank-0 under FSDP - Fix sync_old_model only updating old_model on rank-0 under FSDP - Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest - Add distributed tests for broadcast_state_dict and create_ref_model
40 lines
762 B
Python
40 lines
762 B
Python
from astrai.parallel.executor import (
|
|
AccumOptimizer,
|
|
AccumScheduler,
|
|
BaseExecutor,
|
|
DDPExecutor,
|
|
ExecutorFactory,
|
|
FSDPExecutor,
|
|
GradientState,
|
|
NoneExecutor,
|
|
broadcast_state_dict,
|
|
create_ref_model,
|
|
)
|
|
from astrai.parallel.setup import (
|
|
get_current_device,
|
|
get_rank,
|
|
get_world_size,
|
|
only_on_rank,
|
|
setup_parallel,
|
|
spawn_parallel_fn,
|
|
)
|
|
|
|
__all__ = [
|
|
"get_world_size",
|
|
"get_rank",
|
|
"get_current_device",
|
|
"only_on_rank",
|
|
"setup_parallel",
|
|
"spawn_parallel_fn",
|
|
"ExecutorFactory",
|
|
"BaseExecutor",
|
|
"GradientState",
|
|
"AccumOptimizer",
|
|
"AccumScheduler",
|
|
"NoneExecutor",
|
|
"DDPExecutor",
|
|
"FSDPExecutor",
|
|
"create_ref_model",
|
|
"broadcast_state_dict",
|
|
]
|