aclnnMoeDistributeCombineAddRmsNormV2

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	×
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

算子功能：当存在TP域通信时，先进行ReduceScatterV通信，再进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加）；当不存在TP域通信时，进行AlltoAllV通信，最后将接收的数据整合（乘权重再相加），之后完成Add + RmsNorm融合。算子支持动态缩容场景，支持在创建通信域后，出现故障卡，将故障卡从通信域剔除，算子可正常执行，无需重新编译；支持特殊专家场景。

[object Object][object Object]

计算公式：

$$rsOut = ReduceScatterV(expandX)\\ ataOut = AlltoAllV(rsOut)\\ combineOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)\\ x = combineOut + residualX\\ y = \frac{x}{RMS(x)} * gamma,\quad\text{where}RMS(x) = \sqrt{\frac{1}{H}\sum_{i=1}^{H}x_{i}^{2}+normEps}\\$$

注意该接口必须与aclnnMoeDistributeDispatchV3配套使用，相当于按MoeDistributeDispatchV3接口收集数据的路径原路返还。

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

• 新增支持动态缩容场景，支持在创建通信域后，出现故障卡，将故障卡从通信域剔除，算子可正常执行，无需重新编译，通过传入elasticInfoOptional参数使能本特性。
• 新增支持特殊专家场景
  • zeroExpertNum ≠ 0：通过传入大于0的zeroExpertNum参数使能本特性。$$Moe(oriXOptional) = 0$$
  • copyExpertNum ≠ 0：通过传入大于0的copyExpertNum参数使能本特性，同时还需传入有效的oriXOptional参数。$$Moe(oriXOptional) = oriXOptional$$
  • constExpertNum ≠ 0：通过传入大于0的constExpertNum参数使能本特性，同时还需传入有效的oriXOptional、constExpertAlpha1Optional、constExpertAlpha2Optional、constExpertVOptional参数。$$Moe(oriXOptional) = constExpertAlpha1Optional * oriXOptional + constExpertAlpha2Optional * constExpertVOptional$$

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

函数原型

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

