| import os |
| from typing import List |
| from typing import Union |
|
|
| import sys |
| import torch |
| import torch.distributed as tdist |
| import torch.multiprocessing as mp |
|
|
| __rank, __local_rank, __world_size, __device = 0, 0, 1, 'cpu' |
| __initialized = False |
|
|
|
|
| def initialized(): |
| return __initialized |
|
|
|
|
| def initialize(backend='nccl'): |
| global __device |
| if not torch.cuda.is_available(): |
| print(f'[dist initialize] cuda is not available, use cpu instead', file=sys.stderr) |
| return |
| elif 'RANK' not in os.environ: |
| __device = torch.empty(1).cuda().device |
| print(f'[dist initialize] RANK is not set, use 1 GPU instead', file=sys.stderr) |
| return |
| |
| |
| if mp.get_start_method(allow_none=True) is None: |
| mp.set_start_method('spawn') |
| global_rank, num_gpus = int(os.environ['RANK']), torch.cuda.device_count() |
| local_rank = global_rank % num_gpus |
| torch.cuda.set_device(local_rank) |
| tdist.init_process_group(backend=backend) |
| |
| global __rank, __local_rank, __world_size, __initialized |
| __local_rank = local_rank |
| __rank, __world_size = tdist.get_rank(), tdist.get_world_size() |
| __device = torch.empty(1).cuda().device |
| __initialized = True |
| |
| assert tdist.is_initialized(), 'torch.distributed is not initialized!' |
|
|
|
|
| def get_rank(): |
| return __rank |
|
|
|
|
| def get_local_rank(): |
| return __local_rank |
|
|
|
|
| def get_world_size(): |
| return __world_size |
|
|
|
|
| def get_device(): |
| return __device |
|
|
|
|
| def is_master(): |
| return __rank == 0 |
|
|
|
|
| def is_local_master(): |
| return __local_rank == 0 |
|
|
|
|
| def barrier(): |
| if __initialized: |
| tdist.barrier() |
|
|
|
|
| def parallelize(net, syncbn=False): |
| if syncbn: |
| net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net) |
| net = net.cuda() |
| net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[get_local_rank()], find_unused_parameters=False, broadcast_buffers=False) |
| return net |
|
|
|
|
| def allreduce(t: torch.Tensor) -> None: |
| if __initialized: |
| if not t.is_cuda: |
| cu = t.detach().cuda() |
| tdist.all_reduce(cu) |
| t.copy_(cu.cpu()) |
| else: |
| tdist.all_reduce(t) |
|
|
|
|
| def allgather(t: torch.Tensor, cat=True) -> Union[List[torch.Tensor], torch.Tensor]: |
| if __initialized: |
| if not t.is_cuda: |
| t = t.cuda() |
| ls = [torch.empty_like(t) for _ in range(__world_size)] |
| tdist.all_gather(ls, t) |
| else: |
| ls = [t] |
| if cat: |
| ls = torch.cat(ls, dim=0) |
| return ls |
|
|
|
|
| def broadcast(t: torch.Tensor, src_rank) -> None: |
| if __initialized: |
| if not t.is_cuda: |
| cu = t.detach().cuda() |
| tdist.broadcast(cu, src=src_rank) |
| t.copy_(cu.cpu()) |
| else: |
| tdist.broadcast(t, src=src_rank) |
|
|
