All-reduce the input tensor across model parallel group.
(input_)
| 618 | |
| 619 | |
| 620 | def _reduce(input_): |
| 621 | """All-reduce the input tensor across model parallel group.""" |
| 622 | group = g_mpu.get_model_parallel_group() |
| 623 | |
| 624 | # Bypass the function if we are using only 1 GPU. |
| 625 | if dist.get_world_size(group=group) == 1: |
| 626 | return input_ |
| 627 | |
| 628 | # All-reduce. |
| 629 | dist.all_reduce(input_, group=group) |
| 630 | |
| 631 | return input_ |
| 632 | |
| 633 | |
| 634 | def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split_chunks=False): |
no test coverage detected