Split the tensor along its last dimension and keep the corresponding slice.
(input_)
| 34 | |
| 35 | |
| 36 | def _split(input_): |
| 37 | """Split the tensor along its last dimension and keep the |
| 38 | corresponding slice.""" |
| 39 | group = get_model_parallel_group() |
| 40 | |
| 41 | # Bypass the function if we are using only 1 GPU. |
| 42 | if torch.distributed.get_world_size(group=group) == 1: |
| 43 | return input_ |
| 44 | |
| 45 | # Split along last dimension. |
| 46 | world_size = torch.distributed.get_world_size(group=group) |
| 47 | input_list = split_tensor_along_last_dim(input_, world_size) |
| 48 | |
| 49 | # Note: torch.split does not create contiguous tensors by default. |
| 50 | rank = torch.distributed.get_rank(group=group) |
| 51 | output = input_list[rank].contiguous() |
| 52 | |
| 53 | return output |
| 54 | |
| 55 | |
| 56 | def _gather(input_): |
no test coverage detected