mindspore.ops.communication.all_to_all

View Source On AtomGit
mindspore.ops.communication.all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False)[source]

Scatter and gather list of tensor to/from all rank according to input/output tensor list.

Note

  • Tensor shape in output_tensor_list and input_tensor_list should be match across ranks.

  • Only support PyNative mode, Graph mode is not currently supported.

Parameters

  • output_tensor_list (Union[List(Tensor), List(Tuple(int))]) – List of tensors that indicate the gathered from remote ranks. If the function operates in-place. Otherwise, List of tensors or shape that indicate the gathered tensors shape from remote ranks.

  • input_tensor_list (List[Tensor]) – List of tensors to scatter to the remote rank.

  • 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

  • If the function operates in-place, return CommHandle.

  • If the function operates non in-place, return Tuple(Tensor, CommHandle). The first element stores the output result, and the second element is CommHandle.

Among them, when async_op is True, then CommHandle is an asynchronous working handle; When async_op is False, CommHandle will return None.

Raises

  • TypeError – If not all elements in input_tensor_list or output_tensor_list are Tensor.

  • TypeError – If tensors in input_tensor_list or output_tensor_list are not the same type.

  • TypeError – If group is not str or async_op is not bool.

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 mindspore as ms
>>> from mindspore.ops.communication import init_process_group, get_rank
>>> from mindspore.ops.communication import all_to_all
>>> from mindspore import Tensor
>>>
>>> init_process_group()
>>> this_rank = get_rank()
>>> if this_rank == 0:
...     send_tensor_list = [Tensor(1.), Tensor([[2, 3], [4, 5.]])]
...     recv_tensor_list = [Tensor((0), dtype=ms.float32), Tensor([0, 0.])]
>>> if this_rank == 1:
...     send_tensor_list = [Tensor([2, 2.]), Tensor([4, 5, 6, 7.])]
...     recv_tensor_list = [Tensor([[0, 0.],[0, 0]]), Tensor([0, 0, 0, 0.])]
>>> handle = all_to_all(recv_tensor_list, send_tensor_list)
>>> print(recv_tensor_list)
rank 0:
(Tensor(shape=[], dtype=Float32, value= 1),
 Tensor(shape=[2], dtype=Float32, value= [2.00000000e+00, 2.00000000e+00]))
rank 1:
(Tensor(shape=[2, 2], dtype=Float32, value=
[[2.00000000e+00, 3.00000000e+00],
 [4.00000000e+00, 5.00000000e+00]]),
 Tensor(shape=[4], dtype=Float32, value=[4.00000000e+00, 5.00000000e+00, 6.00000000e+00, 7.00000000e+00]))