aclnnMoeUpdateExpert

产品支持情况

产品	是否支持
[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]	×

功能说明

算子功能：为了解决负载不均衡的场景，MOE网络中常用EPLB(Expert Parallelism Load Balancer)算法进行冗余专家部署，一个逻辑专家在多个卡上都有实例部署（即有多个物理专家），在这种场景下，MoeUpdateExpert算子可以完成每个token的topK个专家逻辑专家号到物理卡号的映射。

函数原型

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

• aclnnStatus aclnnMoeUpdateExpertGetWorkspaceSize(const aclTensor* expertIds, const aclTensor* eplbTable, const aclTensor* expertScalesOptional, const aclTensor* pruningThresholdOptional, const aclTensor* activeMaskOptional, int64_t localRankId, int64_t worldSize, int64_t balanceMode, aclTensor* balancedExpertIds, aclTensor* balancedActiveMask, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnMoeUpdateExpert(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

计算公式： 对于ExpertIds中的第个值，即第i个token:

$$expertId = expertIds[i]\\ tableOffset = expertId * F\\ eplbTable[tableOffset] = 1:\\ newExpertId = eplbTable[tableOffset + 1]\\ eplbTable[tableOffset] > 1 and balanceMode = 0:\\ modeValue = ceil(worldSize, eplbTable[tableOffset])\\ newExpertId = eplbTable[tableOffset + (localRankId / modeValue + 1)]\\ eplbTable[tableOffset] > 1 and balanceMode = 1:\\ newExpertId = eplbTable[tableOffset + (i \% eplbTable[tableOffset])]\\$$

注意该接口必须与aclnnMoeDistributeDispatchV2及aclnnMoeDistributeCombineV2或aclnnMoeDistributeCombineAddRmsNorm算子配套使用，调用顺序为aclnnMoeUpdateExpert，aclnnMoeDistributeDispatchV2，aclnnMoeDistributeCombineV2或aclnnMoeDistributeCombineAddRmsNorm。

aclnnMoeUpdateExpertGetWorkspaceSize

• 参数说明：

  • expertIds（aclTensor*，计算输入）：每个token的topK个专家索引，Device侧的aclTensor，要求为一个2D的Tensor，shape为 (BS, K)。数据类型支持INT32，INT64，undefined要求为ND，支持undefined。
  • eplbTable（aclTensor*，计算输入）：逻辑专家到物理专家的映射表，外部调用者需保证输入Tensor的值正确。world_size张卡，每张卡部署place_per_rank个路由专家实例，一共有world_size*place_per_rank个实例。eplbTable的每行第一列为行号对应逻辑专家部署的实例数count，取值范围[1, world_size]，每行[1, count]列为对应的实例编号，取值范围[0, world_size*place_per_rank)，有效的实例编号不可以重复。Device侧的aclTensor，要求是一个2D的Tensor。数据类型支持INT32，undefined要求为ND，支持undefined。shape为 (moeExperNum, F)。
  • expertScalesOptional（aclTensor*，计算输入）：每个token的topK个专家的scale权重，用户需保证scale在token内部按照降序排列，可选择传入的有效数据或空指针，该参数传入有效数据时，pruningThresholdOptional也需要传入有效数据。Device侧的aclTensor，要求是一个2D的Tensor。数据类型支持FP16、BF16、FLOAT，undefined要求为ND，支持undefined。shape为 (BS, K)。
  • pruningThresholdOptional（aclTensor*，计算输入）：专家scale权重的最小阈值，当某个token对应的某个topK专家scale小于阈值时，该token将对该专家进行剪枝，即token不发送至该专家处理，可选择传入有效数据或空指针，该参数传入有效数据时，expertScalesOptional也需要传入有效数据。Device侧的aclTensor，要求是一个1D或2D的Tensor。数据类型支持FLOAT，undefined要求为ND，支持undefined。shape为(K,)或(1, K)。
  • activeMaskOptional（aclTensor*，计算输入）：表示token是否参与通信，可选择传入有效数据或者空指针。传入有效数据时，expertScalesOptional、pruningThresholdOptional也必须传入有效数据，参数为true表示对应的token参与通信，true必须排到false之前，例：{true, false, true}为非法输入；传入空指针时是表示所有token都会参与通信。Device侧的aclTensor，要求是一个1D的Tensor。数据类型支持BOOL，undefined要求为ND，支持undefined。shape为 (BS，)。
  • localRankId（int64_t，计算输入）：本卡Id，数据类型支持INT64。当balanceMode为0时，取值范围为[0, worldSize)。同一个通信域中各卡的localRankId不重复。
  • worldSize（int64_t，计算输入）：通信域Size，数据类型支持INT64。当balanceMode为0时，取值范围为[2, 768]。
  • balanceMode（int64_t，计算输入）: 均衡规则，传入0时按照rank进行分发，传入1时按照token进行分发，默认值为0。数据类型支持INT64，取值范围为[0, 1]。
  • balancedExpertIds（aclTensor*，计算输出）：映射后每个token的topK个专家所在物理专家的实例编号，Device侧的aclTensor，要求是一个2D的Tensor，shape为（BS, K），数据类型、数据格式与expertIds保持一致。
  • balancedActiveMask（aclTensor*，计算输出）：剪枝后均衡的activeMask，当expertScalesOptional、pruningThresholdOptional传入有效数据时该输出有效。Device侧的aclTensor，要求是一个2D的Tensor，数据类型支持BOOL，undefined要求为ND，支持undefined。shape为（BS, K）。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值

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

  [object Object]

aclnnMoeUpdateExpert

• 参数说明：

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

• 返回值：

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

约束说明

• aclnnMoeUpdateExpert接口必须与aclnnMoeDistributeDispatchV2及aclnnMoeDistributeCombineV2或aclnnMoeDistributeCombineAddRmsNorm接口配套使用，调用顺序为aclnnUpdateExpert，aclnnMoeDistributeDispatchV2，aclnnMoeDistributeCombineV2或aclnnMoeDistributeCombineAddRmsNorm，具体参考undefined。

• 调用接口过程中使用的worldSize、moeExpertNum参数取值所有卡需保持一致，网络中不同层中也需保持一致，且和aclnnMoeDistributeDispatchV2,aclnnMoeDistributeCombineV2或aclnnMoeDistributeCombineAddRmsNorm对应参数也保持一致。

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

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

  • BS：表示batch sequence size，即本卡最终输出的token数量，取值范围为0 < BS ≤ 512。
  • K：表示选取topK个专家，取值范围为0 < K ≤ 16同时满足0 < K ≤ moeExpertNum。
  • moeExpertNum：表示Moe专家数量，取值范围(0, 1024]。
  • F: 表示映射表的列数，第一列为各行号对应Moe专家部署的实例个数(值>0)，后F-1列为该Moe专家部署的物理卡号，取值范围[2, worldSize + 1]。
  • 所有卡部署的moe专家实例总数最多1024，即place_per_rank * world_size ≤ 1024。
  • 每张卡部署的实例数需相同。

调用示例

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

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

• 编译脚本

  [object Object]

• 编译与运行：

  [object Object]

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

  [object Object]