File size: 1,282 Bytes
74da989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | # SPDX-License-Identifier: Apache-2.0
from distributed.communication_op import *
from distributed.parallel_state import (
cleanup_dist_env_and_memory,
get_dp_group,
get_dp_rank,
get_dp_world_size,
get_local_torch_device,
get_sp_group,
get_sp_parallel_rank,
get_sp_world_size,
get_tp_group,
get_tp_rank,
get_tp_world_size,
get_world_group,
get_world_rank,
get_world_size,
init_distributed_environment,
initialize_model_parallel,
maybe_init_distributed_environment_and_model_parallel,
model_parallel_is_initialized,
)
from distributed.utils import *
__all__ = [
# Initialization
"init_distributed_environment",
"initialize_model_parallel",
"cleanup_dist_env_and_memory",
"model_parallel_is_initialized",
"maybe_init_distributed_environment_and_model_parallel",
# World group
"get_world_group",
"get_world_rank",
"get_world_size",
# Data parallel group
"get_dp_group",
"get_dp_rank",
"get_dp_world_size",
# Sequence parallel group
"get_sp_group",
"get_sp_parallel_rank",
"get_sp_world_size",
# Tensor parallel group
"get_tp_group",
"get_tp_rank",
"get_tp_world_size",
# Get torch device
"get_local_torch_device",
]
|