File size: 6,316 Bytes
079c32c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
from functools import lru_cache
from typing import Callable, Tuple, List, Any

import numpy as np
import torch

from .default_helper import error_wrapper
from .fake_linklink import FakeLink
from .import_helper import try_import_link


@lru_cache()
def get_link():
    return try_import_link()


@lru_cache()
def is_fake_link():
    return isinstance(get_link(), FakeLink)


def get_rank() -> int:
    """
    Overview:
        Get the rank of ``linklink`` model, return 0 if use ``FakeLink``.

    .. note::
        Reference ``import_helper.try_import_link`` and ``linklink.get_rank``.
    """
    if is_fake_link():
        return 0
    return error_wrapper(get_link().get_rank, 0, "[WARNING]: call linklink error, return default_ret.")()


def get_world_size() -> int:
    """
    Overview:
        Get the ``world_size`` of ``linklink model``, return 0 if use ``FakeLink``.

    .. note::
        Reference ``import_helper.try_import_link`` and ``linklink.get_world_size``.
    """
    if is_fake_link():
        return 1
    return error_wrapper(get_link().get_world_size, 1, "[WARNING]: call linklink error, return default_ret.")()


def broadcast(value: torch.Tensor, rank: int) -> None:
    """
    Overview:
        Use ``linklink.broadcast`` and raise error when using ``FakeLink``
    Arguments:
        - value (:obj:`obj`): the value to board cast
        - rank (:obj:`int`): the rank to broadcast on
    """
    if is_fake_link():
        raise NotImplementedError
    get_link().broadcast(value, rank)


def allreduce(data: torch.Tensor, op: str = 'sum') -> None:
    """
    Overview:
        Call ``linklink.allreduce`` on the data
    Arguments:
        - data (:obj:`obj`): the data to reduce
        - op (:obj:`str`): the operation to perform on data, support ``['sum', 'max']``
    """
    link_op_map = {'sum': get_link().allreduceOp_t.Sum, 'max': get_link().allreduceOp_t.Max}
    if op not in link_op_map.keys():
        raise KeyError("not support allreduce op type: {}".format(op))
    else:
        link_op = link_op_map[op]
    if is_fake_link():
        return data
    get_link().allreduce(data, reduce_op=link_op)
    if op == 'sum':
        data.div_(get_world_size())


def allreduce_async(data: torch.Tensor, op: str = 'sum') -> None:
    """
    Overview:
        Call ``linklink.allreduce_async`` on the data
    Arguments:
        - data (:obj:`obj`): the data to reduce
        - op (:obj:`str`): the operation to perform on data, support ``['sum', 'max']``
    """
    link_op_map = {'sum': get_link().allreduceOp_t.Sum, 'max': get_link().allreduceOp_t.Max}
    if op not in link_op_map.keys():
        raise KeyError("not support allreduce op type: {}".format(op))
    else:
        link_op = link_op_map[op]
    if is_fake_link():
        return data
    if op == 'sum':
        data.div_(get_world_size())
    get_link().allreduce_async(data, reduce_op=link_op)


def get_group(group_size: int) -> List:
    """
    Overview:
        Get the group segmentation of ``group_size`` each group
    Arguments:
        - group_size (:obj:`int`) the ``group_size``
    """
    rank = get_rank()
    world_size = get_world_size()
    if group_size is None:
        group_size = world_size
    assert (world_size % group_size == 0)
    return simple_group_split(world_size, rank, world_size // group_size)


def dist_mode(func: Callable) -> Callable:
    """
    Overview:
        Wrap the function so that in can init and finalize automatically before each call
    Arguments:
        - func (:obj:`Callable`): the function to wrap
    """

    def wrapper(*args, **kwargs):
        dist_init()
        func(*args, **kwargs)
        dist_finalize()

    return wrapper


def dist_init(method: str = 'slurm', device_id: int = 0) -> Tuple[int, int]:
    """
    Overview:
        Init the distribution
    Arguments:
        - method (:obj:`str`): Support ``['slurm', 'single_node`]``
        - device_id (:obj:`int`): Default device when using ``single_node`` method
    """
    get_link().initialize()
    world_size = get_link().get_world_size()
    rank = get_link().get_rank()

    if method == 'slurm':
        # proc_id = int(os.environ['SLURM_PROCID'])
        # ntasks = int(os.environ['SLURM_NTASKS'])
        # node_list = os.environ['SLURM_NODELIST']
        num_gpus = torch.cuda.device_count()
        torch.cuda.set_device(rank % num_gpus)
    elif method == 'single_node':
        torch.cuda.set_device(device_id)

    return rank, world_size


def dist_finalize() -> None:
    """
    Overview:
        Finalize ``linklink``, see ``linklink.finalize()``
    """
    get_link().finalize()


class DistContext:
    """
    Overview:
        A context manager for ``linklink`` distribution
    Interfaces:
        ``__init__``, ``__enter__``, ``__exit__``
    """

    def __init__(self) -> None:
        """
        Overview:
            Initialize the ``DistContext``
        """

        pass

    def __enter__(self) -> None:
        """
        Overview:
            Initialize ``linklink`` distribution
        """

        dist_init()

    def __exit__(self, *args, **kwargs) -> Any:
        """
        Overview:
            Finalize ``linklink`` distribution
        Arugments:
            - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function.
            - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function.
        """

        dist_finalize()


def simple_group_split(world_size: int, rank: int, num_groups: int) -> List:
    """
    Overview:
        Split the group according to ``worldsize``, ``rank`` and ``num_groups``
    Arguments:
        - world_size (:obj:`int`): The world size
        - rank (:obj:`int`): The rank
        - num_groups (:obj:`int`): The number of groups
    .. note::
        With faulty input, raise ``array split does not result in an equal division``
    """

    groups = []
    rank_list = np.split(np.arange(world_size), num_groups)
    rank_list = [list(map(int, x)) for x in rank_list]
    for i in range(num_groups):
        groups.append(get_link().new_group(rank_list[i]))
    group_size = world_size // num_groups
    return groups[rank // group_size]


def synchronize():
    """
    Overview:
        Synchronize the process
    """

    get_link().synchronize()