File size: 688 Bytes
828992f |
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 |
import os
import random
import torch
import numpy as np
def init(rank):
nranks = 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE'])
nranks = max(1, nranks)
# nranks = -1
# is_distributed = nranks > 0
is_distributed = False
if rank == 0:
print('nranks =', nranks, '\t num_gpus =', torch.cuda.device_count())
if is_distributed:
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(rank % num_gpus)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
return nranks, is_distributed
def barrier(rank):
if rank >= 0:
torch.distributed.barrier()
|