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