Source code for imaginaire.utils.distributed

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import functools
import ctypes

import torch
import torch.distributed as dist


[docs]def init_dist(local_rank, backend='nccl', **kwargs): r"""Initialize distributed training""" if dist.is_available(): if dist.is_initialized(): return torch.cuda.current_device() torch.cuda.set_device(local_rank) dist.init_process_group(backend=backend, init_method='env://', **kwargs) # Increase the L2 fetch granularity for faster speed. _libcudart = ctypes.CDLL('libcudart.so') # Set device limit on the current device # cudaLimitMaxL2FetchGranularity = 0x05 pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int)) _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128)) _libcudart.cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
# assert pValue.contents.value == 128
[docs]def get_rank(): r"""Get rank of the thread.""" rank = 0 if dist.is_available(): if dist.is_initialized(): rank = dist.get_rank() return rank
[docs]def get_world_size(): r"""Get world size. How many GPUs are available in this job.""" world_size = 1 if dist.is_available(): if dist.is_initialized(): world_size = dist.get_world_size() return world_size
[docs]def master_only(func): r"""Apply this function only to the master GPU.""" @functools.wraps(func) def wrapper(*args, **kwargs): r"""Simple function wrapper for the master function""" if get_rank() == 0: return func(*args, **kwargs) else: return None return wrapper
[docs]def is_master(): r"""check if current process is the master""" return get_rank() == 0
[docs]def is_local_master(): return torch.cuda.current_device() == 0
[docs]@master_only def master_only_print(*args): r"""master-only print""" print(*args)
[docs]def dist_reduce_tensor(tensor, rank=0, reduce='mean'): r""" Reduce to rank 0 """ world_size = get_world_size() if world_size < 2: return tensor with torch.no_grad(): dist.reduce(tensor, dst=rank) if get_rank() == rank: if reduce == 'mean': tensor /= world_size elif reduce == 'sum': pass else: raise NotImplementedError return tensor
[docs]def dist_all_reduce_tensor(tensor, reduce='mean'): r""" Reduce to all ranks """ world_size = get_world_size() if world_size < 2: return tensor with torch.no_grad(): dist.all_reduce(tensor) if reduce == 'mean': tensor /= world_size elif reduce == 'sum': pass else: raise NotImplementedError return tensor
[docs]def dist_all_gather_tensor(tensor): r""" gather to all ranks """ world_size = get_world_size() if world_size < 2: return [tensor] tensor_list = [ torch.ones_like(tensor) for _ in range(dist.get_world_size())] with torch.no_grad(): dist.all_gather(tensor_list, tensor) return tensor_list