__init__.py 1.94 KB
Newer Older
Jared Casper's avatar
Jared Casper committed
1
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
2
3
4

"""Model parallel utility interface."""

5
from .initialize import is_unitialized
6
7
8
9
from .initialize import destroy_model_parallel
from .initialize import get_data_parallel_group
from .initialize import get_data_parallel_rank
from .initialize import get_data_parallel_world_size
10
from .initialize import get_embedding_group
Vijay Korthikanti's avatar
Vijay Korthikanti committed
11
from .initialize import get_position_embedding_group
12
from .initialize import get_model_parallel_group
13
14
15
16
17
from .initialize import get_tensor_model_parallel_group
from .initialize import get_pipeline_model_parallel_group
from .initialize import get_tensor_model_parallel_rank, set_tensor_model_parallel_rank
from .initialize import get_pipeline_model_parallel_rank, set_pipeline_model_parallel_rank
from .initialize import is_pipeline_first_stage, is_pipeline_last_stage
18
from .initialize import is_rank_in_embedding_group
Vijay Korthikanti's avatar
Vijay Korthikanti committed
19
from .initialize import is_rank_in_position_embedding_group
20
21
22
from .initialize import is_pipeline_stage_before_split, is_pipeline_stage_after_split
from .initialize import is_pipeline_stage_at_split
from .initialize import get_num_layers
23
from .initialize import get_tensor_model_parallel_src_rank
24
from .initialize import get_data_parallel_src_rank
25
26
from .initialize import get_pipeline_model_parallel_first_rank
from .initialize import get_pipeline_model_parallel_last_rank
27
28
from .initialize import get_pipeline_model_parallel_next_rank
from .initialize import get_pipeline_model_parallel_prev_rank
29
30
from .initialize import get_tensor_model_parallel_world_size, set_tensor_model_parallel_world_size
from .initialize import get_pipeline_model_parallel_world_size, set_pipeline_model_parallel_world_size
31
from .initialize import get_virtual_pipeline_model_parallel_rank, set_virtual_pipeline_model_parallel_rank
32
33
34
35
from .initialize import initialize_model_parallel
from .initialize import model_parallel_is_initialized


36
37
from .utils import divide
from .utils import split_tensor_along_last_dim