aclnnMoeDistributeDispatchTeardown

📄 查看源码

产品支持情况

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

功能说明

  • 接口说明:

    • 接收MOE层EP(Expert Parallelism)域的AllToAllV通信发过来的数据,数据发送端由aclnnMoeDistributeDispatchSetup完成,本接口内完成通信状态确认和数据整理。

    • 注意该接口必须与aclnnMoeDistributeDispatchSetup,aclnnMoeDistributeCombineSetup,aclnnMoeDistributeCombineTeardown配套使用。

函数原型

每个算子分为两段式接口,必须先调用 “aclnnMoeDistributeDispatchTeardownGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnMoeDistributeDispatchTeardown”接口执行计算。

aclnnStatus aclnnMoeDistributeDispatchTeardownGetWorkspaceSize(
    const aclTensor* x, 
    const aclTensor* y, 
    const aclTensor* expertIds, 
    const aclTensor* commCmdInfo,
    const char* groupEp, 
    int64_t epWorldSize, 
    int64_t epRankId,
    int64_t moeExpertNum, 
    int64_t expertShardType, 
    int64_t sharedExpertNum, 
    int64_t sharedExpertRankNum,
    int64_t quantMode, 
    int64_t globalBs, 
    int64_t expertTokenNumsType, 
    int64_t commType, 
    char* commAlg,
    aclTensor* expandXOut, 
    aclTensor* dynamicScalesOut, 
    aclTensor* assistInfoForCombineOut,
    aclTensor* expertTokenNumsOut, 
    uint64_t* workspaceSize, 
    aclOpExecutor** executor)
aclnnStatus aclnnMoeDistributeDispatchTeardown(
    void            *workspace,
    uint64_t        workspaceSize,
    aclOpExecutor   *executor,
    aclrtStream     stream)

aclnnMoeDistributeDispatchTeardownGetWorkspaceSize

  • 参数说明

    参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度(shape) 非连续Tensor
    x 输入 表示本卡发送的token数据。 要求为2D Tensor。 FLOAT16、BFLOAT16 ND (Bs, H)
    y 输入 aclnnMoeDistributeDispatchSetup的输出,表示本卡待发送的通信数据,通信数据对输入token数据做了算法重排;如需量化,先将输入token做量化处理,再对数据做重排。 要求为2D Tensor。 FLOAT16、BFLOAT16、INT8、HIFP8、FP8E5M2、FP8E4M3 ND (BS * (K + sharedExpertNum), tokenMsgSize)
    expertIds 输入 每个token的topK个专家索引。 要求为2D Tensor。 INT32 ND (Bs, K)
    commCmdInfo 输入 aclnnMoeDistributeDispatchSetup的输出,通信的cmd信息。 要求为1D Tensor。 INT32 ND (BS * (K + sharedExpertNum) + epWorldSize * localExpertNum) * 16
    groupEp 输入 EP通信域名称(专家并行通信域)。 字符串长度范围为[1, 128)。 STRING - - -
    epWorldSize 输入 EP通信域大小。 取值范围[2, 384]。 INT64 - - -
    epRankId 输入 EP域本卡Id。 取值范围[0, epWorldSize),同一个EP通信域中各卡的epRankId不重复。 INT64 - - -
    moeExpertNum 输入 MoE专家数量。 取值范围(0, 512]。满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。 INT64 - - -
    expertShardType 输入 表示共享专家卡分布类型。 当前仅支持传入0,表示共享专家卡排在MoE专家卡前面。 INT64 - - -
    sharedExpertNum 输入 表示共享专家数量。 当前版本取值范围[0, 4]。 INT64 - - -
    sharedExpertRankNum 输入 表示共享专家卡数量。 取值范围[0, epWorldSize / 2] INT64 - - -
    quantMode 输入 表示量化模式。 取值范围[0, 4]。0表示非量化,1表示静态量化,2表示Pertoken动态量化,3表示Pergroup动态量化,4表示MX量化,当前仅支持0和4。 INT64 - - -
    globalBs 输入 EP域全局的batch size大小。
    • 各rank Bs一致时,globalBs = Bs * epWorldSize 或 0。
    • 各rank Bs不一致时,globalBs = maxBs * epWorldSize(maxBs为单卡Bs最大值)。
    INT64 - - -
    expertTokenNumsType 输入 表示输出expertTokenNums中值的语义类型。
    • expertTokenNums为0表示expertTokenNums中的输出为每个专家处理的token数的前缀和,1表示expertTokenNums中的输出为每个专家处理的token数量。
    INT64 - - -
    commType 输入 表示通信方案选择。 取值范围[0, 2],0表示AICPU-SDMA方案,1表示CCU方案,2表示URMA方案,当前版本仅支持2。 INT64 - - -
    commAlg 输入 表示通信亲和内存布局算法。 预留字段,当前版本不支持,传空指针或空字符串即可。 STRING - - -
    expandXOut 输出 根据expertIds进行扩展过的token特征。 要求为2D Tensor。 FLOAT16、BFLOAT16、INT8、FP8E5M2、FP8E4M3、HIFP8 ND (A, H)
    dynamicScalesOut 输出 每个token的动态量化参数。 要求为1D Tensor。当quantMode为2、3、4时,才会输出实际有效值。 FLOAT32、FP8E8M0 ND (A, )
    assistInfoForCombineOut 输出 表示给同一专家发送的token个数,对应Combine系列算子中的assistInfoForCombine。 要求为1D Tensor。 INT32 ND (A * 128, )
    expertTokenNumsOut 输出 表示每个专家收到的token个数。 要求为1D Tensor。 INT64 ND (localExpertNum, )
    workspaceSize 输出 返回需要在Device侧申请的workspace大小。 - - - - -
    executor 输出 返回op执行器,包含了算子的计算流程。 - - - - -
    • Ascend 950PR/Ascend 950DT:

      • groupEp 字符串长度范围为[1, 128)。
      • epWorldSize 取值范围[2, 384]。当前仅支持2、8。
      • epRankId 取值范围[0, epWorldSize)。同一个EP通信域中各卡的epRankId不能重复。
      • moeExpertNum 取值范围(0, 512]。
      • expertShardType 当前仅支持传0,表示共享专家卡排在MoE专家卡前面。
      • sharedExpertNum 当前取值范围[0, 4]。
      • sharedExpertRankNum 取值范围[0, epWorldSize / 2]。
      • globalBs 当每个rank的Bs数一致场景下,globalBs = Bs * epWorldSize 或 globalBs = 0;当每个rank的Bs数不一致场景下,globalBs = maxBs * epWorldSize,其中maxBs表示单卡Bs最大值。
      • expertTokenNumsType当前仅支持1。
      • commType 当前仅支持2。
      • commAlg 预留字段,当前版本不支持,传NULL即可。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品:

      • groupEp 字符串长度范围为[1, 128)。
      • epWorldSize 取值范围[2, 384]。
      • epRankId 取值范围[0, epWorldSize)。同一个EP通信域中各卡的epRankId不能重复。
      • moeExpertNum 取值范围(0, 512]。
      • expertShardType 当前仅支持传0,表示共享专家卡排在MoE专家卡前面。
      • sharedExpertNum 当前取值范围[0, 4]。
      • sharedExpertRankNum 取值范围[0, epWorldSize / 2]。
      • globalBs 当每个rank的Bs数一致场景下,globalBs = Bs * epWorldSize 或 globalBs = 0;当每个rank的Bs数不一致场景下,globalBs = maxBs * epWorldSize,其中maxBs表示单卡Bs最大值。
      • commType 当前仅支持0。
      • commAlg 当前版本不支持,传空指针即可。
  • 返回值:

    返回aclnnStatus状态码,具体参见aclnn返回码

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

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