• aclnnStatus aclnnMoeDistributeCombineAddRmsNormV2GetWorkspaceSize(const aclTensor* expandX, const aclTensor* expertIds, const aclTensor* assistInfoForCombine, const aclTensor* epSendCounts, const aclTensor* expertScales, const aclTensor* residualX, const aclTensor* gamma, const aclTensor* tpSendCountsOptional, const aclTensor* xActiveMaskOptional, const aclTensor* activationScaleOptional, const aclTensor* weightScaleOptional, const aclTensor* groupListOptional, const aclTensor* expandScalesOptional, const aclTensor* sharedExpertXOptional, const aclTensor* elasticInfoOptional, const aclTensor* oriXOptional, const aclTensor* constExpertAlpha1Optional, const aclTensor* constExpertAlpha2Optional, const aclTensor* constExpertVOptional,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, float normEps, int64_t zeroExpertNum, int64_t copyExpertNum, int64_t constExpertNum, aclTensor* yOut, aclTensor* rstdOut, aclTensor* xOut, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnMoeDistributeCombineAddRmsNormV2(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnMoeDistributeCombineAddRmsNormV2GetWorkspaceSize

• 参数说明：

  • expandX（aclTensor*，计算输入）：根据expertIds进行扩展过的token特征，Device侧的aclTensor，要求为一个2D的Tensor，shape为 (max(tpWorldSize, 1) * A , H)，数据类型支持BFLOAT16，要求为ND，支持。
  • expertIds（aclTensor*，计算输入）：每个token的topK个专家索引，Device侧的aclTensor，要求为一个2D的Tensor，shape为 (BS, K)。数据类型支持INT32，要求为ND，支持。
  • assistInfoForCombine（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV3中的assistInfoForCombineOut输出，Device侧的aclTensor，要求是一个1D的Tensor，数据类型支持INT32，要求为ND，支持。shape为 (A * 128, )。
  • epSendCounts（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV3中的epRecvCounts输出，Device侧的aclTensor，要求是一个1D的Tensor。数据类型支持INT32，要求为ND，支持。shape为 (epWorldSize * max(tpWorldSize, 1) * localExpertNum, )。
  • expertScales（aclTensor*，计算输入）：每个token的topK个专家的权重，Device侧的aclTensor，要求是一个2D的Tensor，shape为 (BS, K)。类型支持FLOAT32，要求为ND，支持。
  • residualX（aclTensor*，计算输入）：AddRmsNorm中Add的右矩阵，Device侧的aclTensor，要求是一个3D的Tensor，shape为(BS，1，H)。类型为BFLOAT16，要求为ND，支持非连续的Tensor。
  • gamma（aclTensor*，计算输入）：RmsNorm中的gamma输入，Device侧的aclTensor，要求是一个1D的Tensor，shape为(H, )。类型为BFLOAT16，要求为ND，支持非连续的Tensor。
  • tpSendCountsOptional（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV3中的tpRecvCounts输出，Device侧的aclTensor，若有TP域通信需要传参，若无TP域通信，传空指针即可。当有TP域通信时，要求是一个1D的Tensor，shape为 (tpWorldSize, )。数据类型支持INT32，要求为ND，支持。
  • xActiveMaskOptional（aclTensor*，计算输入）：表示token是否参与通信，Device侧的aclTensor，要求是一个1D或者2D Tensor。当输入为1D时，shape为(BS, ); 当输入为2D时，shape为(BS, K)。数据类型支持BOOL；可选择传入有效数据或传入空指针。当输入为1D时，参数为true表示对应的token参与通信，true必须排到false之前，例：{true, false, true} 为非法输入；当输入为2D时，参数为true表示当前token对应的expert_ids参与通信。 若当前token对应的K个BOOL值全为false，表示当前token不会参与通信。默认所有token都会参与通信。当每张卡的BS数量不一致时，所有token必须全部有效。要求为ND，支持。
  • activationScaleOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • weightScaleOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • groupListOptional（aclTensor*，计算输入）：Device侧的aclTensor，预留参数，当前版本不支持，传空指针即可。
  • expandScalesOptional（aclTensor*，计算输入）：对应aclnnMoeDistributeDispatchV3中的expandScales输出，Device侧的aclTensor。预留参数，当前版本不支持，传空指针即可。
  • sharedExpertXOptional（aclTensor*, 计算输入）：表示共享专家计算后的token，Device侧的aclTensor。要求是一个2D或3D的Tensor，当Tensor为2D时，shape为 (BS, H)；当Tensor为3D时，shape为 (BS, 1, H)。数据类型需跟expandX保持一致。可传/可不传。要求为ND，支持。
  • elasticInfoOptional（aclTensor*, 计算输入）：Device侧的aclTensor，表示EP通信域的动态缩容信息。当某些通信卡因异常从通信域内被剔除，实际参与通信的卡数与创建通信域时不一致，可从本参数中获取当前部署信息。可选择传入有效数据或填空指针，传入空指针时表示不使能动态缩容功能；当传入有效数据时，要求是一个1D的Tensor，shape为(4 + 2 * epWorldSize,)，数据类型支持INT32；要求ND，支持。Tensor中的前四个数字分别表示（是否缩容，缩容后实际rank数，缩容后共享专家使用rank数，缩容后moe专家个数），后2*epWorldSize表示2个rank映射表，第一个Table的映射关系为Table[epRankId] = localEpRankId或-1，-1表示epRankId这张卡从通信域中剔除，localEpRankId表示缩容后本卡在新EP通信域中的rankIndex，第二个Table映射关系为Table2[localEpRankId] = epRankId
  • oriXOptional（aclTensor*, 计算输入）：Device侧的aclTensor，表示未经过FFN（Feed-Forward Neural network）的token数据，对应上述中的oriXOptional，在使能copyExpert或使能constExpert的场景下需要本输入数据。可选择传入有效数据或填空指针，当copyExpertNum不为0或constExpertNum不为0时必须传入有效输入；当传入有效数据时，要求是一个2D的Tensor，shape为（BS, H），数据类型需跟expandX保持一致；要求为ND，支持。
  • constExpertAlpha1Optional（aclTensor*，计算输入）：Device侧的aclTensor，在使能constExpert的场景下需要输入的计算系数，对应上述中的constExpertAlpha1Optional。可选择传入有效数据或填空指针，当constExpertNum不为0时必须传入有效输入；当传入有效数据时，要求是一个1D的Tensor，shape为(constExpertNum,)，数据类型需跟expandX保持一致；要求为ND，支持。
  • constExpertAlpha2Optional（aclTensor*，计算输入）：Device侧的aclTensor，在使能constExpert的场景下需要输入的计算系数，对应上述中的constExpertAlpha2Optional。可选择传入有效数据或填空指针，当constExpertNum不为0时必须传入有效输入；当传入有效数据时，要求是一个1D的Tensor，shape为(constExpertNum,)，数据类型需跟expandX保持一致；要求为ND，支持。
  • constExpertVOptional（aclTensor*，计算输入）：Device侧的aclTensor，在使能constExpert的场景下需要输入的计算系数，对应上述中的constExpertVOptional。可选择传入有效数据或填空指针，当constExpertNum不为0时必须传入有效输入；当传入有效数据时，要求是一个2D的Tensor，shape为(constExpertNum, H)，数据类型需跟expandX保持一致；要求为ND，支持。
  • groupEp（char*，计算输入）：EP通信域名称，专家并行的通信域，string数据类型。字符串长度范围为[1, 128)，不能和groupTp相同。
  • epWorldSize（int64_t，计算输入）：EP通信域size，数据类型支持INT64。取值支持(1, 768]。
  • epRankId（int64_t，计算输入）：EP域本卡Id，数据类型支持INT64，取值范围[0, epWorldSize)。同一个EP通信域中各卡的epRankId不重复。
  • moeExpertNum（int64_t，计算输入）: MoE专家数量，数据类型支持INT64，取值范围(0, 1024]，并且满足moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0。
  • groupTp（char*，计算输入）：TP通信域名称，数据并行的通信域，string数据类型。字符串长度范围为[1, 128)，不能和groupEp相同。
  • tpWorldSize（int64_t，计算输入）：TP通信域size，int数据类型。取值范围[0, 2]，0和1表示无tp域通信，有tp域通信时仅支持2。
  • tpRankId（int64_t，计算输入）：TP域本卡Id，数据类型支持INT64。取值范围[0, 1]，同一个TP通信域中各卡的tpRankId不重复。无TP域通信时，传0即可。
  • expertShardType（int64_t，计算输入）：表示共享专家卡分布类型，数据类型支持INT64。当前仅支持传0，表示共享专家卡排在MoE专家卡前面。
  • sharedExpertNum (int64_t，计算输入）：表示共享专家数量，当前版本不支持，传0即可。
  • sharedExpertRankNum（int64_t，计算输入）：表示共享专家卡数量，当前版本不支持，仅支持传入0。
  • 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，计算输入）：通信量化类型。当前版本不支持，传0即可。
  • groupListType（int64_t，计算输入）：group List格式，预留参数，当前版本不支持，传0即可。
  • commAlg（char*，计算输入）：表示通信亲和内存布局算法，string数据类型，预留字段，当前版本不支持，传入空指针即可。
  • normEps（float，计算输入）：用于防止AddRmsNorm除0错误，数据类型为float，可以为1e-6。
  • zeroExpertNum （int64_t，计算输入）：表示零专家的数量。取值范围:[0, MAX_INT32)，其中MAX_INT32值为2147483647，合法的零专家ID的值是[moeExpertNum, moeExpertNum + zeroExpertNum)。
  • copyExpertNum （int64_t，计算输入）：表示拷贝专家的数量。取值范围:[0, MAX_INT32)，其中MAX_INT32值为2147483647，合法的拷贝专家ID的值是[moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum)。
  • constExpertNum （int64_t，计算输入）：表示常量专家的数量。取值范围:[0, MAX_INT32)，其中MAX_INT32值为2147483647，合法的常量专家ID的值是[moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum)。
  • yOut（aclTensor*，计算输出）：RmsNorm后的输出结果，Device侧的aclTensor，要求是一个3D的Tensor，shape为（BS，1，H），数据类型、数据格式与residualX保持一致。
  • rstdOut（aclTensor*，计算输出）：RmsNorm后的输出结果，Device侧的aclTensor，要求是一个3D的Tensor，shape为（Bs，1，1），数据类型支持FLOAT32，数据格式支持ND。
  • xOut（aclTensor*，计算输出）：Add后的输出结果，Device侧的aclTensor，要求是一个3D的Tensor，shape为(BS, 1，H)，数据类型、数据格式与residualX保持一致。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值

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

  [object Object]

aclnnMoeDistributeCombineAddRmsNormV2

• 参数说明：

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

• 返回值：

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

约束说明

• aclnnMoeDistributeDispatchV3接口与aclnnMoeDistributeCombineAddRmsNormV2接口必须配套使用，具体参考。

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

• 动态缩容后的部署信息通过elasticInfoOptional参数传递给算子，无需修改其他参数。动态缩容后，MOE专家卡上的本卡部署MOE专家数需与缩容前保持一致。

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

• moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum < MAX_INT32，其中MAX_INT32值为2147483647。

• 动态缩容功能不支持在TP并行场景下使能。

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

  • A：表示本卡需要分发的最大token数量，取值范围如下：

    • 当globalBs为0时，要满足A >= BS * epWorldSize * min(localExpertNum, K)；
    • 当globalBs非0时，要满足A >= globalBs * min(localExpertNum, K)。

  • H：表示hidden size隐藏层大小，取值范围为[1024, 8192]。

  • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为0 < BS ≤ 512。

  • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum。

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

    • 对于共享专家卡，localExpertNum = 1
    • 对于MoE专家卡，localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)，localExpertNum > 1时，不支持TP域通信。

• HCCL_BUFFSIZE： 调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理，该环境变量表示单个通信域占用内存大小，单位MB，不配置时默认为200MB。要求 >= 2且满足1024 ^ 2 * (HCCL_BUFFSIZE - 2) / 2 >= (BS * 2 * (H + 128) * (epWorldSize * localExpertNum + K + 1)，localExpertNum表示MoE专家卡的本卡专家数。

• 通信域使用约束：

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

调用示例

以[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]为例，调用aclnnMoeDistributeDispatchV3和aclnnMoeDistributeCombineAddRmsNormV2接口

• 文件准备：
  1.新建combineAddRmsNormV2Demo目录，按照下方指导在combineAddRmsNormV2Demo下新建aclnnCombineAddRmsNormDemoV2.cpp，buildCombineAddRmsNormV2.sh，文件并修改。 2.将combineAddRmsNormV2Demo项目拷贝到服务器中。 3.安装cann包，并根据下方指导编译运行combineAddRmsNormV2Demo。

• 编译脚本

  [object Object]

• 编译与运行：

  [object Object]

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

  [object Object]