All-reduce the the input tensor across model parallel group.
(input_)
| 24 | |
| 25 | |
| 26 | def _reduce(input_): |
| 27 | """All-reduce the the input tensor across model parallel group.""" |
| 28 | |
| 29 | # Bypass the function if we are using only 1 GPU. |
| 30 | if get_tensor_model_parallel_world_size() == 1: |
| 31 | return input_ |
| 32 | |
| 33 | # All-reduce. |
| 34 | torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group()) |
| 35 | |
| 36 | return input_ |
| 37 | |
| 38 | |
| 39 | def _split(input_): |
no test coverage detected