Opposite of above function, gather values from model parallel ranks.
(tensor)
| 119 | |
| 120 | |
| 121 | def gather_split_1d_tensor(tensor): |
| 122 | """Opposite of above function, gather values from model parallel ranks.""" |
| 123 | world_size = get_tensor_model_parallel_world_size() |
| 124 | numel = torch.numel(tensor) |
| 125 | numel_gathered = world_size * numel |
| 126 | gathered = torch.empty( |
| 127 | numel_gathered, |
| 128 | dtype=tensor.dtype, |
| 129 | device=torch.cuda.current_device(), |
| 130 | requires_grad=False, |
| 131 | ) |
| 132 | chunks = [gathered[i * numel : (i + 1) * numel] for i in range(world_size)] |
| 133 | torch.distributed.all_gather( |
| 134 | chunks, tensor, group=get_tensor_model_parallel_group() |
| 135 | ) |
| 136 | return gathered |
| 137 | |
| 138 | |
| 139 | class CudaRNGStatesTracker: |
no test coverage detected