0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
cann-robotcann-robot补充CombineSetup的opapi和tiling UT
89ecea20创建于 4月3日历史提交

aclnnMoeDistributeCombineTeardown

📄 查看源码

产品支持情况

产品 是否支持
Ascend 950PR/Ascend 950DT √
Atlas A3 训练系列产品/Atlas A3 推理系列产品 ×
Atlas A2 训练系列产品/Atlas A2 推理系列产品 ×
Atlas 200I/500 A2 推理产品 ×
Atlas 推理系列产品 ×
Atlas 训练系列产品 ×

功能说明

  • 接口功能：接收aclnnMoeDistributeCombineSetup发来的数据并整合数据（乘权重再相加）。

  • 计算公式：

    ataOut=AllToAllV(expandX)xOut=Sum(expertScales∗ataOut+expertScales∗sharedExpertX)ataOut = AllToAllV(expandX)\\ xOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)

    按MoeDistributeDispatchSetup和MoeDistributeDispatchTeardown算子收集数据的路径原路返还，本算子只做通信状态确认和数据整理，对aclnnMoeDistributeCombineSetup接口发送的数据做通信状态确认，并按照专家的scale权重做加权求和。

  • 注意：该接口必须与aclnnMoeDistributeDispatchSetup、aclnnMoeDistributeDispatchTeardown及aclnnMoeDistributeCombineSetup配套使用。

函数原型

每个算子分为两段式接口，必须先调用 “aclnnMoeDistributeCombineTeardownGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMoeDistributeCombineTeardown”接口执行计算。

aclnnStatus aclnnMoeDistributeCombineTeardownGetWorkspaceSize(
    const aclTensor* expandX,
    const aclTensor* quantExpandX,
    const aclTensor* expertIds,
    const aclTensor* expandIdx,
    const aclTensor* expertScales,
    const aclTensor* commCmdInfo,
    const aclTensor* xActiveMaskOptional,
    const aclTensor* sharedExpertXOptional,
    const char* groupEp,
    int64_t epWorldSize,
    int64_t epRankId,
    int64_t moeExpertNum,
    int64_t expertShardType,
    int64_t sharedExpertNum,
    int64_t sharedExpertRankNum,
    int64_t globalBs,
    int64_t commQuantMode,
    int64_t commType,
    const char* commAlg,
    aclTensor* xOut,
    uint64_t* workspaceSize,
    aclOpExecutor** executor)

aclnnStatus aclnnMoeDistributeCombineTeardown(
    void* workspace,
    uint64_t workspaceSize,
    aclOpExecutor* executor,
    aclrtStream stream)

aclnnMoeDistributeCombineTeardownGetWorkspaceSize

  • 参数说明

    参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度 非连续Tensor
    expandX（aclTensor*） 输入 根据expertIds进行扩展过的token特征 不支持空Tensor。 FLOAT16、BFLOAT16 ND (A, H) √
    quantExpandX（aclTensor*） 输入 对应aclnnMoeDistributeCombineSetup中的quantExpandX输出 不支持空Tensor。 INT8 ND (A, tokenMsgSize) √
    expertIds（aclTensor*） 输入 每个token的topK个专家索引 不支持空Tensor。 INT32 ND (Bs, K) √
    expandIdx（aclTensor*） 输入 对应aclnnMoeDistributeDispatchSetup中的expandIdx输出 不支持空Tensor。 INT32 ND (BS * K, ) √
    expertScales（aclTensor*） 输入 每个token的topK个专家的权重 不支持空Tensor。 FLOAT32 ND (BS, K) √
    commCmdInfo（aclTensor*） 输入 aclnnMoeDistributeCombineSetup的输出，通信的cmd信息 不支持空Tensor。 INT32 ND ((A + epWorldSize) * 16, ) √
    xActiveMaskOptional（aclTensor*） 可选输入 表示token是否参与通信 不支持空Tensor。可选择传入有效数据或填空指针，传空指针时表示所有token都参与通信。 BOOL ND (BS, ) √
    sharedExpertXOptional（aclTensor*） 可选输入 表示共享专家计算后的Token 不支持空Tensor。要求是一个2D或3D的Tensor，当Tensor为2D时，shape为 (BS,H)；当Tensor为3D时，前两位的乘积需等于BS，第三维需等于H。数据类型需跟expandX保持一致。可选择传入有效数据或填空指针，传入有效数据时，sharedExpertNum需为0。 FLOAT16、BFLOAT16 ND (BS, H)或(a, b, H) √
    groupEp 输入 EP通信域名称 字符串长度范围为[1, 128) - - - -
    epWorldSize 输入 EP通信域size 取值支持[2, 384] - - - -
    epRankId 输入 EP域本卡Id 取值范围[0, epWorldSize)。
    同一个EP通信域中各卡的epRankId不重复。    		 - - - -
    moeExpertNum 输入 MoE专家数量 取值范围(0, 512]。
    满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。    		 - - - -
    expertShardType 输入 共享专家卡分布类型 当前仅支持传0，表示共享专家卡排在MoE专家卡前面。 - - - -
    sharedExpertNum 输入 共享专家数量 当前取值范围[0, 4]。0表示无共享专家。当前仅支持无共享专家。 - - - -
    sharedExpertRankNum 输入 共享专家卡数量 当前取值范围[0, epWorldSize / 2]。 - - - -
    globalBs 输入 EP域全局的batch size大小 当每个rank的Bs数一致场景下，globalBs = Bs * epWorldSize 或 globalBs = 0；当每个rank的Bs数不一致场景下，globalBs = maxBs * epWorldSize，其中maxBs表示单卡Bs最大值。 - - - -
    commQuantMode 输入 通信量化类型 取值范围[0, 2]，0表示通信时不进行量化，当前仅支持0。 - - - -
    commType 输入 通信方案选择 取值范围[0, 2]，0表示AICPU-SDMA方案，1表示CCU方案，2表示URMA方案，当前仅支持2。 - - - -
    commAlg 输入 通信算法选择 仅支持传入空指针或空字符串 - - - -
    xOut（aclTensor*） 输出 表示处理后的token 不支持空Tensor。数据类型与expandX保持一致。 FLOAT16、BFLOAT16 ND (BS, H) √
    workspaceSize 输出 返回需要在Device侧申请的workspace大小。 - - - - -
    executor 输出 返回op执行器，包含了算子计算流程。 - - - - -

    • Ascend 950PR/Ascend 950DT：

      • 不支持共享专家场景。
      • epWorldSize当前取值仅支持2、8。
      • moeExpertNum表示MoE专家数量，当前仅能传入32。
      • expertShardType当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
      • sharedExpertNum表示共享专家数量，当前不支持共享专家，仅能传入0。
      • sharedExpertRankNum表示共享专家卡数，当前不支持共享专家，仅能传入0。
      • commQuantMode当前仅支持传入0，表示不进行量化。
      • commType取值范围[0, 2]，当前仅支持2，表示URMA通路。
      • commAlg 当前版本不支持，传空指针即可。

    • Atlas A3 训练系列产品/Atlas A3 推理系列产品：

      • 不支持共享专家场景。
      • epWorldSize当前取值仅支持2、8。
      • moeExpertNum表示MoE专家数量，当前仅能传入32。
      • expertShardType当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
      • sharedExpertNum表示共享专家数量，当前不支持共享专家，仅能传入0。
      • sharedExpertRankNum表示共享专家卡数，当前不支持共享专家，仅能传入0。
      • commQuantMode当前仅支持传入0，表示不进行量化。
      • commType取值范围[0, 2]，当前仅支持2，表示URMA通路。
      • commAlg 当前版本不支持，传空指针即可。

  • 返回值

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

    第一段接口完成入参校验，出现以下场景时报错：

    返回值 错误码 描述
    ACLNN_ERR_PARAM_NULLPTR 161001 输入和输出的必选参数Tensor是空指针。
    ACLNN_ERR_PARAM_INVALID 161002 输入和输出的数据类型不在支持的范围内。
    ACLNN_ERR_INNER_TILING_ERROR 561002 输入和输出的shape不在支持的范围内。
    参数的取值不在支持的范围内。

