Spaces:
Sleeping
Sleeping
| import os | |
| import pickle | |
| import torch | |
| import torch.distributed as dist | |
| # import torch.distributed.elastic.multiprocessing.api | |
| def setup_for_distributed(is_master): | |
| """ | |
| This function disables printing when not in master process | |
| """ | |
| import builtins as __builtin__ | |
| builtin_print = __builtin__.print | |
| def new_print(*args, **kwargs): | |
| force = kwargs.pop('force', False) | |
| if is_master or force: | |
| builtin_print(*args, **kwargs) | |
| __builtin__.print = new_print | |
| def is_dist_avail_and_initialized(): | |
| return dist.is_available() and dist.is_initialized() | |
| def get_world_size(): | |
| if not is_dist_avail_and_initialized(): | |
| return 1 | |
| return dist.get_world_size() | |
| def get_rank(): | |
| if not is_dist_avail_and_initialized(): | |
| return 0 | |
| return dist.get_rank() | |
| def get_local_size(): | |
| if not is_dist_avail_and_initialized(): | |
| return 1 | |
| return int(os.environ['LOCAL_SIZE']) | |
| def get_local_rank(): | |
| if not is_dist_avail_and_initialized(): | |
| return 0 | |
| return int(os.environ['LOCAL_RANK']) | |
| def is_main_process(): | |
| return get_rank() == 0 | |
| def save_on_master(*args, **kwargs): | |
| if is_main_process(): | |
| torch.save(*args, **kwargs) | |
| def init_distributed_mode(args): | |
| if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: | |
| args.rank = int(os.environ["RANK"]) | |
| args.world_size = int(os.environ['WORLD_SIZE']) | |
| args.gpu = int(os.environ['LOCAL_RANK']) | |
| args.dist_url = 'env://' | |
| os.environ['LOCAL_SIZE'] = str(torch.cuda.device_count()) | |
| else: | |
| print('Not using distributed mode') | |
| args.world_size = 1 | |
| args.distributed = False | |
| return | |
| args.distributed = True | |
| torch.cuda.set_device(args.gpu) | |
| args.dist_backend = 'nccl' | |
| print('| distributed init (rank {}): {}'.format( | |
| args.rank, args.dist_url), flush=True) | |
| torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, | |
| world_size=args.world_size, rank=args.rank) | |
| torch.distributed.barrier() | |
| setup_for_distributed(args.rank == 0) | |
| if args.batch_size % args.world_size != 0: | |
| raise Exception("Batch size should be divisible by world size.") | |
| args.batch_size //= args.world_size | |
| args.eval_batch_size //= args.world_size | |
| def all_gather(data): | |
| """ | |
| Run all_gather on arbitrary picklable data (not necessarily tensors) | |
| Args: | |
| data: any picklable object | |
| Returns: | |
| list[data]: list of data gathered from each rank | |
| """ | |
| world_size = get_world_size() | |
| if world_size == 1: | |
| return [data] | |
| # serialized to a Tensor | |
| buffer = pickle.dumps(data) | |
| storage = torch.ByteStorage.from_buffer(buffer) | |
| tensor = torch.ByteTensor(storage).to("cuda") | |
| # obtain Tensor size of each rank | |
| local_size = torch.tensor([tensor.numel()], device="cuda") | |
| size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] | |
| dist.all_gather(size_list, local_size) | |
| size_list = [int(size.item()) for size in size_list] | |
| max_size = max(size_list) | |
| # receiving Tensor from all ranks | |
| # we pad the tensor because torch all_gather does not support | |
| # gathering tensors of different shapes | |
| tensor_list = [] | |
| for _ in size_list: | |
| tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) | |
| if local_size != max_size: | |
| padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") | |
| tensor = torch.cat((tensor, padding), dim=0) | |
| dist.all_gather(tensor_list, tensor) | |
| data_list = [] | |
| for size, tensor in zip(size_list, tensor_list): | |
| buffer = tensor.cpu().numpy().tobytes()[:size] | |
| data_list.append(pickle.loads(buffer)) | |
| return data_list | |