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

"""Model parallel utility interface."""


from .data import broadcast_data

8
from .initialize import is_unitialized
9
10
11
12
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
13
from .initialize import get_embedding_group
Vijay Korthikanti's avatar
Vijay Korthikanti committed
14
from .initialize import get_position_embedding_group
15
from .initialize import get_model_parallel_group
16
17
18
19
20
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
21
from .initialize import is_rank_in_embedding_group
Vijay Korthikanti's avatar
Vijay Korthikanti committed
22
from .initialize import is_rank_in_position_embedding_group
23
24
25
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
26
from .initialize import get_tensor_model_parallel_src_rank
27
from .initialize import get_data_parallel_src_rank
28
29
from .initialize import get_pipeline_model_parallel_first_rank
from .initialize import get_pipeline_model_parallel_last_rank
30
31
from .initialize import get_pipeline_model_parallel_next_rank
from .initialize import get_pipeline_model_parallel_prev_rank
32
33
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
34
from .initialize import get_virtual_pipeline_model_parallel_rank, set_virtual_pipeline_model_parallel_rank
35
36
37
from .initialize import initialize_model_parallel
from .initialize import model_parallel_is_initialized

Vijay Korthikanti's avatar
Vijay Korthikanti committed
38
from .layers import LinearWithGradAccumulationAndAsyncCommunication
39
40
41
from .layers import ColumnParallelLinear
from .layers import RowParallelLinear
from .layers import VocabParallelEmbedding
42
43
from .layers import (set_tensor_model_parallel_attributes,
                     set_defaults_if_not_set_tensor_model_parallel_attributes,
mohammad's avatar
mohammad committed
44
                     copy_tensor_model_parallel_attributes)
45
46
47
 
from .mappings import  copy_to_tensor_model_parallel_region
from .mappings import  reduce_from_tensor_model_parallel_region
48
49
50
from .mappings import  scatter_to_tensor_model_parallel_region
from .mappings import  gather_from_tensor_model_parallel_region
from .mappings import  scatter_to_sequence_parallel_region
51
from .mappings import  gather_from_sequence_parallel_region
52
from .mappings import  reduce_scatter_to_sequence_parallel_region
53
54
55

from .random import checkpoint
from .random import get_cuda_rng_tracker
56
from .random import model_parallel_cuda_manual_seed
57
58
from .random import gather_split_1d_tensor
from .random import split_tensor_into_1d_equal_chunks
Lawrence McAfee's avatar
Lawrence McAfee committed
59
60
61
from .random import make_viewless_tensor
from .random import assert_viewless_tensor
from .random import safely_set_viewless_tensor_data
62

63
64
from .utils import divide
from .utils import split_tensor_along_last_dim