alltoall¶

paddle.distributed. alltoall ( in_tensor_list, out_tensor_list, group=None, use_calc_stream=True ) [源代码] ¶

将in_tensor_list里面的tensors按照卡数均分并按照卡的顺序分发到所有参与的卡并将结果tensors汇总到out_tensor_list。如下图所示，GPU0卡的in_tensor_list会按照两张卡拆分成0_0和0_1， GPU1卡的in_tensor_list同样拆分成1_0和1_1，经过alltoall算子后, GPU0卡的0_0会发送给GPU0，GPU0卡的0_1会发送给GPU1，GPU1卡的1_0会发送给GPU0，GPU1卡的1_1会发送给GPU1，所以GPU0卡的out_tensor_list包含0_0和1_0， GPU1卡的out_tensor_list包含0_1和1_1。

参数¶

in_tensor_list (list) - 包含所有输入Tensors的一个列表。在列表里面的所有元素都必须是一个Tensor，Tensor的数据类型必须是float16、float32、 float64、int32、int64。

out_tensor_list (Tensor) - 包含所有输出Tensors的一个列表。在列表里面的所有元素数据类型要和输入的Tensors数据类型一致。

group (Group, 可选) - new_group返回的Group实例，或者设置为None表示默认地全局组。默认值：None。

use_calc_stream (bool，可选) - 标识使用计算流还是通信流。默认值：True。

返回¶

无

代码示例¶

          # required: distributed
import numpy as np
import paddle
from paddle.distributed import init_parallel_env

init_parallel_env()
out_tensor_list = []
if paddle.distributed.ParallelEnv().rank == 0:
    np_data1 = np.array([[1, 2, 3], [4, 5, 6]])
    np_data2 = np.array([[7, 8, 9], [10, 11, 12]])
else:
    np_data1 = np.array([[13, 14, 15], [16, 17, 18]])
    np_data2 = np.array([[19, 20, 21], [22, 23, 24]])
data1 = paddle.to_tensor(np_data1)
data2 = paddle.to_tensor(np_data2)
paddle.distributed.alltoall([data1, data2], out_tensor_list)
# out for rank 0: [[[1, 2, 3], [4, 5, 6]], [[13, 14, 15], [16, 17, 18]]]
# out for rank 1: [[[7, 8, 9], [10, 11, 12]], [[19, 20, 21], [22, 23, 24]]]

         