mindspore.communication.comm_func.all_to_all_with_output_shape

View Source On Gitee
mindspore.communication.comm_func.all_to_all_with_output_shape(output_shape_list, input_tensor_list, group=None)[source]

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

Note

tensor shape in output_shape_list and input_tensor_list should be match across ranks. Only support PyNative mode, Graph mode is not currently supported.

Parameters

  • output_shape_list (Union[Tuple(Tensor), List(Tensor), Tuple(Tuple(int))]) – List of shape that indicate the gathered tensors shape from remote ranks.

  • input_tensor_list (Union[Tuple(Tensor), 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" on Ascend, "nccl_world_group" on GPU.

Returns

Tuple(Tensor), the tensors gathered from remote ranks.

Raises

  • TypeError – If input_tensor_list is not list of tensors.

  • TypeError – If output_shape_list is not list of tuple or tensors.

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

Supported Platforms:

Ascend

Examples

Note

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

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

This example should be run with 2 devices.

>>> import numpy as np
>>> import mindspore
>>> from mindspore.communication import init, get_rank, get_group_size
>>> from mindspore.communication.comm_func import all_to_all_with_output_shape
>>> from mindspore import Tensor
>>> from mindspore.ops import zeros
>>>
>>> init()
>>> this_rank = get_rank()
>>> if this_rank == 0:
>>>     send_tensor_list = [Tensor(1.), Tensor([[2, 3], [4, 5.]])]
>>>     recv_tensor_list = [(), (2,)]
>>> if this_rank == 1:
>>>     send_tensor_list = [Tensor([2, 2.]), Tensor([4, 5, 6, 7.])]
>>>     recv_tensor_list = [(2, 2), (4,)]
>>> output = all_to_all_with_output_shape(recv_tensor_list, send_tensor_list)
>>> print(output)
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]))