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