30 lines
561 B
Python
30 lines
561 B
Python
from khaosz.parallel.utils import (
|
|
get_world_size,
|
|
get_rank,
|
|
get_device_count,
|
|
get_current_device,
|
|
get_available_backend,
|
|
setup_parallel,
|
|
only_main_procs,
|
|
spawn_parallel_fn
|
|
)
|
|
|
|
from khaosz.parallel.module import (
|
|
RowParallelLinear,
|
|
ColumnParallelLinear
|
|
)
|
|
|
|
__all__ = [
|
|
"get_world_size",
|
|
"get_rank",
|
|
"get_device_count",
|
|
"get_current_device",
|
|
"get_available_backend",
|
|
"setup_parallel",
|
|
"only_main_procs",
|
|
"spawn_parallel_fn",
|
|
|
|
"RowParallelLinear",
|
|
"ColumnParallelLinear"
|
|
]
|