r"""Each process scatter input tensor to all processes and return gathered tensor. Args: inp: Input tensor. group: The process group to work on. The default group is WORLD which means all processes available. You can use a list of process ranks to create
(
inp: Tensor,
group: Optional[Group] = WORLD,
device: Optional[str] = None,
split_axis: int = 0,
concat_axis: int = 0,
)
| 693 | |
| 694 | |
| 695 | def all_to_all( |
| 696 | inp: Tensor, |
| 697 | group: Optional[Group] = WORLD, |
| 698 | device: Optional[str] = None, |
| 699 | split_axis: int = 0, |
| 700 | concat_axis: int = 0, |
| 701 | ) -> Tensor: |
| 702 | r"""Each process scatter input tensor to all processes and return gathered tensor. |
| 703 | |
| 704 | Args: |
| 705 | inp: Input tensor. |
| 706 | group: The process group to work on. |
| 707 | The default group is WORLD which means all processes available. |
| 708 | You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5]. |
| 709 | device: The specific device to execute this operator. |
| 710 | None default device means the device of inp will be used. |
| 711 | Specify "gpu0:1" to execute this operator on diffrent cuda stream, |
| 712 | 1 is stream id, and default stream id is 0. |
| 713 | split_axis: The axis that collectivecomm will split data |
| 714 | the default axis is 0 |
| 715 | |
| 716 | Returns: |
| 717 | Result tensor. |
| 718 | |
| 719 | Examples: |
| 720 | |
| 721 | .. code-block:: |
| 722 | |
| 723 | input = Tensor([0 1]) + rank*2 |
| 724 | # Rank 0 # input: Tensor([0 1]) |
| 725 | # Rank 1 # input: Tensor([2 3]) |
| 726 | output = all_to_all(input) |
| 727 | # Rank 0 # output: Tensor([0 2]) |
| 728 | # Rank 1 # output: Tensor([1 3]) |
| 729 | |
| 730 | input = Tensor([0 1]) + rank*2 |
| 731 | group = Group([1, 0]) |
| 732 | output = all_to_all(input, group) |
| 733 | # Rank 0 # output: Tensor([0 3]) |
| 734 | # Rank 1 # output: Tensor([2 1]) |
| 735 | """ |
| 736 | group_size = group.size if group is not None else 1 |
| 737 | assert ( |
| 738 | list(inp._tuple_shape)[split_axis] % group_size == 0 |
| 739 | ), "current axis: {} can't devided by group size".format(split_axis) |
| 740 | origin_shape = inp._tuple_shape |
| 741 | if split_axis != 0: |
| 742 | k_new_shape = list(inp._tuple_shape) |
| 743 | k_new_shape[split_axis] //= group_size |
| 744 | k_new_shape[0] *= group_size |
| 745 | new_shape = list(inp._tuple_shape) |
| 746 | new_shape[split_axis] //= group_size |
| 747 | new_shape.insert(split_axis, group_size) |
| 748 | index = ( |
| 749 | [split_axis] |
| 750 | + [_ for _ in range(0, split_axis)] |
| 751 | + [_ for _ in range(split_axis + 1, inp.ndim + 1)] |
| 752 | ) |