You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
4.2 KiB
138 lines
4.2 KiB
""" Distributed training/validation utils
|
|
|
|
Hacked together by / Copyright 2020 Ross Wightman
|
|
"""
|
|
import os
|
|
|
|
import torch
|
|
from torch import distributed as dist
|
|
|
|
try:
|
|
import horovod.torch as hvd
|
|
except ImportError:
|
|
hvd = None
|
|
|
|
from .model import unwrap_model
|
|
|
|
|
|
def reduce_tensor(tensor, n):
|
|
rt = tensor.clone()
|
|
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
|
|
rt /= n
|
|
return rt
|
|
|
|
|
|
def distribute_bn(model, world_size, reduce=False):
|
|
# ensure every node has the same running bn stats
|
|
for bn_name, bn_buf in unwrap_model(model).named_buffers(recurse=True):
|
|
if ('running_mean' in bn_name) or ('running_var' in bn_name):
|
|
if reduce:
|
|
# average bn stats across whole group
|
|
torch.distributed.all_reduce(bn_buf, op=dist.ReduceOp.SUM)
|
|
bn_buf /= float(world_size)
|
|
else:
|
|
# broadcast bn stats from rank 0 to whole group
|
|
torch.distributed.broadcast(bn_buf, 0)
|
|
|
|
|
|
def is_global_primary(args):
|
|
return args.rank == 0
|
|
|
|
|
|
def is_local_primary(args):
|
|
return args.local_rank == 0
|
|
|
|
|
|
def is_primary(args, local=False):
|
|
return is_local_primary(args) if local else is_global_primary(args)
|
|
|
|
|
|
def is_distributed_env():
|
|
if 'WORLD_SIZE' in os.environ:
|
|
return int(os.environ['WORLD_SIZE']) > 1
|
|
if 'SLURM_NTASKS' in os.environ:
|
|
return int(os.environ['SLURM_NTASKS']) > 1
|
|
return False
|
|
|
|
|
|
def world_info_from_env():
|
|
local_rank = 0
|
|
for v in ('LOCAL_RANK', 'MPI_LOCALRANKID', 'SLURM_LOCALID', 'OMPI_COMM_WORLD_LOCAL_RANK'):
|
|
if v in os.environ:
|
|
local_rank = int(os.environ[v])
|
|
break
|
|
|
|
global_rank = 0
|
|
for v in ('RANK', 'PMI_RANK', 'SLURM_PROCID', 'OMPI_COMM_WORLD_RANK'):
|
|
if v in os.environ:
|
|
global_rank = int(os.environ[v])
|
|
break
|
|
|
|
world_size = 1
|
|
for v in ('WORLD_SIZE', 'PMI_SIZE', 'SLURM_NTASKS', 'OMPI_COMM_WORLD_SIZE'):
|
|
if v in os.environ:
|
|
world_size = int(os.environ[v])
|
|
break
|
|
|
|
return local_rank, global_rank, world_size
|
|
|
|
|
|
def init_distributed_device(args):
|
|
# Distributed training = training on more than one GPU.
|
|
# Works in both single and multi-node scenarios.
|
|
args.distributed = False
|
|
args.world_size = 1
|
|
args.rank = 0 # global rank
|
|
args.local_rank = 0
|
|
|
|
# TBD, support horovod?
|
|
# if args.horovod:
|
|
# assert hvd is not None, "Horovod is not installed"
|
|
# hvd.init()
|
|
# args.local_rank = int(hvd.local_rank())
|
|
# args.rank = hvd.rank()
|
|
# args.world_size = hvd.size()
|
|
# args.distributed = True
|
|
# os.environ['LOCAL_RANK'] = str(args.local_rank)
|
|
# os.environ['RANK'] = str(args.rank)
|
|
# os.environ['WORLD_SIZE'] = str(args.world_size)
|
|
dist_backend = getattr(args, 'dist_backend', 'nccl')
|
|
dist_url = getattr(args, 'dist_url', 'env://')
|
|
if is_distributed_env():
|
|
if 'SLURM_PROCID' in os.environ:
|
|
# DDP via SLURM
|
|
args.local_rank, args.rank, args.world_size = world_info_from_env()
|
|
# SLURM var -> torch.distributed vars in case needed
|
|
os.environ['LOCAL_RANK'] = str(args.local_rank)
|
|
os.environ['RANK'] = str(args.rank)
|
|
os.environ['WORLD_SIZE'] = str(args.world_size)
|
|
torch.distributed.init_process_group(
|
|
backend=dist_backend,
|
|
init_method=dist_url,
|
|
world_size=args.world_size,
|
|
rank=args.rank,
|
|
)
|
|
else:
|
|
# DDP via torchrun, torch.distributed.launch
|
|
args.local_rank, _, _ = world_info_from_env()
|
|
torch.distributed.init_process_group(
|
|
backend=dist_backend,
|
|
init_method=dist_url,
|
|
)
|
|
args.world_size = torch.distributed.get_world_size()
|
|
args.rank = torch.distributed.get_rank()
|
|
args.distributed = True
|
|
|
|
if torch.cuda.is_available():
|
|
if args.distributed:
|
|
device = 'cuda:%d' % args.local_rank
|
|
else:
|
|
device = 'cuda:0'
|
|
torch.cuda.set_device(device)
|
|
else:
|
|
device = 'cpu'
|
|
|
|
args.device = device
|
|
device = torch.device(device)
|
|
return device
|