-
Notifications
You must be signed in to change notification settings - Fork 21.3k
/
distributed_c10d.py
3605 lines (3053 loc) · 140 KB
/
distributed_c10d.py
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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
import itertools
import collections.abc
import contextlib
import io
import logging
import os
import pickle
import time
import warnings
from collections import namedtuple
from datetime import timedelta
from typing import Any, Dict, Optional, Tuple, Union
import torch
from torch._C._distributed_c10d import (
AllreduceCoalescedOptions,
AllreduceOptions,
AllToAllOptions,
_DistributedBackendOptions,
BarrierOptions,
BroadcastOptions,
GatherOptions,
PrefixStore,
ProcessGroup,
ReduceOp,
ReduceOptions,
ReduceScatterOptions,
ScatterOptions,
Store,
DebugLevel,
get_debug_level,
Work
)
from torch._six import string_classes
from .constants import default_pg_timeout
from .rendezvous import register_rendezvous_handler, rendezvous # noqa: F401
__all__ = [
'Backend', 'GroupMember', 'P2POp', 'all_gather', 'all_gather_coalesced',
'all_gather_multigpu', 'all_gather_object', 'all_reduce',
'all_reduce_coalesced', 'all_reduce_multigpu', 'all_to_all',
'all_to_all_single', 'barrier', 'batch_isend_irecv', 'broadcast',
'broadcast_multigpu', 'broadcast_object_list', 'destroy_process_group',
'dist_backend', 'gather', 'gather_object', 'get_backend', 'get_rank',
'get_world_size', 'group', 'init_process_group', 'irecv',
'is_gloo_available', 'is_initialized', 'is_mpi_available',
'is_nccl_available', 'is_torchelastic_launched', 'is_ucc_available',
'isend', 'monitored_barrier', 'new_group', 'new_subgroups',
'new_subgroups_by_enumeration', 'recv', 'reduce', 'reduce_multigpu',
'reduce_scatter', 'reduce_scatter_multigpu', 'scatter',
'scatter_object_list', 'send', 'supports_complex',
'AllreduceCoalescedOptions', 'AllreduceOptions', 'AllToAllOptions',
'BarrierOptions', 'BroadcastOptions', 'GatherOptions', 'PrefixStore',
'ProcessGroup', 'ReduceOp', 'ReduceOptions', 'ReduceScatterOptions',
'ScatterOptions', 'Store', 'DebugLevel', 'get_debug_level', 'Work',
'default_pg_timeout', 'get_group_rank', 'get_global_rank', 'get_process_group_ranks',
'reduce_op', 'all_gather_into_tensor', 'reduce_scatter_tensor',
]
_MPI_AVAILABLE = True
_NCCL_AVAILABLE = True
_GLOO_AVAILABLE = True
_UCC_AVAILABLE = True
_pickler = pickle.Pickler
_unpickler = pickle.Unpickler
# Change __module__ of all imported types from torch._C._distributed_c10d that are public
def _export_c_types():
_public_types_to_change_module = [
AllreduceCoalescedOptions,
AllreduceOptions,
AllToAllOptions,
BarrierOptions,
BroadcastOptions,
GatherOptions,
PrefixStore,
ProcessGroup,
ReduceOp,
ReduceOptions,
ReduceScatterOptions,
ScatterOptions,
Store,
DebugLevel,
get_debug_level,
Work
]
for type in _public_types_to_change_module:
type.__module__ = "torch.distributed.distributed_c10d"
_export_c_types()
try:
from torch._C._distributed_c10d import ProcessGroupMPI
ProcessGroupMPI.__module__ = "torch.distributed.distributed_c10d"
__all__ += ["ProcessGroupMPI"]
except ImportError:
_MPI_AVAILABLE = False
try:
from torch._C._distributed_c10d import ProcessGroupNCCL
ProcessGroupNCCL.__module__ = "torch.distributed.distributed_c10d"
__all__ += ["ProcessGroupNCCL"]
except ImportError:
_NCCL_AVAILABLE = False
try:
from torch._C._distributed_c10d import ProcessGroupGloo
from torch._C._distributed_c10d import _ProcessGroupWrapper
ProcessGroupGloo.__module__ = "torch.distributed.distributed_c10d"
__all__ += ["ProcessGroupGloo"]
except ImportError:
_GLOO_AVAILABLE = False
try:
from torch._C._distributed_c10d import ProcessGroupUCC
ProcessGroupUCC.__module__ = "torch.distributed.distributed_c10d"
__all__ += ["ProcessGroupUCC"]
except ImportError:
_UCC_AVAILABLE = False
logger = logging.getLogger(__name__)
PG_WRAPPER_STORE_PREFIX = "pg_wrapper"
# Some reduce ops are not supported by complex numbers and will result in an error.
# We currently provide complex support to the distributed API by viewing
# complex tensors as real (torch.view_as_real), meaning that calling
# these unsupported ops will return garbage values rather than error out.
# (e.g. max(2+3i, 3+2i) = 3+3i)
# We'd like calls to unsupported ops to error out accordingly,
# rather than returning garbage values.
def supports_complex(reduceOp: ReduceOp) -> bool:
denyList = [
ReduceOp.MAX,
ReduceOp.MIN,
ReduceOp.PRODUCT,
ReduceOp.BAND,
ReduceOp.BOR,
ReduceOp.BXOR,
]
return reduceOp not in denyList
class Backend(object):
"""
An enum-like class of available backends: GLOO, NCCL, UCC, MPI, and other registered
backends.
The values of this class are lowercase strings, e.g., ``"gloo"``. They can
be accessed as attributes, e.g., ``Backend.NCCL``.
This class can be directly called to parse the string, e.g.,
``Backend(backend_str)`` will check if ``backend_str`` is valid, and
return the parsed lowercase string if so. It also accepts uppercase strings,
e.g., ``Backend("GLOO")`` returns ``"gloo"``.
.. note:: The entry ``Backend.UNDEFINED`` is present but only used as
initial value of some fields. Users should neither use it directly
nor assume its existence.
"""
UNDEFINED = "undefined"
GLOO = "gloo"
NCCL = "nccl"
UCC = "ucc"
MPI = "mpi"
TCP = "tcp"
_BackendPlugin = namedtuple("_BackendPlugin", ["creator_fn", "extended_api"])
_plugins: Dict[str, _BackendPlugin] = {}
def __new__(cls, name: str):
if not isinstance(name, string_classes):
raise ValueError("Backend name must be a string, but got: {}".format(name))
value = getattr(Backend, name.upper(), Backend.UNDEFINED)
if value == Backend.TCP:
raise ValueError(
"TCP backend has been deprecated. Please use "
"Gloo or MPI backend for collective operations "
"on CPU tensors."
)
elif value == Backend.UNDEFINED:
raise ValueError("Invalid backend: '{}'".format(name))
elif value != Backend.GLOO and value != Backend.NCCL and value != Backend.UCC and value != Backend.MPI:
value = name.lower()
return value
@classmethod
def register_backend(cls, name, func, extended_api=False):
"""
Registers a new backend with the given name and instantiating function.
This class method is used by 3rd party ``ProcessGroup`` extension to
register new backends.
Args:
name (str): Backend name of the ``ProcessGroup`` extension. It
should match the one in ``init_process_group()``.
func (function): Function handler that instantiates the backend.
The function should be implemented in the backend
extension and takes four arguments, including
``store``, ``rank``, ``world_size``, and ``timeout``.
extended_api (bool, optional): Whether the backend supports extended argument structure.
Default: ``False``. If set to ``True``, the backend
will get an instance of ``c10d::DistributedBackendOptions``, and
a process group options object as defined by the backend implementation.
.. note:: This support of 3rd party backend is experimental and subject to change.
"""
# Allow UCC plugin if Pytorch is not built with native support.
# TODO: remove this exception once UCC plugin is fully deprecated.
if (name != Backend.UCC or (name == Backend.UCC and is_ucc_available())):
assert not hasattr(Backend, name.upper()), (
f"{name.upper()} c10d backend already exist"
)
assert name.upper() not in Backend._plugins, (
f"{name.upper()} c10d backend creator function already exist"
)
setattr(Backend, name.upper(), name.upper())
Backend._plugins[name.upper()] = Backend._BackendPlugin(func, extended_api)
# `_backend`, `dist_backend`, and `reduce_op` are here to maintain backward
# compatibility with pre-c10d distributed package.
# TODO: remove them when users are ready to take a hard dependency on PyTorch 1.
_backend: str = Backend.UNDEFINED
dist_backend = Backend
# NOTE(crcrpar): [ReduceOp static class attributes to support `isinstance`]
# A ReduceOp instance of `PREMUL_SUM` is supposed to be created via `_make_nccl_premul_sum`
# while the other `op`s (meaning RedOpType members) can be directly passed to c10d reduce collectives.
# I changed `ReduceOp` to struct from enum class and introduced RedOpType enum class for PREMUL_SUM,
# which broke an implicit contract of ReduceOp being enum-like with which users apply isinstance to
# `op`, for example, `isinstance(ReduceOp.SUM, ReduceOp)`: https://github.com/pytorch/pytorch/issues/87191
DENY_LIST = ("PREMUL_SUM", )
for _red_op_name, _red_op_value in ReduceOp.RedOpType.__members__.items():
setattr(ReduceOp, _red_op_name, _red_op_value if _red_op_name in DENY_LIST else ReduceOp(_red_op_value))
class _reduce_op(object):
r"""
Deprecated enum-like class for reduction operations: ``SUM``, ``PRODUCT``,
``MIN``, and ``MAX``.
:class:`~torch.distributed.ReduceOp` is recommended to use instead.
"""
def __init__(self):
# __members__ is a dict storing key-value pairs for enum classes
for k, v in ReduceOp.RedOpType.__members__.items():
setattr(self, k, v)
self.__members__ = ReduceOp.RedOpType.__members__
def __getattribute__(self, key):
warnings.warn(
"torch.distributed.reduce_op is deprecated, please use "
"torch.distributed.ReduceOp instead"
)
return object.__getattribute__(self, key)
reduce_op = _reduce_op()
class group(object):
# Points to the default PG once initialized.
WORLD: Optional[ProcessGroup] = None
class GroupMember(object):
# Alias to group.WORLD for backward compatibility
WORLD = group.WORLD
NON_GROUP_MEMBER = object()
# Cached process groups
# For NCCL and GLOO pg, it is a map from ProcessGroup to (Backend, Store)
# For MPI pg, it is a map from ProcessGroup to (Backend, None)
_pg_map: Dict[ProcessGroup, Tuple[str, Optional[Store]]] = {}
# Process group's names, map from ProcessGroup to str
_pg_names: Dict[ProcessGroup, str] = {}
# Process group's global rank to local rank mapping
_pg_group_ranks: Dict[ProcessGroup, Dict[int, int]] = {}
# Default process group state
_default_pg_init_method = None
# Process group count for default naming
_group_count = 0
STORE_BASED_BARRIER_PREFIX = "store_based_barrier_key"
def _get_pg_device(group: ProcessGroup):
"""
Returns the device to use with ``group``.
This is cuda for NCCL and CPU for everything else
"""
if _check_for_nccl_backend(group):
return torch.device("cuda", torch.cuda.current_device())
return torch.device("cpu")
def _store_based_barrier(rank, store, timeout):
"""
Barrier based on store which is used for synchronizing processes after
``init_process_group`` or ``new_group``. Intended to be used only with
those two methods and is not a generic alternative to ``barrier()``.
"""
store_key = "{}:{}".format(STORE_BASED_BARRIER_PREFIX, _group_count)
store.add(store_key, 1)
logger.info("Added key: {} to store for rank: {}".format(store_key, rank))
# Now wait for all workers to check in with the store.
world_size = get_world_size()
# Use 'add' instead of 'get' since for some store implementations 'add'
# doesn't work well with 'get'. Ideally the store implementations should
# be fixed, but for backward compatiblity reasons it is risky to change
# the store implementations. Once, we completely migrate away from these
# legacy stores, we can use 'get' here instead.
worker_count = store.add(store_key, 0)
start = time.time()
log_time = time.time()
while worker_count != world_size:
time.sleep(0.01)
worker_count = store.add(store_key, 0)
# Print status periodically to keep track.
if timedelta(seconds=(time.time() - log_time)) > timedelta(seconds=10):
logger.info(
"Waiting in store based barrier to initialize process group for "
"rank: {}, key: {} (world_size={}, worker_count={}, timeout={})".format(
rank, store_key, world_size, worker_count, timeout
)
)
log_time = time.time()
if timedelta(seconds=(time.time() - start)) > timeout:
raise RuntimeError(
"Timed out initializing process group in store based barrier on "
"rank: {}, for key: {} (world_size={}, worker_count={}, timeout={})".format(
rank, store_key, world_size, worker_count, timeout
)
)
logger.info(
f"Rank {rank}: Completed store-based barrier for key:{store_key} with {world_size} nodes."
)
def _rank_not_in_group(group: ProcessGroup):
"""
Helper that checks if the current process's rank is not in a given group.
"""
if group is None:
return False
return group == GroupMember.NON_GROUP_MEMBER
def _warn_not_in_group(op_name):
global_rank = -1 if GroupMember.WORLD is None else GroupMember.WORLD.rank()
warnings.warn(
f"Running {op_name} on global rank {global_rank} which does not "
"belong to the given group."
)
def get_group_rank(group: ProcessGroup, global_rank: int) -> int:
"""
Translate a global rank into a group rank.
``global_rank`` must be part of ``group`` otherwise this raises RuntimeError.
Args:
group (ProcessGroup): ProcessGroup to find the relative rank.
global_rank (int): Global rank to query.
Returns:
Group rank of ``global_rank`` relative to ``group``
N.B. calling this function on the default process group returns identity
"""
if group is GroupMember.WORLD:
return global_rank
if group not in _pg_group_ranks:
raise RuntimeError(f"Group {group} is not registered, please create group with torch.distributed.new_group API")
group_ranks = _pg_group_ranks[group]
if global_rank not in group_ranks:
raise RuntimeError(f"Global rank {global_rank} is not part of group {group}")
return group_ranks[global_rank]
def get_global_rank(group: ProcessGroup, group_rank: int) -> int:
"""
Translate a group rank into a global rank.
``group_rank`` must be part of `group` otherwise this raises RuntimeError.
Args:
group (ProcessGroup): ProcessGroup to find the global rank from.
group_rank (int): Group rank to query.
Returns:
Global rank of ``group_rank`` relative to ``group``
N.B. calling this function on the default process group returns identity
"""
if group is GroupMember.WORLD:
return group_rank
if group not in _pg_group_ranks:
raise RuntimeError(f"Group {group} is not registered, please create group with torch.distributed.new_group API")
for rank, grp_rank in _pg_group_ranks[group].items():
if grp_rank == group_rank:
return rank
raise RuntimeError(f"Group rank {group_rank} is not part of group {group}")
# TODO: remove this once the ecosystem moves away from it.
def _get_global_rank(group, rank):
"""
This method is deprecated, please use get_global_rank.
"""
warnings.warn(
"torch.distributed.distributed_c10d._get_global_rank is deprecated "
"please use torch.distributed.distributed_c10d.get_global_rank instead"
)
return get_global_rank(group, rank)
def get_process_group_ranks(group: ProcessGroup):
"""
Get all ranks associated with ``group``.
Args:
group (ProcessGroup): ProcessGroup to get all ranks from.
Returns:
List of global ranks ordered by group rank.
"""
return list(_pg_group_ranks[group].keys())
def _get_group_size(group):
"""
Helper that gets a given group's world size.
"""
if group is GroupMember.WORLD or group is None:
default_pg = _get_default_group()
return default_pg.size()
return group.size()
def _check_single_tensor(param, param_name):
"""
Helper to check that the parameter ``param_name`` is a single tensor.
"""
if not isinstance(param, torch.Tensor):
raise RuntimeError(
"Invalid function argument. Expected parameter `{}` "
"to be of type torch.Tensor.".format(param_name)
)
def _check_tensor_list(param, param_name):
"""
Helper to check that the parameter ``param_name`` is a list of tensors.
"""
if not isinstance(param, list) or not all(
isinstance(p, torch.Tensor) for p in param
):
raise RuntimeError(
"Invalid function argument. Expected parameter `{}` "
"to be of type List[torch.Tensor].".format(param_name)
)
def _as_iterable(obj) -> collections.abc.Iterable:
return obj if isinstance(obj, list) else (obj,)
def _ensure_all_tensors_same_dtype(*tensors) -> None:
last_dtype = None
for tensor in itertools.chain(*map(_as_iterable, tensors)):
tensor_dtype = tensor.dtype
# Mixing complex and its element type is allowed
if tensor_dtype.is_complex:
tensor_dtype = torch.float32 if tensor_dtype == torch.complex64 else torch.complex128
if last_dtype is None:
last_dtype = tensor_dtype
else:
if last_dtype != tensor_dtype:
raise RuntimeError(
"Invalid usage of tensors with different dtypes"
f"Found {last_dtype} and {tensor.dtype}"
)
def _check_op(op):
"""
Helper to check that the ``op`` is either isend or irecv.
"""
if op not in [isend, irecv]:
raise RuntimeError(
"Invalid ``op``. Expected ``op`` "
"to be of type ``torch.distributed.isend`` or "
"``torch.distributed.irecv``."
)
def _check_p2p_op_list(p2p_op_list):
"""
Helper to check that the ``p2p_op_list`` is a list of P2POp instances and
all ops use the same group.
"""
if not isinstance(p2p_op_list, list) or not all(
isinstance(p2p_op, P2POp) for p2p_op in p2p_op_list
):
raise RuntimeError(
"Invalid ``p2p_op_list``. Each op is expected to "
"to be of type ``torch.distributed.P2POp``."
)
group = p2p_op_list[0].group
if not all(group == p2p_op.group for p2p_op in p2p_op_list):
raise RuntimeError("All ops need to use the same group.")
def is_mpi_available() -> bool:
"""
Checks if the MPI backend is available.
"""
return _MPI_AVAILABLE
def is_nccl_available() -> bool:
"""
Checks if the NCCL backend is available.
"""
return _NCCL_AVAILABLE
def is_gloo_available() -> bool:
"""
Checks if the Gloo backend is available.
"""
return _GLOO_AVAILABLE
def is_ucc_available() -> bool:
"""
Checks if the UCC backend is available.
"""
return _UCC_AVAILABLE
def is_initialized() -> bool:
"""
Checking if the default process group has been initialized
"""
return GroupMember.WORLD is not None
def is_torchelastic_launched() -> bool:
"""
Checks whether this process was launched with ``torch.distributed.elastic``
(aka torchelastic). The existence of ``TORCHELASTIC_RUN_ID`` environment
variable is used as a proxy to determine whether the current process
was launched with torchelastic. This is a reasonable proxy since
``TORCHELASTIC_RUN_ID`` maps to the rendezvous id which is always a
non-null value indicating the job id for peer discovery purposes..
"""
return os.getenv("TORCHELASTIC_RUN_ID") is not None
def _get_default_group():
"""
Getting the default process group created by init_process_group
"""
if not is_initialized():
raise RuntimeError(
"Default process group has not been initialized, "
"please make sure to call init_process_group."
)
return GroupMember.WORLD
def _get_default_store():
"""
Getting the default store created by init_process_group
"""
if not is_initialized():
raise RuntimeError(
"Default process group has not been initialized, "
"please make sure to call init_process_group."
)
default_pg = _get_default_group()
_, default_store = _pg_map[default_pg]
return default_store
def _update_default_pg(pg):
GroupMember.WORLD = group.WORLD = pg
def get_backend(group: Optional[ProcessGroup] = None) -> str:
"""
Returns the backend of the given process group.
Args:
group (ProcessGroup, optional): The process group to work on. The
default is the general main process group. If another specific group
is specified, the calling process must be part of :attr:`group`.
Returns:
The backend of the given process group as a lower case string.
"""
if group is None:
pg = _get_default_group()
else:
pg = group
if _rank_not_in_group(pg):
raise RuntimeError("Invalid process group specified")
pg_store = _pg_map.get(pg, None)
assert pg_store is not None
return pg_store[0]
def init_process_group(
backend: Union[str, Backend],
init_method: Optional[str] = None,
timeout: timedelta = default_pg_timeout,
world_size: int = -1,
rank: int = -1,
store: Optional[Store] = None,
group_name: str = "",
pg_options: Optional[Any] = None,
):
"""
Initializes the default distributed process group, and this will also
initialize the distributed package.
There are 2 main ways to initialize a process group:
1. Specify ``store``, ``rank``, and ``world_size`` explicitly.
2. Specify ``init_method`` (a URL string) which indicates where/how
to discover peers. Optionally specify ``rank`` and ``world_size``,
or encode all required parameters in the URL and omit them.
If neither is specified, ``init_method`` is assumed to be "env://".
Args:
backend (str or Backend): The backend to use. Depending on
build-time configurations, valid values include ``mpi``, ``gloo``,
``nccl``, and ``ucc``. This field should be given as a lowercase
string (e.g., ``"gloo"``), which can also be accessed via
:class:`Backend` attributes (e.g., ``Backend.GLOO``). If using
multiple processes per machine with ``nccl`` backend, each process
must have exclusive access to every GPU it uses, as sharing GPUs
between processes can result in deadlocks. ``ucc`` backend is
experimental.
init_method (str, optional): URL specifying how to initialize the
process group. Default is "env://" if no
``init_method`` or ``store`` is specified.
Mutually exclusive with ``store``.
world_size (int, optional): Number of processes participating in
the job. Required if ``store`` is specified.
rank (int, optional): Rank of the current process (it should be a
number between 0 and ``world_size``-1).
Required if ``store`` is specified.
store(Store, optional): Key/value store accessible to all workers, used
to exchange connection/address information.
Mutually exclusive with ``init_method``.
timeout (timedelta, optional): Timeout for operations executed against
the process group. Default value equals 30 minutes.
This is applicable for the ``gloo`` backend. For ``nccl``, this is
applicable only if the environment variable ``NCCL_BLOCKING_WAIT``
or ``NCCL_ASYNC_ERROR_HANDLING`` is set to 1. When
``NCCL_BLOCKING_WAIT`` is set, this is the duration for which the
process will block and wait for collectives to complete before
throwing an exception. When ``NCCL_ASYNC_ERROR_HANDLING`` is set,
this is the duration after which collectives will be aborted
asynchronously and the process will crash. ``NCCL_BLOCKING_WAIT``
will provide errors to the user which can be caught and handled,
but due to its blocking nature, it has a performance overhead. On
the other hand, ``NCCL_ASYNC_ERROR_HANDLING`` has very little
performance overhead, but crashes the process on errors. This is
done since CUDA execution is async and it is no longer safe to
continue executing user code since failed async NCCL operations
might result in subsequent CUDA operations running on corrupted
data. Only one of these two environment variables should be set.
For ``ucc``, blocking wait is supported similar to NCCL. However,
async error handling is done differently since with UCC we have
progress thread and not watch-dog thread.
group_name (str, optional, deprecated): Group name.
pg_options (ProcessGroupOptions, optional): process group options
specifying what additional options need to be passed in during
the construction of specific process groups. As of now, the only
options we support is ``ProcessGroupNCCL.Options`` for the ``nccl``
backend, ``is_high_priority_stream`` can be specified so that
the nccl backend can pick up high priority cuda streams when
there're compute kernels waiting.
.. note:: To enable ``backend == Backend.MPI``, PyTorch needs to be built from source
on a system that supports MPI.
"""
global _pg_group_ranks
global _backend
global _default_pg_init_method
if not isinstance(timeout, timedelta):
raise RuntimeError(
"Expected timeout argument to be of type" "datetime.timedelta"
)
if GroupMember.WORLD is not None:
raise RuntimeError("trying to initialize the default process group " "twice!")
assert (store is None) or (
init_method is None
), "Cannot specify both init_method and store."
if store is not None:
assert world_size > 0, "world_size must be positive if using store"
assert rank >= 0, "rank must be non-negative if using store"
elif init_method is None:
init_method = "env://"
backend = Backend(backend)
if backend == Backend.MPI:
if world_size != -1 or rank != -1:
warnings.warn(
"For MPI backend, world_size ({}) and rank ({}) "
"are ignored since they are assigned by the "
"MPI runtime.".format(world_size, rank)
)
default_pg = _new_process_group_helper(
-1, -1, [], Backend.MPI, None, group_name=group_name, timeout=timeout
)
_update_default_pg(default_pg)
else:
# backward compatible API
if store is None:
rendezvous_iterator = rendezvous(
init_method, rank, world_size, timeout=timeout
)
store, rank, world_size = next(rendezvous_iterator)
store.set_timeout(timeout)
# Use a PrefixStore to avoid accidental overrides of keys used by
# different systems (e.g. RPC) in case the store is multi-tenant.
store = PrefixStore("default_pg", store)
default_pg = _new_process_group_helper(
world_size,
rank,
[],
backend,
store,
pg_options=pg_options,
group_name=group_name,
timeout=timeout,
)
_update_default_pg(default_pg)
_pg_group_ranks[GroupMember.WORLD] = {i: i for i in range(GroupMember.WORLD.size())} # type: ignore[attr-defined, index]
_backend = _pg_map[GroupMember.WORLD][0] # type: ignore[index]
_default_pg_init_method = init_method
# barrier at the end to ensure that once we return from this method, all
# process groups including global variables are updated correctly on all
# ranks.
if backend == Backend.MPI:
# MPI backend doesn't use store.
barrier()
else:
# Use store based barrier here since barrier() used a bunch of
# default devices and messes up NCCL internal state.
_store_based_barrier(rank, store, timeout)
# Set sequence numbers for gloo and nccl process groups.
if get_backend(default_pg) in [Backend.GLOO, Backend.NCCL]:
default_pg._set_sequence_number_for_group()
def _new_process_group_helper(
group_size,
group_rank,
global_ranks_in_group,
backend,
store,
pg_options=None,
group_name=None,
timeout=default_pg_timeout,
):
"""
Create a new distributed process group.
This function must be called by ALL processes in the global group, even if
the calling process is not part of the newly created group. In that case,
this function returns GroupMember.NON_GROUP_MEMBER.
This function is called with ``group_ranks == []`` for the default group.
"""
global _pg_map
global _group_count
global _pg_names
if not group_name:
group_name = str(_group_count)
_group_count += 1
if group_name in _pg_names.values():
raise RuntimeError(
"The specified group name has already been "
"created, please use a different group name"
)
if not isinstance(timeout, timedelta):
raise RuntimeError(
"Expected timeout argument to be of type" "datetime.timedelta"
)
# The list of group ranks is empty if we're creating the default group.
is_default_group = len(global_ranks_in_group) == 0
backend = Backend(backend)
pg: Union[ProcessGroupGloo, ProcessGroupMPI, ProcessGroupNCCL, ProcessGroupUCC]
if backend == Backend.MPI:
if not is_mpi_available():
raise RuntimeError(
"Distributed package doesn't have MPI built in."
" MPI is only included if you build PyTorch from"
" source on a host that has MPI installed."
)
pg = ProcessGroupMPI.create(global_ranks_in_group)
if not pg:
return GroupMember.NON_GROUP_MEMBER
_pg_map[pg] = (Backend.MPI, None)
_pg_names[pg] = group_name
else:
# If this is a subgroup (which means group_ranks is specified),
# we check if the current process is a member of the new group.
if not is_default_group:
global_rank = _get_default_group().rank()
if global_rank not in global_ranks_in_group:
return GroupMember.NON_GROUP_MEMBER
# Use the group name as prefix in the default store, such that
# a single store can be reused by multiple groups.
prefix_store = PrefixStore(group_name, store)
if backend == Backend.GLOO:
if pg_options is not None:
raise RuntimeError("GLOO options not supported")
pg = ProcessGroupGloo(prefix_store, group_rank, group_size, timeout=timeout)
# In debug mode and if GLOO is available, wrap in a wrapper PG that
# enables enhanced collective checking for debugability.
if get_debug_level() == DebugLevel.DETAIL:
if not _GLOO_AVAILABLE:
logger.info(
"""TORCH_DISTRIBUTED_DEBUG was set to DETAIL, but
GLOO is not available. Build with Gloo to
create a wrapper process group in debug mode
to aid collective desynchronization debugging."""
)
else:
pg = _create_process_group_wrapper(
wrapped_pg=pg,
store_prefix=group_name,
store=store,
rank=group_rank,
world_size=group_size,
timeout=timeout,
)
_pg_map[pg] = (Backend.GLOO, store)
_pg_names[pg] = group_name
elif backend == Backend.NCCL:
if not is_nccl_available():
raise RuntimeError("Distributed package doesn't have NCCL " "built in")
if pg_options is not None:
assert isinstance(
pg_options, ProcessGroupNCCL.Options
), "Expected pg_options argument to be of type ProcessGroupNCCL.Options"
else:
# default pg_options for NCCL
pg_options = ProcessGroupNCCL.Options()
pg_options.is_high_priority_stream = False
pg_options._timeout = timeout
pg = ProcessGroupNCCL(prefix_store, group_rank, group_size, pg_options)
# In debug mode and if GLOO is available, wrap in a wrapper PG that
# enables enhanced collective checking for debugability.
if get_debug_level() == DebugLevel.DETAIL:
if not _GLOO_AVAILABLE:
logger.info(
"""TORCH_DISTRIBUTED_DEBUG was set to DETAIL, but
GLOO is not available. Build with Gloo to
create a wrapper process group in debug mode
to aid collective desynchronization debugging."""
)
else:
pg = _create_process_group_wrapper(
wrapped_pg=pg,
store_prefix=group_name,
store=store,
rank=group_rank,
world_size=group_size,
timeout=timeout,
)
_pg_map[pg] = (Backend.NCCL, store)
_pg_names[pg] = group_name
elif backend == Backend.UCC and is_ucc_available():
# TODO: once UCC plugin is fully deprecated, remove
# is_ucc_available() from above elif-condition and raise
# RuntimeError if is_ucc_available() returns false.
pg = ProcessGroupUCC(prefix_store, group_rank, group_size, timeout=timeout)
# In debug mode and if GLOO is available, wrap in a wrapper PG that
# enables enhanced collective checking for debugability.
if get_debug_level() == DebugLevel.DETAIL:
if not _GLOO_AVAILABLE:
logger.info(
"""TORCH_DISTRIBUTED_DEBUG was set to DETAIL, but
GLOO is not available. Build with Gloo to
create a wrapper process group in debug mode
to aid collective desynchronization debugging."""
)
else:
pg = _create_process_group_wrapper(
wrapped_pg=pg,
store_prefix=group_name,
store=store,
rank=group_rank,
world_size=group_size,
timeout=timeout,
)
_pg_map[pg] = (Backend.UCC, store)
_pg_names[pg] = group_name
else:
assert backend.upper() in Backend._plugins, (
f"Unknown c10d backend type {backend.upper()}"
)
backend_plugin = Backend._plugins[backend.upper()]
creator_fn = backend_plugin.creator_fn
extended_api = backend_plugin.extended_api
if not extended_api:
pg = creator_fn(prefix_store, group_rank, group_size, timeout)
else:
dist_backend_opts = _DistributedBackendOptions()
dist_backend_opts.store = prefix_store
dist_backend_opts.group_rank = group_rank
dist_backend_opts.group_size = group_size
dist_backend_opts.timeout = timeout
dist_backend_opts.group_id = group_name
dist_backend_opts.global_ranks_in_group = global_ranks_in_group
pg = creator_fn(dist_backend_opts, pg_options)
_pg_map[pg] = (backend, store)
_pg_names[pg] = group_name
return pg
def destroy_process_group(group: Optional[ProcessGroup] = None):
"""
Destroy a given process group, and deinitialize the distributed package
Args:
group (ProcessGroup, optional): The process group to be destroyed, if
group.WORLD is given, all process
groups including the default one will
be destroyed.
"""
global _pg_map
global _pg_names
global _pg_group_ranks
global _default_pg_init_method
global _group_count
if group == GroupMember.NON_GROUP_MEMBER:
return
if group is None:
pg = GroupMember.WORLD
else:
pg = group
assert pg is not None
if _pg_map.get(pg, None) is None:
raise RuntimeError("Invalid process group specified")