/
custom_utils.py
30 lines (28 loc) · 1017 Bytes
/
custom_utils.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
r"""
Utils to play with PyTorch.
"""
import torch.distributed as dist
# pylint: disable=broad-except
# pylint: disable=protected-access
def get_torch_default_comm():
r"""
The NCCL communicator is needed so that Fast MoE can perform customized
communication operators in the C code. However, it is not a publicly
available variable. Therefore, a hacking class of the `ProcessGroupNCCL`
in Fast MoE's C code takes the `_default_pg` and tries to dig the
communicator out from the object. As PyTorch's private interface varies from
time to time, different hacking techniques are tried one-by-one to be
compatible with various versions of PyTorch.
"""
try:
comm = dist.distributed_c10d._get_default_group()
return comm
except Exception as _:
pass
try:
comm = dist.distributed_c10d._default_pg
if comm is not None:
return comm
except Exception as _:
pass
raise RuntimeError("Unsupported PyTorch version")