mindspore.communication.comm_func.all_to_all_v_c

mindspore.communication.comm_func.all_to_all_v_c(output, input, send_count_matrix, group=None, async_op=False)[源代码]

根据用户输入的切分大小，把输入tensor切分后，发送到其他的设备上，并从其他设备接收切分块，然后合并到一个输出tensor中。

说明

仅支持PyNative模式，目前不支持Graph模式。

参数：

output (Tensor) - 表示从远端收集的tensor结果。
input (Tensor) - 要发送到远端设备的tensor。
send_count_matrix (list[int]) - 所有rank的收发大小列表。 \(\text{send_count_matrix}[i*\text{rank_size}+j]\) 表示rank i发给rank j的数据量，基本单位是输入的第一个维度尺寸。其中， rank_size 表示通信组大小。
group (str, 可选) - 通信组名称。如果为 None ，在Ascend平台中表示为 "hccl_world_group" 。默认值： None 。
async_op (bool, 可选) - 本算子是否是异步算子。默认值： False 。

返回：

CommHandle。若 async_op 是True，则CommHandle是一个异步工作句柄；若 async_op 是False，则CommHandle将返回None。

异常：

TypeError - input 或者 output 不是tensor类型， group 不是str， async_op 不是bool。

支持平台：

Ascend

样例：

说明

运行以下样例之前，需要配置好通信环境变量。

针对Ascend/GPU/CPU设备，推荐使用msrun启动方式，无第三方以及配置文件依赖。详见 msrun启动。

该样例需要在2卡环境下运行。

>>> import numpy as np
>>> import mindspore
>>> from mindspore.mint.distributed import init_process_group, get_rank
>>> from mindspore.communication.comm_func import all_to_all_v_c
>>> from mindspore import Tensor
>>> from mindspore.ops import zeros
>>>
>>> init_process_group()
>>> this_rank = get_rank()
>>> if this_rank == 0:
...     output = Tensor(np.zeros([3]).astype(np.float32))
...     tensor = Tensor([0, 1, 2.]) * this_rank
...     result = all_to_all_v_c(output, tensor, [0, 3, 3, 0])
...     print(output)
>>> if this_rank == 1:
...     output = Tensor(np.zeros([3]).astype(np.float32))
...     tensor = Tensor([0, 1, 2.]) * this_rank
...     result = all_to_all_v_c(output, tensor, [0, 3, 3, 0])
...     print(output)
rank 0:
[0. 1. 2]
rank 1:
[0. 0. 0]