Gather tensors and concatenate along the last dimension.
(input_)
| 673 | |
| 674 | |
| 675 | def _gather(input_): |
| 676 | """Gather tensors and concatenate along the last dimension.""" |
| 677 | group = g_mpu.get_model_parallel_group() |
| 678 | |
| 679 | # Bypass the function if we are using only 1 GPU. |
| 680 | if dist.get_world_size(group=group) == 1: |
| 681 | return input_ |
| 682 | |
| 683 | # Size and dimension. |
| 684 | last_dim = input_.dim() - 1 |
| 685 | rank = dist.get_rank(group=group) |
| 686 | world_size = dist.get_world_size(group=group) |
| 687 | |
| 688 | tensor_list = [torch.empty_like(input_) for _ in range(world_size)] |
| 689 | tensor_list[rank] = input_ |
| 690 | dist.all_gather(tensor_list, input_, group=group) |
| 691 | |
| 692 | # Note: torch.cat already creates a contiguous tensor. |
| 693 | output = torch.cat(tensor_list, dim=last_dim).contiguous() |
| 694 | |
| 695 | return output |
| 696 | |
| 697 | |
| 698 | class _CopyToModelParallelRegion(torch.autograd.Function): |
no test coverage detected