AllToAllOperation

功能

向通信域内所有通信卡发送相同数据量（输入切分成rankSize份）的数据，并从所有通信卡接收相同数据量的数据。

算子上下文

图1 通信INT8量化中AllToAll的位置

模型传输一个int8的量化输入tensor数据X，首先使用AlltoAll进行通信发送到各个节点上（这里使用int8进行通信，提升了通信速度）。

然后使用reduce的sum操作对x进行求和，将int8数据反量化为float16，最后使用AllGather进行通信将计算结果传输到各个节点上。（需求图中AllGather前后的量化，有可能会有精度损失，当前未实现。）

算子功能实现描述

向通信域内所有通信卡发送相同数据量（输入切分成rankSize份）的数据，并从所有通信卡接收相同数据量的数据。如下图所示，对于rankSize=4的情况，各个通信卡将自身的数据均匀分成rankSize份，发送给其他的通信卡。

图2 AllToAll算子功能示意图

计算过程示意（Python）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

goldenTensors = []
for i in range(rankSize):
    golden_out = []
    for j in range(rankSize):
        golden_out_list = intensors[j].reshape(-1).tolist() # intensors[j]代表第j张卡上的intesor，计算时视为一维向量
        split = golden_out_list[i*len(golden_out_list) // rankSize:(i+1)*len(golden_out_list) // rankSize]
        golden_out += split
    golden_out_tensor = torch.tensor(golden_out).reshape(shape) # golden_out_tensor是第i张卡上的输出tensor，输出时shape和输入tensor一致
    goldenTensors.append(golden_out_tensor)
return [goldenTensors[rank].cpu()]

使用场景

用于将数据发送到各个节点上，多对多。AllToAll是对Allgather的扩展，相比于AllGather，AllToAll不同的节点从某一节点收集到的数据是不同的。

图3 AllToAll和AllGather使用对比

应用于模型并行；模型并行里的矩阵转置；数据并行到模型并行的矩阵转置。

定义

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

struct AllToAllParam {
    int rank = 0;
    int rankSize = 0;
    int rankRoot = 0;
    std::string backend = "hccl";
    HcclComm hcclComm = nullptr;
    CommMode commMode = COMM_MULTI_PROCESS;
    std::string rankTableFile;
    std::string commDomain;
    bool transpose = false;
    uint8_t rsv[62] = {0};
};

参数列表

基础场景输入输出

transpose置为false

转置场景输入输出

transpose置为true

规格约束

• 当前只支持 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品 和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品 。
• backend为lccl时不支持多机通信和多通信域，且只支持偶数的rankSize。
• 各个通信卡的输入tensor维度相等。
• 参数维度相乘要能整除rankSize。
• 进行AllToAll计算时会将输入tensor视为一维向量计算，输出的shape则会和输入tensor一致。
• 转置场景下有如下约束：
  • w需要被rankSize整除。
  • w/rankSize的值不能超过90KB。
  • 整个x tensor的大小不能超过190MB。
  • rankSize只能为2、4、8。
• rank、rankSize、rankRoot需满足以下条件。
  • 0 ≤ rank < rankSize
  • 0 ≤ rankRoot < rankSize

• 多用户使用时需要使用ATB_SHARE_MEMORY_NAME_SUFFIX环境变量（请参见Transformer加速库环境变量说明）进行共享内存的区分，以进行初始化信息同步。
• 当使用加速库的通信算子异常退出时，需要清空残留数据，避免影响之后的使用，命令参考如下：

  rm -rf /dev/shm/sem.lccl*
  rm -rf /dev/shm/sem.hccl*
  ipcrm -a