AllToAllVV2Operation
产品支持情况
硬件型号 |
是否支持 |
---|---|
|
x |
|
√ |
|
x |
|
√ |
|
x |
功能说明
向通信域内所有通信卡发送数据(数据量可以通过参数定制),并从所有通信卡接收数据(数据量可以通过参数定制)。

1 2 3 4 5 |
# 计算goldTensor gold_outtensor = [] for j in range(len(recvout[rank])): gold_outtensor.append(tensorafters[j][senddisp[j][rank]:sendcount[j][rank] + senddisp[j][rank]]) gold_outtensor = [i for arr in gold_outtensor for i in arr] |
算子上下文

模型传输一个int8的量化输入tensor数据X,首先使用AlltoAll进行通信发送到各个节点上(这里使用int8进行通信,提升了通信速度)
然后使用reduce的sum操作对x进行求和,将int8数据反量化为float16, 最后使用AllGather进行通信将计算结果传输到各个节点上。(需求图中AllGather前后的量化,有可能会有精度损失,当前未实现。)
使用场景
用于将数据发送到各个节点上,多对多。AllToAll是对Allgather的扩展,相比于Allgather,AllToAll不同的节点从某一节点收集到的数据是不同的。

应用于模型并行;模型并行里的矩阵转置;数据并行到模型并行的矩阵转置。
使用示例:
>>> rank0 input tensor([[0,1,2,3], [4,5,6,7]], device='npu:0') shape[2,4] >>> rank0 sendcount tensor([2, 4], device='npu:0') shape[2] >>> rank0 sdispls tensor([0, 2]) >>> rank0 recvCounts tensor([2, 2]) >>> rank0 rdispls tensor([0, 2]) >>> rank0 tensorForInferShape tensor([3, 3, 3, 3]) >>> rank1 input tensor([[0, 1, 2], [3, 4, 6]], device='npu:1') shape[2,3] >>> rank1 sendcount tensor([2, 1], device='npu:1') shape[2] >>> rank1 sdispls tensor([0, 2]) >>> rank1 recvCounts tensor([4, 1]) >>> rank1 rdispls tensor([0, 4]) >>> rank1 tensorForInferShape tensor([3, 3, 3, 3, 3]) >>> rank0 output tensor([[0, 1, 0, 1]], device='npu:0') shape[1,4] >>> rank1 output tensor([[2, 3, 4, 5, 2]], device='npu:1') shape[1,5]
定义
1 2 3 4 5 6 7 8 9 10 11 |
struct AllToAllVV2Param { int rank = -1; int rankSize = 0; int rankRoot = 0; std::string backend = "hccl"; HcclComm hcclComm = nullptr; CommMode commMode = COMM_MULTI_PROCESS; std::string rankTableFile; std::string commDomain; uint8_t rsv[64] = {0}; }; |
参数列表
输入
参数 |
维度 |
数据类型 |
格式 |
是否必选 |
描述 |
---|---|---|---|---|---|
x |
[dim_0, dim_1, ..., dim_n] |
"hccl": float16/int8 |
ND |
是 |
输入tensor。 |
sendCount |
1[rankSize] |
int64 |
ND |
否 |
表示发送数据量的数组,为host侧tensor。 例如,若发送的数据类型为float16,sendCounts[i] = n 表示本rank发给rank_i n个float16数据。 |
sdispls |
1[rankSize] |
int64 |
ND |
否 |
表示发送偏移量的数组,为host侧tensor。sdispls[i] = n表示本rank从相对于输入起始位置的偏移量为n的位置开始发送数据给rank_i。 |
recvCounts |
1[rankSize] |
int64 |
ND |
否 |
表示接收数据量的数组,为host侧tensor。例如,若发送的数据类型为float16,recvCounts[i] = n 表示本rank从rank_i收到n个float16数据。 |
rdispls |
1[rankSize] |
int64 |
ND |
否 |
表示接收偏移量的数组,为host侧tensor。rdispls[i] = n表示本rank从相对于输入起始位置的偏移量为n的位置开始接收rank_i的数据。 |
tensorForInferShape |
[recvCountsSum] |
int8 |
ND |
否 |
shape为recvCounts的所有元素之和,用于infer shape。 |
输出
参数 |
维度 |
数据类型 |
格式 |
是否必选 |
描述 |
---|---|---|---|---|---|
output |
[1,recvCountsSum] |
"hccl": float16/int8 |
ND |
否 |
输出tensor,最后一维的shape为参数recvCounts的所有元素之和。数据类型和输入相同。 |
约束说明
- 多卡场景要指定rank。
- sendCounts、recvCounts、sdispls、rdispls均可看作长度为ranksize的一维数组,且数组中元素值都大于等于0。
- sendCounts、recvCounts数组元素之和不能溢出int64,对于rdispls中的任一元素,recvCounts[i] + rdispls[i] 不能大于recvCountsSum(即output的最后一维),对于sdispls中的任一元素,sendCounts[i] + sdispls[i]不能大于输入tensor的数据量。如输入tensor shape为[3, 4, 5],数据量为3*4*5 = 60,在进行AllToAllV时会被视为shape为[60]的tensor进行计算。
- rank、rankSize、rankRoot需满足以下条件。
- 0 ≤ rank < rankSize
- 0 ≤ rankRoot < rankSize