97 lines
2.7 KiB
Python
Executable File
97 lines
2.7 KiB
Python
Executable File
import pickle
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def get_world_size():
|
|
if not dist.is_available():
|
|
return 1
|
|
if not dist.is_initialized():
|
|
return 1
|
|
return dist.get_world_size()
|
|
|
|
|
|
def get_rank():
|
|
if not dist.is_available():
|
|
return 0
|
|
if not dist.is_initialized():
|
|
return 0
|
|
return dist.get_rank()
|
|
|
|
|
|
def is_main_process():
|
|
return get_rank() == 0
|
|
|
|
|
|
def synchronize():
|
|
"""
|
|
Helper function to synchronize (barrier) among all processes when
|
|
using distributed training
|
|
"""
|
|
if not dist.is_available():
|
|
return
|
|
if not dist.is_initialized():
|
|
return
|
|
world_size = dist.get_world_size()
|
|
if world_size == 1:
|
|
return
|
|
dist.barrier()
|
|
|
|
|
|
def _encode(encoded_data, data):
|
|
# gets a byte representation for the data
|
|
encoded_bytes = pickle.dumps(data)
|
|
# convert this byte string into a byte tensor
|
|
storage = torch.ByteStorage.from_buffer(encoded_bytes)
|
|
tensor = torch.ByteTensor(storage).to("cuda")
|
|
# encoding: first byte is the size and then rest is the data
|
|
s = tensor.numel()
|
|
assert s <= 255, "Can't encode data greater than 255 bytes"
|
|
# put the encoded data in encoded_data
|
|
encoded_data[0] = s
|
|
encoded_data[1: (s + 1)] = tensor
|
|
|
|
|
|
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.LongTensor([tensor.numel()]).to("cuda")
|
|
size_list = [torch.LongTensor([0]).to("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.ByteTensor(size=(max_size,)).to("cuda"))
|
|
if local_size != max_size:
|
|
padding = torch.ByteTensor(size=(max_size - local_size,)).to("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
|