mindspore.ops.AlltoAllV

class mindspore.ops.AlltoAllV(group=GlobalComm.WORLD_COMM_GROUP, block_size=1)[source]

AllToAllV which support uneven scatter and gather compared with AllToAll.

Note

Only support flatten tensor as input. input tensor should be flattened and concatenated before call this primitive.

Parameters

group (str, optional) – The communication group to work on. Default: GlobalComm.WORLD_COMM_GROUP, which means "hccl_world_group" in Ascend.
block_size (int, optional) – The basic units for scatter and gather numel by send_numel_list and recv_numel_list. Default: 1.

Inputs:

input_x (Tensor) - flatten tensor to scatter. The shape of tensor is $(x_{1})$ .
send_numel_list (Union[tuple[int], list[int], Tensor]) - split numel to scatter to different remote rank. The actual distributed data numel is $(s e n d_{n} u m e l_{l} i s t * b l o c k_{s} i z e * i n p u t_{x} . d t y p e)$ .
recv_numel_list (Union[tuple[int], list[int], Tensor]) - split numel to gather from different remote rank. The actual aggregated data numel is $(r e c v_{n} u m e l_{l} i s t * b l o c k_{s} i z e * i n p u t_{x} . d t y p e)$ .

Outputs:

Tensor, flattened and concatenated tensor gather from remote ranks. If gather result is empty, it will return a Tensor with value 0, which has no actual meaning.

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.

>>> from mindspore import ops
>>> import mindspore.nn as nn
>>> from mindspore.communication import init, get_rank
>>> from mindspore import Tensor
>>>
>>> init()
>>> rank = get_rank()
>>> class Net(nn.Cell):
...     def __init__(self):
...         super(Net, self).__init__()
...         self.all_to_all = ops.AlltoAllV()
...
...     def construct(self, x, send_numel_list, recv_numel_list):
...         return self.all_to_all(x, send_numel_list, recv_numel_list)
>>> send_numel_list = []
>>> recv_numel_list = []
>>> if rank == 0:
...    send_tensor = Tensor([0, 1, 2.])
...    send_numel_list = [1, 2]
...    recv_numel_list = [1, 2]
>>> elif rank == 1:
...    send_tensor = Tensor([3, 4, 5.])
...    send_numel_list = [2, 1]
...    recv_numel_list = [2, 1]
>>> net = Net()
>>> output = net(send_tensor, send_numel_list, recv_numel_list)
>>> print(output)
rank 0:
[0. 3. 4]
rank 1:
[1. 2. 5]