aclnnMoeDistributeCombineTeardown

  • 参数说明：

    参数名 输入/输出 描述
    workspace 输入 在Device侧申请的workspace内存地址。
    workspaceSize 输入 在Device侧申请的workspace大小，由第一段接口aclnnMoeDistributeCombineTeardownGetWorkspaceSize获取。
    executor 输入 op执行器，包含了算子计算流程。
    指定执行任务的stream流。

  • 返回值：

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

约束说明

  • 确定性计算：

    • aclnnMoeDistributeCombineTeardown默认确定性实现。

  • aclnnMoeDistributeDispatchSetup接口，aclnnMoeDistributeDispatchTeardown接口，aclnnMoeDistributeCombineSetup接口，aclnnMoeDistributeCombineTeardown接口必须配套使用。

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

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

    • A：表示本卡需要分发的最大token数量，取值范围如下：
      • 对于共享专家，要满足A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum；当globalBs非0时，要满足A = globalBs* sharedExpertNum / sharedExpertRankNum。
      • 对于MoE专家，当globalBs为0时，要满足A >= BS * epWorldSize * min(localExpertNum, K)；当globalBs非0时，要满足A >= globalBs * min(localExpertNum, K)。
    • H：表示hidden size隐藏层大小。取值为[1024, 8196]。当前仅支持4096、7168。
    • HS：表示通信时的数据大小，Align512( Align32(H) + Align8(H) / 8 *sizeof(float) )，其中Align512(x) = ((x + 512 - 1) / 512) * 512，Align32(x) = ((x + 32 - 1) / 32) * 32，Align8(x) = ((x + 8 - 1) / 8) * 8。
    • BS：表示batch sequence size，即本卡最终输出的token数量。取值范围为0 < BS ≤ 512。当前仅支持8、16、256。
    • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum。当前仅支持6、8。
    • localExpertNum：表示本卡专家数量。
      • 对于共享专家卡，localExpertNum = 1
      • 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)。moeExpertNum当前仅支持32。
    • tokenMsgSize：表示每个token在数据通信时的维度信息，计算公式为Align512(Align32(H)+Align8(H)/8*sizeof(float))，其中AlignN(x)=((x+N-1)/N*N)。
    • 当前不支持共享专家。sharedExpertNum和sharedExpertRankNum当前仅支持0。

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

    • Ascend 950PR/Ascend 950DT：
      • 要求 >= 2且满足>= 4 * (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
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品：
      • 要求 >= 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

  • 通信域使用约束：

    • 一个模型中的aclnnMoeDistributeDispatchSetup接口，aclnnMoeDistributeDispatchTeardown接口，aclnnMoeDistributeCombineSetup接口，aclnnMoeDistributeCombineTeardown接口仅支持相同EP通信域，且该通信域中不允许有其他算子。

调用示例

  • Ascend 950PR/Ascend 950DT：aclnnMoeDistributeCombineTeardown必须在aclnnMoeDistributeCombineSetup调用之后调用，请参考aclnnMoeDistributeCombineSetup中的调用示例。