alltoall¶
将 in_tensor_list 中的一组 tensor 分发到每个进程,随后在每个进程上将分发结果聚合到 out_tensor_list。
如下图所示,2 个 GPU 分别开启 1 个进程,rank=0 的进程的 in_tensor_list 包含 0_0 和 0_1 两个 tensor,rank=1 的进程的 in_tensor_list 包含 1_0 和 1_1 两个 tensor。 操作后,rank=0 的进程的 out_tensor_list 会包含 0_0 和 1_0 两个 tensor,rank=1 的进程的 out_tensor_list 会包含 0_0 和 1_1 两个 tensor。
简单来说,该操作类似于 scatter + gather。更直观地,如果将全部进程上的数据看作一个矩阵,该操作类似于对矩阵进行转置。
参数¶
in_tensor_list (List[Tensor]) - 输入的 tensor 列表。支持的数据类型包括:float16、float32、float64、int32、int64、int8、uint8、bool、bfloat16。
out_tensor_list (List[Tensor]) - 用于保存操作结果的 tensor 列表。其中每个 tensor 的数据类型必须与输入的 tensor 保持一致。
group (Group,可选) - 执行该操作的进程组实例(通过
new_group
创建)。默认为 None,即使用全局默认进程组。sync_op (bool,可选) - 该操作是否为同步操作。默认为 True,即同步操作。
返回¶
无返回值。
代码示例¶
# required: distributed
import paddle
import paddle.distributed as dist
dist.init_parallel_env()
out_tensor_list = []
if dist.get_rank() == 0:
data1 = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])
data2 = paddle.to_tensor([[7, 8, 9], [10, 11, 12]])
else:
data1 = paddle.to_tensor([[13, 14, 15], [16, 17, 18]])
data2 = paddle.to_tensor([[19, 20, 21], [22, 23, 24]])
dist.alltoall([data1, data2], out_tensor_list)
print(out_tensor_list)
# [[[1, 2, 3], [4, 5, 6]], [[13, 14, 15], [16, 17, 18]]] (2 GPUs, out for rank 0)
# [[[7, 8, 9], [10, 11, 12]], [[19, 20, 21], [22, 23, 24]]] (2 GPUs, out for rank 1)