File size: 787 Bytes
45cf443 | 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 31 32 33 34 35 36 37 38 39 40 | import os
import torch
import torch.distributed as dist
def init_distributed_mode():
if int(os.environ.get("RANK", -1)) == -1:
return 0
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
def is_distributed():
return dist.is_initialized()
def is_main_process():
return not dist.is_initialized() or dist.get_rank() == 0
def get_rank():
return dist.get_rank() if dist.is_initialized() else 0
def get_world_size():
return dist.get_world_size() if dist.is_initialized() else 1
def barrier():
if dist.is_initialized():
dist.barrier()
def destroy():
if dist.is_initialized():
dist.barrier()
dist.destroy_process_group()
|