| 246 | return int(partition_size) |
| 247 | |
| 248 | def get_full_inputs(tensors): |
| 249 | inputs=[] |
| 250 | for i in range(int(len(tensors)/2)-1): |
| 251 | item = tensors[2 * i] |
| 252 | size = tensors[2* i + 1] |
| 253 | partition_size = item.numel() |
| 254 | tensor_size = partition_size * mp_size |
| 255 | flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=item.device) |
| 256 | partitions=[] |
| 257 | for i in range(mp_size): |
| 258 | part_i = flat_tensor.narrow(0, partition_size * i , partition_size) |
| 259 | if i == mp_rank: |
| 260 | part_i.copy_(item) |
| 261 | partitions.append(part_i) |
| 262 | dist.all_gather(partitions,partitions[mp_rank], group=mp_group) |
| 263 | input_tensor = flat_tensor.view(list(size.numpy())) |
| 264 | item.data=input_tensor.data |
| 265 | |
| 266 | inputs.append(item) |
| 267 | inputs.append(tensors[-2]) |
| 268 | |
| 269 | return tuple(inputs) |
| 270 | |
| 271 | |
| 272 | |