|
|
|
|
|
|
|
|
"""Distributed helpers.""" |
|
|
|
|
|
import torch |
|
|
import torch.distributed as dist |
|
|
|
|
|
_LOCAL_PROCESS_GROUP = None |
|
|
|
|
|
|
|
|
def all_reduce(tensors, average=True): |
|
|
""" |
|
|
All reduce the provided tensors from all processes across machines. |
|
|
Args: |
|
|
tensors (list): tensors to perform all reduce across all processes in |
|
|
all machines. |
|
|
average (bool): scales the reduced tensor by the number of overall |
|
|
processes across all machines. |
|
|
""" |
|
|
|
|
|
for tensor in tensors: |
|
|
dist.all_reduce(tensor, async_op=False) |
|
|
if average: |
|
|
world_size = dist.get_world_size() |
|
|
for tensor in tensors: |
|
|
tensor.mul_(1.0 / world_size) |
|
|
return tensors |
|
|
|
|
|
|
|
|
def is_master_proc(num_gpus=8): |
|
|
""" |
|
|
Determines if the current process is the master process. |
|
|
""" |
|
|
if torch.distributed.is_initialized(): |
|
|
return dist.get_rank() % num_gpus == 0 |
|
|
else: |
|
|
return True |
|
|
|
|
|
|
|
|
def get_world_size(): |
|
|
""" |
|
|
Get the size of the world. |
|
|
""" |
|
|
if not dist.is_available(): |
|
|
return 1 |
|
|
if not dist.is_initialized(): |
|
|
return 1 |
|
|
return dist.get_world_size() |
|
|
|
|
|
|
|
|
def init_distributed_training(cfg): |
|
|
""" |
|
|
Initialize variables needed for distributed training. |
|
|
""" |
|
|
if cfg.NUM_GPUS <= 1: |
|
|
return |
|
|
num_gpus_per_machine = cfg.NUM_GPUS |
|
|
num_machines = dist.get_world_size() // num_gpus_per_machine |
|
|
for i in range(num_machines): |
|
|
ranks_on_i = list( |
|
|
range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine) |
|
|
) |
|
|
pg = dist.new_group(ranks_on_i) |
|
|
if i == cfg.SHARD_ID: |
|
|
global _LOCAL_PROCESS_GROUP |
|
|
_LOCAL_PROCESS_GROUP = pg |
|
|
|
