mindspore.ops.communication.all_to_all

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

根据输入/输出张量列表，在所有rank之间分散和收集张量列表。

说明

  • output_tensor_list 和 input_tensor_list 中的张量shape应在rank之间匹配。

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

参数：

  • output_tensor_list (Union[List(Tensor), List(Tuple(int))]) - 如果函数以in-place模式运行，则表示从远程rank收集的张量列表。否则，表示从远程rank收集的张量或shape列表。

  • input_tensor_list (List[Tensor]) - 要分散到远程rank的张量列表。

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

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

返回：

  • 若函数以in-place模式运行，返回CommHandle。

  • 若函数以非in-place模式运行，返回Tuple(Tensor, CommHandle)。第一个元素存储输出结果，第二个元素是CommHandle。

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

异常：

  • TypeError - 如果 input_tensor_list 或 output_tensor_list 中的元素不全是Tensor。

  • TypeError - 如果 input_tensor_list 或 output_tensor_list 中的张量不是同一类型。

  • TypeError - 如果 group 不是str或 async_op 不是bool。

支持平台：

Ascend

样例：

说明

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

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

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

>>> 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]))