Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,251 Bytes
2d9a728 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 |
"""
Copyright (c) Microsoft Corporation.
Licensed under the MIT license.
distributed API using Horovod
Modified from OpenNMT's native pytorch distributed utils
(https://github.com/OpenNMT/OpenNMT-py)
"""
import math
import pickle
import torch
import torch.distributed as dist
from time import time
from torch.autograd import Function
from torch.utils.data.distributed import DistributedSampler
class ddp_allgather_with_grads(Function):
@staticmethod
def forward(ctx, x):
tmp_input = x.cuda()
size = torch.tensor(tmp_input.shape[0]).cuda()
size_list = [torch.zeros_like(size) for i in range(dist.get_world_size())]
dist.all_gather(size_list, size)
max_size = max(size_list).item()
padding_size = max_size - size
if padding_size > 0 :
padding_tensor = torch.zeros(padding_size,*tmp_input.shape[1:]).to(tmp_input)
tmp_input = torch.cat((tmp_input, padding_tensor), dim = 0)
tmp_list = [torch.zeros_like(tmp_input) for i in range(dist.get_world_size())]
dist.all_gather(tmp_list, tmp_input)
ctx.size = size_list
output = []
for t, s in zip(tmp_list, size_list):
output.append(t[:s])
output = torch.cat(output,dim=0)
output.requires_grad = True
return output
@staticmethod
def backward(ctx, grad_output):
grad_x = None
if grad_output is not None:
grad_output.detach()
#grad_x = grad_output.chunk(dist.get_world_size(),dim=0)[dist.get_rank()]
start = sum(ctx.size[:dist.get_rank()])
end = start + ctx.size[dist.get_rank()]
grad_x = grad_output[start:end]
return grad_x
def ddp_allgather(input):
tmp_input = input.cuda()
size = torch.tensor(tmp_input.shape[0]).cuda()
size_list = [torch.zeros_like(size) for i in range(dist.get_world_size())]
dist.all_gather(size_list, size)
max_size = max(size_list).item()
padding_size = max_size - size
if padding_size > 0 :
padding_tensor = torch.zeros(padding_size,*tmp_input.shape[1:]).to(tmp_input)
tmp_input = torch.cat((tmp_input, padding_tensor), dim = 0)
tmp_list = [torch.zeros_like(tmp_input) for i in range(dist.get_world_size())]
dist.all_gather(tmp_list, tmp_input)
output = []
for t, s in zip(tmp_list, size_list):
output.append(t[:s])
output = torch.cat(output,dim=0)
return output
def _encode(enc, max_size, use_max_size=False):
enc_size = len(enc)
enc_byte = max(math.floor(math.log(max_size, 256)+1), 1)
if use_max_size:
# this is used for broadcasting
buffer_ = torch.cuda.ByteTensor(max_size+enc_byte)
else:
buffer_ = torch.cuda.ByteTensor(enc_size+enc_byte)
remainder = enc_size
for i in range(enc_byte):
base = 256 ** (enc_byte-i-1)
buffer_[i] = remainder // base
remainder %= base
buffer_[enc_byte:enc_byte+enc_size] = torch.ByteTensor(list(enc))
return buffer_, enc_byte
def _decode(buffer_, enc_byte):
size = sum(256 ** (enc_byte-i-1) * buffer_[i].item()
for i in range(enc_byte))
bytes_list = bytes(buffer_[enc_byte:enc_byte+size].tolist())
shift = size + enc_byte
return bytes_list, shift
_BUFFER_SIZE = 4096
def all_gather_list(data):
"""Gathers arbitrary data from all nodes into a list."""
enc = pickle.dumps(data)
enc_size = len(enc)
max_size = ddp_allgather(torch.tensor([enc_size]).cuda()).max().item()
in_buffer, enc_byte = _encode(enc, max_size)
out_buffer = ddp_allgather(in_buffer[:enc_byte+enc_size])
results = []
for _ in range(dist.get_world_size()):
bytes_list, shift = _decode(out_buffer, enc_byte)
out_buffer = out_buffer[shift:]
result = pickle.loads(bytes_list)
results.append(result)
return results
def any_broadcast(data, root_rank):
"""broadcast arbitrary data from root_rank to all nodes."""
enc = pickle.dumps(data)
max_size = ddp_allgather(torch.tensor([len(enc)]).cuda()).max().item()
buffer_, enc_byte = _encode(enc, max_size, use_max_size=True)
dist.broadcast(buffer_, root_rank)
bytes_list, _ = _decode(buffer_, enc_byte)
result = pickle.loads(bytes_list)
return result
class DistributedSampler_wopadding(DistributedSampler):
def __iter__(self):
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type]
else:
indices = list(range(len(self.dataset))) # type: ignore[arg-type]
if self.drop_last:
indices = indices[:self.total_size]
#assert len(indices) == self.total_size
# subsample
indices = indices[self.rank:len(indices):self.num_replicas]
# assert len(indices) == self.num_samples
return iter(indices)
class GatherLayer(torch.autograd.Function):
"""
Gather tensors from all workers with support for backward propagation:
This implementation does not cut the gradients as torch.distributed.all_gather does.
"""
@staticmethod
def forward(ctx, x):
output = [
torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())
]
torch.distributed.all_gather(output, x)
return tuple(output)
@staticmethod
def backward(ctx, *grads):
all_gradients = torch.stack(grads)
torch.distributed.all_reduce(all_gradients)
return all_gradients[torch.distributed.get_rank()]
def all_gather_with_grad(tensors):
"""
Performs all_gather operation on the provided tensors.
Graph remains connected for backward grad computation.
"""
# Queue the gathered tensors
world_size = torch.distributed.get_world_size()
# There is no need for reduction in the single-proc case
if world_size == 1:
return tensors
# tensor_all = GatherLayer.apply(tensors)
tensor_all = GatherLayer.apply(tensors)
return torch.cat(tensor_all, dim=0) |