Gather tensors and concatinate along the last dimension.
(input_)
| 54 | |
| 55 | |
| 56 | def _gather(input_): |
| 57 | """Gather tensors and concatinate along the last dimension.""" |
| 58 | group = get_model_parallel_group() |
| 59 | |
| 60 | # Bypass the function if we are using only 1 GPU. |
| 61 | if torch.distributed.get_world_size(group=group) == 1: |
| 62 | return input_ |
| 63 | |
| 64 | # Size and dimension. |
| 65 | last_dim = input_.dim() - 1 |
| 66 | rank = torch.distributed.get_rank(group=group) |
| 67 | world_size = torch.distributed.get_world_size(group=group) |
| 68 | |
| 69 | tensor_list = [torch.empty_like(input_) for _ in range(world_size)] |
| 70 | tensor_list[rank] = input_ |
| 71 | torch.distributed.all_gather(tensor_list, input_, group=group) |
| 72 | |
| 73 | # Note: torch.cat already creates a contiguous tensor. |
| 74 | output = torch.cat(tensor_list, dim=last_dim).contiguous() |
| 75 | |
| 76 | return output |
| 77 | |
| 78 | |
| 79 | class _CopyToModelParallelRegion(torch.autograd.Function): |
no test coverage detected