aclnnMoeDistributeDispatchTeardown

  • 参数说明

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

    返回aclnnStatus状态码,具体参见aclnn返回码

约束说明

  1. 确定性计算:

    • aclnnMoeDistributeDispatchTeardown默认确定性实现。
  2. aclnnMoeDistributeDispatchTeardown接口与aclnnMoeDistributeDispatchSetup,aclnnMoeDistributeCombineSetup,aclnnMoeDistributeCombineTeardown接口必须配套使用。

  3. 调用接口过程中使用的groupEpepWorldSizemoeExpertNumexpertShardTypesharedExpertNumsharedExpertRankNumglobalBscommQuantModecommTypecommAlg参数取值所有卡需保持一致,groupEpepWorldSizeexpertShardTypesharedExpertNumsharedExpertRankNumglobalBscommQuantModecommTypecommAlg参数取值在网络中不同层中也需保持一致,且和aclnnMoeDistributeDispatchTeardownaclnnMoeDistributeCombineSetupaclnnMoeDistributeCombineTeardown对应参数也保持一致。

  4. Atlas A3 训练系列产品/Atlas A3 推理系列产品:该场景下单卡包含双DIE(简称为“晶粒”或“裸片”),因此参数说明里的“本卡”均表示单DIE。

  5. 参数说明里shape格式说明:

    • A:表示本卡可能接收的最大token数量,取值范围如下:

      • 对于MoE专家,当globalBs为0时,要满足A >= BS * epWorldSize * min(localExpertNum, K);当globalBs非0时,要满足A >= globalBs * min(localExpertNum, K)。
      • 对于共享专家,当globalBs为0时,要满足A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum;当globalBs非0时,要满足A = globalBs * sharedExpertNum / sharedExpertRankNum
    • H:表示hidden size隐藏层大小,取值范围[1024, 8192]。当前仅支持4096、7168。

    • BS:表示batch sequence size,即本卡最终输出的token数量,取值范围为0 < BS ≤ 512。当前仅支持8、16、256。

    • K:表示选取topK个专家,取值范围为0 < K ≤ 16同时满足0 < KmoeExpertNum。当前仅支持6、8。

    • localExpertNum:表示本卡专家数量。

      • 对于共享专家卡,localExpertNum = 1
      • 对于MoE专家卡,localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)。moeExpertNum当前仅支持32。
    • tokenMsgSize:表示每个token在数据通信时的维度信息。

      • 非量化场景下,tokenMsgSize = Align256(H)。
      • 量化场景下,tokenMsgSize = Align512(Align32(H) + 4 ),其中AlignN(x) = ((x + N - 1) / N) * N。
    • 当前版本暂不支持共享专家。sharedExpertNum和sharedExpertRankNum当前仅支持0。

  6. HCCL_BUFFSIZE:

    • Ascend 950PR/Ascend 950DT: 调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理,该环境变量表示单个通信域占用内存大小,单位MB,不配置时默认为200MB。要求 >= 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 推理系列产品: 调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理,该环境变量表示单个通信域占用内存大小,单位MB,不配置时默认为200MB。要求 >= 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。

  7. 通信域使用约束:

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

调用示例

  • Ascend 950PR/Ascend 950DT:aclnnMoeDistributeDispatchTeardown必须在aclnnMoeDistributeDispatchSetup调用之后调用,请参考aclnnMoeDistributeDispatchSetup中的调用示例。