aclnnMoeDistributeCombineV2

支持的产品型号

• [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]

功能说明

算子功能：当存在TP域通信时，先进行ReduceScatterV通信，再进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加）；当不存在TP域通信时，进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加）。

$$rsOut = ReduceScatterV(expandX)\\ ataOut = AllToAllV(rsOut)\\ xOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)$$

注意该接口必须与aclnnMoeDistributeDispatchV2配套使用，相当于按MoeDistributeDispatchV2算子收集数据的路径原路返回

相较于aclnnMoeDistributeCombine接口，该接口变更如下：

• 输入了更详细的token信息辅助aclnnMoeDistributeCombineV2高效地进行全卡同步，因此原接口中shape为(Bs * K,)的expandIdx入参替换为shape为(A * 128,)的assistInfoForCombine参数；
• 新增sharedExpertXOptional入参，支持在sharedExpertNum为0时，由用户输入共享专家计算后的token；
• 新增commAlg入参，预留字段，暂未使用。

详细说明请参考以下参数说明。

函数原型

每个算子分为undefined，必须先调用 “aclnnMoeDistributeCombineV2GetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMoeDistributeCombineV2”接口执行计算。

• aclnnStatus aclnnMoeDistributeCombineV2GetWorkspaceSize(const aclTensor* expandX, const aclTensor* expertIds, const aclTensor* assistInfoForCombine, const aclTensor* epSendCounts, const aclTensor* expertScales, const aclTensor* tpSendCountsOptional, const aclTensor* xActiveMaskOptional, const aclTensor* activationScaleOptional, const aclTensor* weightScaleOptional, const aclTensor* groupListOptional, const aclTensor* expandScalesOptional, const aclTensor* sharedExpertXOptional, const char* groupEp, int64_t epWorldSize, int64_t epRankId, int64_t moeExpertNum, const char* groupTp, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, int64_t outDtype, int64_t commQuantMode, int64_t groupListType, const char* commAlg, aclTensor* xOut, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnMoeDistributeCombineV2(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnMoeDistributeCombineV2GetWorkspaceSize

• 参数说明：

  • expandX（aclTensor*，计算输入）：根据expertIds进行扩展过的token特征，Device侧的aclTensor，要求为一个2D的Tensor，shape为 (max(tpWorldSize, 1) * A , H)，数据类型支持FLOAT16、BFLOAT16，undefined要求为ND，支持undefined。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：不支持共享专家场景。
  • expertIds（aclTensor*，计算输入）：每个token的topK个专家索引，Device侧的aclTensor，要求为一个2D的Tensor，shape为 (Bs, K)。数据类型支持INT32，undefined要求为ND，支持undefined。
  • assistInfoForCombine（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV2中的assistInfoForCombineOut输出，Device侧的aclTensor，要求是一个1D的Tensor，要求shape为 (A * 128, )，数据类型支持INT32，undefined要求为ND，支持undefined。
  • epSendCounts（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV2中的epRecvCounts输出，Device侧的aclTensor，要求是一个1D的Tensor。数据类型支持INT32，undefined要求为ND，支持undefined。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：要求shape为 (moeExpertNum + 2 * globalBs * K * serverNum, )，前moeExpertNum个数表示从EP通信域各卡接收的token数，2 * globalBs * K * serverNum存储了机间机内做通信前combine可以提前做reduce的token个数和token在通信区中的偏移，globalBs传入0时在此处应当按照Bs * epWorldSize计算。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求shape为 (epWorldSize * max(tpWorldSize, 1) * localExpertNum, )。
  • expertScales（aclTensor*，计算输入）：每个token的topK个专家的权重，Device侧的aclTensor，要求是一个2D的Tensor，shape为 (Bs, K)。类型支持FLOAT32，undefined要求为ND，支持undefined。
  • tpSendCountsOptional（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV2中的tpRecvCounts输出，Device侧的aclTensor，若有TP域通信需要传参，若无TP域通信，传空指针即可。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前不支持TP域通信。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当有TP域通信时，要求是一个1D的Tensor，shape为 (tpWorldSize, )。数据类型支持INT32，undefined要求为ND，支持undefined。
  • xActiveMaskOptional（aclTensor*，计算输入）：表示token是否参与通信，Device侧的aclTensor。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：预留参数，当前版本不支持，传空指针即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求是一个1D的Tensor，shape为 (Bs, )，数据类型支持BOOL；可选择传入有效数据或传入空指针。传入有效数据时，参数为true表示对应的token参与通信，true必须排到false之前，例：{true, false, true} 为非法输入；传入空指针时是表示所有token都会参与通信。undefined要求为ND，支持undefined。
  • activationScaleOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • weightScaleOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • groupListOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • expandScalesOptional（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV2中的expandScales输出，Device侧的aclTensor。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：要求是一个1D的Tensor，shape为 (A, )，数据类型支持FLOAT32，undefined要求为ND，支持undefined。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：预留参数，当前版本不支持，传空指针即可。
  • sharedExpertXOptional（aclTensor*, 计算输入）：表示共享专家计算后的token，Device侧的aclTensor。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：预留参数，当前版本不支持，传空指针即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求是一个2D或3D的Tensor，当Tensor为2D时，shape为 (Bs, H)；当Tensor为3D时，前两位的乘积需等于Bs，第三维需等于H。数据类型需跟expandX保持一致。可传/可不传，传入时，sharedExpertRankNum需为0。undefined要求为ND，支持undefined。
  • groupEp（char*，计算输入）：EP通信域名称，专家并行的通信域，string数据类型。字符串长度范围为[1, 128)，不能和groupTp相同。
  • epWorldSize（int64_t，计算输入）：EP通信域size，数据类型支持INT64。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：取值支持16、32、64。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值支持[2, 384]。
  • epRankId（int64_t，计算输入）：EP域本卡Id，数据类型支持INT64，取值范围[0, epWorldSize)。同一个EP通信域中各卡的epRankId不重复。
  • moeExpertNum（int64_t，计算输入）: MoE专家数量，数据类型支持INT64，取值范围(0, 512]，并且满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：还需满足moeExpertNum / (epWorldSize - sharedExpertRankNum) <= 24。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：还需满足moeExpertNum / (epWorldSize - sharedExpertRankNum) * epWorldSize <= 1280。
  • groupTp（char*，计算输入）：TP通信域名称，数据并行的通信域，string数据类型。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传空字符即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：字符串长度范围为[1, 128)，不能和groupEp相同。
  • tpWorldSize（int64_t，计算输入）：TP通信域size，int数据类型。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传0即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围[0, 2]，0和1表示无TP域通信，有TP域通信时仅支持2。
  • tpRankId（int64_t，计算输入）：TP域本卡Id，数据类型支持INT64。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传0即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围[0, 1]，同一个TP通信域中各卡的tpRankId不重复。无TP域通信时，传0即可。
  • expertShardType（int64_t，计算输入）：表示共享专家卡分布类型，数据类型支持INT64。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传0即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
  • sharedExpertNum（int64_t，计算输入）：表示共享专家数量，一个共享专家可以复制部署到多个卡上，数据类型支持INT64。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传0即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前取值范围[0, 4]。0表示无共享专家。
  • sharedExpertRankNum（int64_t，计算输入）：表示共享专家卡数量，数据类型支持INT64。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当前版本不支持，传0即可。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前取值范围[0, epWorldSize - 1)，不为0时需满足epWorldSize % (sharedExpertRankNum / sharedExpertNum) = 0。当sharedExpertRankNum为0时，算子会忽略sharedExpertNum传入数值。
  • globalBs（int64_t，计算输入）：EP域全局的batch size大小，数据类型支持INT64。当每个rank的Bs数一致时，globalBs = Bs * epWorldSize 或 globalBs = 0；当每个rank的Bs数不一致时，globalBs = maxBs * epWorldSize，其中maxBs表示单卡Bs最大值。
  • outDtype（int64_t，计算输入）：用于指定输出x的数据类型，预留参数，当前版本不支持，传0即可。
  • commQuantMode（int64_t，计算输入）：通信量化类型。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：取值范围0或者2，0表示通信时不进行量化，2表示通信时进行int8量化，2仅当HCCL_INTRA_PCIE_ENABLE为1且HCCL_INTRA_ROCE_ENABLE为0且驱动版本不低于25.0.RC1.1时支持。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围0或者2，0表示通信时不进行量化，2表示通信时进行int8量化，int8量化当且仅当tpWorldSize < 2时可使能。
  • groupListType（int64_t，计算输入）：group List格式，预留参数，当前版本不支持，传0即可。
  • commAlg（char*，计算输入）：表示通信亲和内存布局算法，string数据类型，预留字段，当前版本不支持，传入空指针即可。
  • xOut（aclTensor*，计算输出）：表示处理后的token，Device侧的aclTensor，要求是一个2D的Tensor，shape为(Bs, H)，数据类型、数据格式与expandX保持一致。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值

  返回aclnnStatus状态码，具体参见undefined。

  [object Object]

aclnnMoeDistributeCombineV2

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnMoeDistributeCombineV2GetWorkspaceSize获取。
  • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
  • stream（aclrtStream，入参）：指定执行任务的AscendCL stream流。

• 返回值：

  返回aclnnStatus状态码，具体参见undefined。

约束说明

• aclnnMoeDistributeDispatchV2接口与aclnnMoeDistributeCombineV2接口必须配套使用，具体参考undefined。

• 在不同产品型号、不同通信算法或不同版本中，aclnnMoeDistributeDispatchV2的Tensor输出assistInfoForCombineOut、epRecvCounts、tpRecvCounts、expandScales中的元素值可能不同，使用时直接将上述Tensor传给aclnnMoeDistributeCombineV2对应参数即可，模型其他业务逻辑不应对其存在依赖。

• 调用接口过程中使用的groupEp、epWorldSize、moeExpertNum、groupTp、tpWorldSize、expertShardType、sharedExpertNum、sharedExpertRankNum、globalBs参数取值所有卡需保持一致，网络中不同层中也需保持一致，且和aclnnMoeDistributeDispatchV2对应参数也保持一致。

• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。

• 参数说明里shape格式说明：

  • A：表示本卡需要分发的最大token数量，取值范围如下：
    • 对于共享专家，要满足A = Bs * epWorldSize * sharedExpertNum / sharedExpertRankNum。
    • 对于MoE专家，当globalBs为0时，要满足A >= Bs * epWorldSize * min(localExpertNum, K)；当globalBs非0时，要满足A >= globalBs * min(localExpertNum, K)。
  • H：表示hidden size隐藏层大小。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：取值范围(0, 7168]，且保证是32的整数倍。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值为[1024, 7168]。
  • Bs：表示batch sequence size，即本卡最终输出的token数量。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：取值范围为0 < Bs ≤ 256。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围为0 < Bs ≤ 512。
  • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum。
  • serverNum：表示服务器的节点数，取值只支持2、4、8。
  • localExpertNum：表示本卡专家数量。
    • 对于共享专家卡，localExpertNum = 1
    • 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)，localExpertNum > 1时，不支持TP域通信。

• HCCL_BUFFSIZE： 调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。

  • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：要求 >= 2 * (Bs * epWorldSize * min(localExpertNum, K) * H * sizeof(uint16) + 2MB)。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求 >= 2且满足>= 2 * (localExpertNum * maxBs * epWorldSize * Align512(Align32(2 * H) + 44) + (K + sharedExpertNum) * maxBs * Align512(2 * H))，localExpertNum需使用MoE专家卡的本卡专家数，其中Align512(x) = ((x + 512 - 1) / 512) * 512，Align32(x) = ((x + 32 - 1) / 32) * 32。

• HCCL_INTRA_PCIE_ENABLE和HCCL_INTRA_ROCE_ENABLE：

  • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：当K = 8 且 Bs ≤ 128时，设置环境变量HCCL_INTRA_PCIE_ENABLE = 1和HCCL_INTRA_ROCE_ENABLE = 0可以减少跨机通信数据量，可能提升算子性能。 此时，HCCL_BUFFSIZE要求 >= moeExpertNum * Bs * (H * sizeof(dtypeX) + 4 * ((K + 7) / 8 * 8) * sizeof(uint32)) + 4MB + 100MB。并且，对于入参moeExpertNum，只要求moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0，不要求moeExpertNum / (epWorldSize - sharedExpertRankNum) <= 24。

• 本文公式中的"/"表示整除。

• 通信域使用约束：

  • 一个模型中的aclnnMoeDistributeCombineV2和aclnnMoeDistributeDispatchV2仅支持相同EP通信域，且该通信域中不允许有其他算子。
  • 一个模型中的aclnnMoeDistributeCombineV2和aclnnMoeDistributeDispatchV2仅支持相同TP通信域或都不支持TP通信域，有TP通信域时该通信域中不允许有其他算子。

调用示例

以[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]为例，调起MoeDistributeCombineV2和MoeDistributeDispatchV2算子。

• 文件准备：
  1.新建combineDemo目录，按照下方指导在combineDemo下新建aclnnCombineDemo.cpp，buildCombine.sh，文件并参考如下代码修改。

  2.安装cann包，并根据下方指导编译运行combineDemo。

• 编译脚本

  [object Object]

• 编译与运行：

  [object Object]

• 示例代码如下，仅供参考

  [object Object]