mindspore.mint.distributed.all_to_all_single

mindspore.mint.distributed.all_to_all_single(output, input, output_split_sizes=None, input_split_sizes=None, group=None, async_op=False)[source]

Exchanges data among all processes (ranks) in the communication group. The function first splits input along dimension 0 according to input_split_sizes, then sends each chunk in order to the corresponding rank. Meanwhile, each rank receives data from all ranks in the group and concatenates the received data along dimension 0 according to output_split_sizes into a single output tensor.

Note

Only supports PyNative mode. Graph mode is not currently supported.

Parameters:

output (Tensor) – The output tensor is gathered and concatenated from remote ranks.
input (Tensor) – The input tensor to be scattered to remote rank.
output_split_sizes (Union(Tuple(int), List(int)), optional) – The output split size at dim 0. Default: None, which means equally split by world_size.
input_split_sizes (Union(Tuple(int), List(int)), optional) – The input split size at dim 0. Default: None, which means equally split by world_size.
group (str, optional) – The communication group to work on. Default: None, which means "hccl_world_group" in Ascend.
async_op (bool, optional) – Whether this operator should be an async operator. Default: False .

Returns:

CommHandle. CommHandle is an async work handle, if async_op is set to True. CommHandle will be None, when async_op is False.

Raises:

TypeError – If input or output is not tensor, group is not a str, or async_op is not bool.
ValueError – When input_split_sizes is empty, input dim 0 can not be divided by world_size.
ValueError – When output_split_sizes is empty, output dim 0 can not be divided by world_size.

Supported Platforms:: Ascend

Examples

Note

Before running the following examples, you need to configure the communication environment variables.

For Ascend devices, it is recommended to use the msrun startup method without any third-party or configuration file dependencies. Please see the msrun startup for more details.

This example should be run with 2 devices.

>>> import numpy as np
>>> from mindspore.mint.distributed import init_process_group, get_rank
>>> from mindspore.mint.distributed import all_to_all_single
>>> from mindspore import Tensor
>>>
>>> init_process_group()
>>> this_rank = get_rank()
>>> if this_rank == 0:
...     output = Tensor(np.zeros([3, 3]).astype(np.float32))
...     tensor = Tensor([[0, 1, 2.], [3, 4, 5], [6, 7, 8]])
...     result = all_to_all_single(output, tensor, [2, 1], [2, 1])
...     print(output)
>>> if this_rank == 1:
...     output = Tensor(np.zeros([2, 3]).astype(np.float32))
...     tensor = Tensor([[9, 10., 11], [12, 13, 14]])
...     result = all_to_all_single(output, tensor, [1, 1], [1, 1])
...     print(output)
rank 0:
[[ 0.  1.  2.]
 [ 3.  4.  5.]
 [ 9. 10. 11.]]
rank 1:
[[ 6.  7.  8.]
 [12. 13. 14.]]