mindspore.mint.distributed.all_to_all

查看源文件
mindspore.mint.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False)[源代码]

根据用户输入的张量列表，将对应的张量发送到远端设备，并从其他设备接收张量，返回一个接收的张量列表。

说明

  • 各个设备之间发送和接收的张量形状需要互相匹配。

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

参数：

  • output_tensor_list (List[Tensor]) - 包含接收张量的列表。

  • input_tensor_list (List[Tensor]) - 包含发送到其他设备张量的列表。

  • group (str, 可选) - 通信组名称，如果为 None ，Ascend平台表示为 "hccl_world_group" 。 默认值： None 。

  • async_op (bool, 可选) - 本算子是否是异步算子。默认值： False 。

返回：

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

异常：

  • TypeError - input_tensor_list 和 output_tensor_list 中不全是张量类型。

  • TypeError - input_tensor_list 和 output_tensor_list 中张量的数据类型不全部一致。

  • TypeError - group 不是str， async_op 不是bool。

支持平台：

Ascend

样例：

说明

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

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

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

>>> import mindspore as ms
>>> from mindspore.mint.distributed import init_process_group, get_rank
>>> from mindspore.mint.distributed 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]))