Break a tensor into equal 1D chunks.
(tensor)
| 110 | |
| 111 | |
| 112 | def split_tensor_into_1d_equal_chunks(tensor): |
| 113 | """Break a tensor into equal 1D chunks.""" |
| 114 | data = tensor.view(-1) |
| 115 | partition_size = torch.numel(data) // get_tensor_model_parallel_world_size() |
| 116 | start_index = partition_size * get_tensor_model_parallel_rank() |
| 117 | end_index = start_index + partition_size |
| 118 | return data[start_index:end_index] |
| 119 | |
| 120 | |
| 121 | def gather_split_1d_tensor(tensor): |
no test coverage detected