Split the tensor along its last dimension and keep the corresponding slice.
(input_)
| 653 | |
| 654 | |
| 655 | def _split(input_): |
| 656 | """Split the tensor along its last dimension and keep the |
| 657 | corresponding slice.""" |
| 658 | group = g_mpu.get_model_parallel_group() |
| 659 | |
| 660 | # Bypass the function if we are using only 1 GPU. |
| 661 | if dist.get_world_size(group=group) == 1: |
| 662 | return input_ |
| 663 | |
| 664 | # Split along last dimension. |
| 665 | world_size = dist.get_world_size(group=group) |
| 666 | input_list = split_tensor_along_last_dim(input_, world_size) |
| 667 | |
| 668 | # Note: torch.split does not create contiguous tensors by default. |
| 669 | rank = dist.get_rank(group=group) |
| 670 | output = input_list[rank].contiguous() |
| 671 | |
| 672 | return output |
| 673 | |
| 674 | |
| 675 | def _gather(input_): |
no test coverage detected