Returns: list[int]: size of the tensor, on each rank Tensor: padded tensor that has the max size
(tensor, group)
| 289 | |
| 290 | |
| 291 | def _pad_to_largest_tensor(tensor, group): |
| 292 | """ |
| 293 | Returns: |
| 294 | list[int]: size of the tensor, on each rank |
| 295 | Tensor: padded tensor that has the max size |
| 296 | """ |
| 297 | world_size = dist.get_world_size(group=group) |
| 298 | assert ( |
| 299 | world_size >= 1 |
| 300 | ), 'comm.gather/all_gather must be called from ranks within the group!' |
| 301 | local_size = torch.tensor([tensor.numel()], |
| 302 | dtype=torch.int64, |
| 303 | device=tensor.device) |
| 304 | size_list = [ |
| 305 | torch.zeros([1], dtype=torch.int64, device=tensor.device) |
| 306 | for _ in range(world_size) |
| 307 | ] |
| 308 | dist.all_gather(size_list, local_size, group=group) |
| 309 | size_list = [int(size.item()) for size in size_list] |
| 310 | |
| 311 | max_size = max(size_list) |
| 312 | |
| 313 | # we pad the tensor because torch all_gather does not support |
| 314 | # gathering tensors of different shapes |
| 315 | if local_size != max_size: |
| 316 | padding = torch.zeros((max_size - local_size, ), |
| 317 | dtype=torch.uint8, |
| 318 | device=tensor.device) |
| 319 | tensor = torch.cat((tensor, padding), dim=0) |
| 320 | return size_list, tensor |
| 321 | |
| 322 | |
| 323 | def all_gather(data, group=None): |
no test coverage detected
searching dependent graphs…